from __future__ import annotations
import ast
from pathlib import Path
from .abyss import ASSETS_NAME as _ABYSS_ASSETS_NAME, DISPATCH_NAME as _ABYSS_DISPATCH_NAME, AbyssTransformer, build_runtime_stmts as build_abyss_runtime, encode_assets as encode_abyss_assets
from .lazy import encrypt_user_functions
from .loader import build_loader
from .packer import pack
from .transforms import FStringLowerer, IdentifierRenamer, MatchLowerer, 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'
_CONST_POOL_NAME = '_pw_literal_pool'
_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

def _normalize_future_imports(tree: ast.Module) -> None:
    if not tree.body:
        return
    doc: list[ast.stmt] = []
    rest = tree.body
    if (
        rest
        and isinstance(rest[0], ast.Expr)
        and isinstance(rest[0].value, ast.Constant)
        and isinstance(rest[0].value.value, str)
    ):
        doc = [rest[0]]
        rest = rest[1:]
    future = [stmt for stmt in rest if isinstance(stmt, ast.ImportFrom) and stmt.module == '__future__']
    other = [stmt for stmt in rest if not (isinstance(stmt, ast.ImportFrom) and stmt.module == '__future__')]
    tree.body = [*doc, *future, *other]

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, abyss: bool=False, abyss_functions: list[str] | set[str] | tuple[str, ...] | None=None, lower_fstrings: bool=True, lower_match: bool=True):
        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 ())
        self.abyss = bool(abyss)
        self.abyss_functions = set(abyss_functions or ())
        self.lower_fstrings = bool(lower_fstrings)
        self.lower_match = bool(lower_match)
        self.abyss_stats: dict[str, object] = {"protected": [], "skipped": []}

    def obfuscate(self, source: str) -> str:
        tree = ast.parse(source)
        skip = collect_skip_names(tree)
        if self.lower_fstrings:
            tree = FStringLowerer().visit(tree)
        if self.lower_match:
            tree = MatchLowerer().visit(tree)
        abyss_runtime_stmts: list[ast.stmt] = []
        abyss_keep: set[str] = set()
        if self.abyss or self.abyss_functions:
            abyss_transformer = AbyssTransformer(self.rng, targets=self.abyss_functions or None, auto=self.abyss)
            tree = abyss_transformer.protect(tree)
            abyss_keep = abyss_transformer.keep_names | {_ABYSS_DISPATCH_NAME, _ABYSS_ASSETS_NAME}
            self.abyss_stats = {
                "protected": list(abyss_transformer.protected),
                "skipped": list(abyss_transformer.skipped),
            }
            if abyss_transformer.assets:
                abyss_runtime_stmts = build_abyss_runtime(encode_abyss_assets(abyss_transformer.assets, self.rng), self.rng)
        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, _CONST_POOL_NAME)
            tree = enc.visit(tree)
            runtime_stmts = []
            pool_stmt = enc.build_pool_stmt()
            if pool_stmt is not None:
                runtime_stmts.append(pool_stmt)
            runtime_stmts.extend(build_decrypt_helper(_DEC_S_NAME, _DEC_B_NAME))
            _insert_runtime_stmts(tree, runtime_stmts)
        if self.rename:
            renamer = IdentifierRenamer(self.rng, keep=skip | self.keep_extra | abyss_keep | {_RESOLVE_NAME})
            tree = renamer.visit(tree)
        if abyss_runtime_stmts:
            _insert_runtime_stmts(tree, abyss_runtime_stmts)
        _normalize_future_imports(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