from __future__ import annotations

import random
import os
import subprocess
import sys
from pathlib import Path
import tempfile
import textwrap

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from patchwork import Obfuscator


GENERATED_SEEDS = [3, 11, 29, 101, 1009]
OBFUSCATION_SEEDS = [5, 17, 257]


def run(path: Path) -> str:
    env = dict(os.environ)
    env["PYTHONIOENCODING"] = "utf-8"
    result = subprocess.run([sys.executable, str(path)], capture_output=True, text=True, timeout=20, env=env)
    if result.returncode != 0:
        raise RuntimeError(f"{path} exited {result.returncode}\nstderr:\n{result.stderr}")
    return result.stdout


def generated_program(seed: int) -> str:
    rng = random.Random(seed)
    numbers = [rng.randint(-5000, 5000) for _ in range(16)]
    strings = [f"s{idx}_{rng.randrange(1 << 20):x}_\u2603_\U0001f680" for idx in range(6)]
    pairs = [(rng.randint(-20, 20), rng.randint(-20, 20)) for _ in range(8)]
    return textwrap.dedent(
        f"""
        from __future__ import annotations

        NUMBERS = {numbers!r}
        STRINGS = {strings!r}
        PAIRS = {pairs!r}

        def fold(values):
            acc = 0
            for idx, value in enumerate(values):
                if idx % 3 == 0:
                    acc ^= value
                elif idx % 3 == 1:
                    acc += value << (idx % 5)
                else:
                    acc -= value >> (idx % 4)
            return acc

        def classify(item):
            match item:
                case (0, y):
                    return f"zero:{{y}}"
                case (x, y) if x == y:
                    return f"same:{{x}}"
                case (x, y):
                    return f"pair:{{x + y}}"
                case _:
                    return "other"

        def closure(base):
            state = base
            def inner(delta):
                nonlocal state
                state = (state + delta) ^ (delta << 2)
                return state
            return inner

        def run():
            fn = closure(7)
            print("fold", fold(NUMBERS))
            print("strings", "|".join(s[::-1] for s in STRINGS))
            print("pairs", [classify(pair) for pair in PAIRS])
            print("closure", [fn(x) for x in range(6)])
            print("bytes", (bytes(range(8)) + b"literal-marker").hex())

        if __name__ == "__main__":
            run()
        """
    ).lstrip()


def check_source(source: str, *, label: str, seed: int, **kwargs) -> None:
    obf_source = Obfuscator(seed=seed, **kwargs).obfuscate(source)
    with tempfile.TemporaryDirectory() as raw:
        directory = Path(raw)
        src_path = directory / f"{label}.py"
        obf_path = directory / f"{label}_obf.py"
        src_path.write_text(source, encoding="utf-8")
        obf_path.write_text(obf_source, encoding="utf-8")
        expected = run(src_path)
        actual = run(obf_path)
    if actual != expected:
        raise AssertionError(f"{label} seed={seed} mismatch\n--- expected ---\n{expected}--- actual ---\n{actual}")


def main() -> int:
    failures = 0
    matrix = [
        {},
        {"lazy_funcs": False},
        {"anti_debug": False},
        {"rename": False},
        {"encrypt_strings": False},
        {"opaque_predicates": False, "junk_branches": False},
    ]
    for generated_seed in GENERATED_SEEDS:
        source = generated_program(generated_seed)
        for obfuscation_seed in OBFUSCATION_SEEDS:
            for options in matrix:
                label = f"generated_{generated_seed}_{obfuscation_seed}_{len(options)}"
                try:
                    check_source(source, label=label, seed=obfuscation_seed, **options)
                    print(f"ok: {label} {options}")
                except Exception as exc:
                    failures += 1
                    print(f"FAIL: {label} {options}: {exc}")
    print(f"\n{failures} failures" if failures else "\nstress matrix all good")
    return 1 if failures else 0


if __name__ == "__main__":
    raise SystemExit(main())