from __future__ import annotations
import marshal
import random
import types
from .crypto import stream_xor
from .util import gen_bytes
LAZY_PREFIX = '<pw'
LAZY_SUFFIX = '>'
_CO_GENERATOR = 0x20
_CO_COROUTINE = 0x100
_CO_ASYNC_GENERATOR = 0x200


def _classify_code(code: types.CodeType) -> str:
    if not code.co_flags & 1:
        return 'class'
    if code.co_name.startswith('<') and code.co_name.endswith('>'):
        return 'special'
    return 'function'


def _stub_source(kind: str, idx: int, resolver_name: str) -> str:
    if kind == 'async_gen':
        return (
            f'async def _stub(*_a, **_k):\n'
            f'    async for _v in {resolver_name}({idx}, _a, _k):\n'
            f'        yield _v\n'
        )
    if kind == 'coroutine':
        return (
            f'async def _stub(*_a, **_k):\n'
            f'    return await {resolver_name}({idx}, _a, _k)\n'
        )
    if kind == 'generator':
        return (
            f'def _stub(*_a, **_k):\n'
            f'    yield from {resolver_name}({idx}, _a, _k)\n'
        )
    return (
        f'def _stub(*_a, **_k):\n'
        f'    return {resolver_name}({idx}, _a, _k)\n'
    )


def _kind_for(code: types.CodeType) -> str:
    f = code.co_flags
    if f & _CO_ASYNC_GENERATOR:
        return 'async_gen'
    if f & _CO_COROUTINE:
        return 'coroutine'
    if f & _CO_GENERATOR:
        return 'generator'
    return 'function'


def _make_stub_code(idx: int, original: types.CodeType, resolver_name: str) -> types.CodeType:
    kind = _kind_for(original)
    src = _stub_source(kind, idx, resolver_name)
    mod = compile(src, '<pwstub>', 'exec')
    inner = next((c for c in mod.co_consts if isinstance(c, types.CodeType)))
    return inner.replace(
        co_name=f'{LAZY_PREFIX}{idx:x}{LAZY_SUFFIX}',
        co_freevars=original.co_freevars,
        co_cellvars=original.co_cellvars,
    )


def encrypt_user_functions(top: types.CodeType, rng: random.Random, resolver_name: str, key_size: int=24) -> tuple[types.CodeType, list[tuple[bytes, bytes]]]:
    blobs: list[tuple[bytes, bytes]] = []

    def transform(code: types.CodeType, parent_kind: str) -> types.CodeType:
        new_consts: list = []
        for c in code.co_consts:
            if isinstance(c, types.CodeType):
                ck = _classify_code(c)
                if parent_kind in ('module', 'class') and ck == 'function':
                    idx = len(blobs)
                    key = gen_bytes(rng, key_size)
                    enc = stream_xor(marshal.dumps(c), key)
                    blobs.append((key, enc))
                    new_consts.append(_make_stub_code(idx, c, resolver_name))
                else:
                    new_consts.append(transform(c, ck))
            else:
                new_consts.append(c)
        return code.replace(co_consts=tuple(new_consts))
    new_top = transform(top, 'module')
    return (new_top, blobs)