from __future__ import annotations
import ast
from pathlib import Path
from .lazy import encrypt_user_functions
from .loader import build_loader
from .packer import pack
from .transforms import IdentifierRenamer, NumberObfuscator, OpaquePredicateInjector, StringEncryptor, build_decrypt_helper, collect_skip_names
from .transforms.junk import JunkBranchInjector
from .transforms.mba import MBATransformer
from .transforms.opaque import build_seed_stmt
from .util import make_rng
_SEED_NAME = '_pw_seed_value'
_DEC_S_NAME = '_pw_dec_str'
_DEC_B_NAME = '_pw_dec_bytes'
_RESOLVE_NAME = '__pw_resolve_lazy__'

def _insert_runtime_stmts(tree: ast.Module, stmts: list[ast.stmt]) -> None:
    if not stmts:
        return
    insert_at = 0
    if (
        tree.body
        and isinstance(tree.body[0], ast.Expr)
        and isinstance(tree.body[0].value, ast.Constant)
        and isinstance(tree.body[0].value.value, str)
    ):
        insert_at = 1
    while (
        insert_at < len(tree.body)
        and isinstance(tree.body[insert_at], ast.ImportFrom)
        and tree.body[insert_at].module == '__future__'
    ):
        insert_at += 1
    tree.body[insert_at:insert_at] = stmts

class Obfuscator:

    def __init__(self, *, seed: int | None=None, rename: bool=True, encrypt_strings: bool=True, obfuscate_numbers: bool=True, opaque_predicates: bool=True, mba: bool=True, junk_branches: bool=True, lazy_funcs: bool=True, anti_debug: bool=True, layers: int=3, stage2_layers: int=3, keep: set[str] | None=None):
        self.rng, self.seed = make_rng(seed)
        self.rename = rename
        self.encrypt_strings = encrypt_strings
        self.obfuscate_numbers = obfuscate_numbers
        self.opaque_predicates = opaque_predicates
        self.mba = mba
        self.junk_branches = junk_branches
        self.lazy_funcs = lazy_funcs
        self.anti_debug = anti_debug
        self.layers = max(1, int(layers))
        self.stage2_layers = max(1, int(stage2_layers))
        self.keep_extra = set(keep or ())

    def obfuscate(self, source: str) -> str:
        tree = ast.parse(source)
        skip = collect_skip_names(tree)
        seed_used = False
        if self.opaque_predicates:
            opi = OpaquePredicateInjector(self.rng, _SEED_NAME)
            tree = opi.visit(tree)
            seed_used = seed_used or opi.injected
        if self.junk_branches:
            jbi = JunkBranchInjector(self.rng, _SEED_NAME)
            tree = jbi.visit(tree)
            seed_used = seed_used or jbi.injected
        if seed_used:
            _insert_runtime_stmts(tree, [build_seed_stmt(_SEED_NAME)])
        if self.mba:
            tree = MBATransformer(self.rng).visit(tree)
        if self.obfuscate_numbers:
            tree = NumberObfuscator(self.rng).visit(tree)
        if self.encrypt_strings:
            enc = StringEncryptor(self.rng, _DEC_S_NAME, _DEC_B_NAME)
            tree = enc.visit(tree)
            _insert_runtime_stmts(tree, build_decrypt_helper(_DEC_S_NAME, _DEC_B_NAME))
        if self.rename:
            renamer = IdentifierRenamer(self.rng, keep=skip | self.keep_extra | {_RESOLVE_NAME})
            tree = renamer.visit(tree)
        ast.fix_missing_locations(tree)
        user_code = compile(tree, '<patchwork>', 'exec')
        if self.lazy_funcs:
            user_code, lazy_blobs = encrypt_user_functions(user_code, self.rng, _RESOLVE_NAME)
        else:
            lazy_blobs = []
        user_payload, user_keys = pack(user_code, self.rng, layers=self.layers)
        return build_loader(user_payload, user_keys, lazy_blobs, resolver_name=_RESOLVE_NAME, rng=self.rng, anti_debug=self.anti_debug, stage2_layers=self.stage2_layers)

def obfuscate(source: str, **kwargs) -> str:
    return Obfuscator(**kwargs).obfuscate(source)

def obfuscate_file(input_path: str | Path, output_path: str | Path | None=None, **kwargs) -> Path:
    inp = Path(input_path)
    src = inp.read_text(encoding='utf-8')
    out = obfuscate(src, **kwargs)
    if output_path is None:
        outp = inp.with_name(inp.stem + '_obf.py')
    else:
        outp = Path(output_path)
    outp.write_text(out, encoding='utf-8')
    return outp