"""核心包加密发布脚本。

优先 PyArmor；本机装不上时使用内置 marshal 加密后端（去可读源码）。

用法（仓库根）:
  python scripts/obfuscate_release.py --dry-run
  python scripts/obfuscate_release.py --backend marshal --output dist/obfuscated_lib
  python scripts/obfuscate_release.py --backend marshal --apply

不加密: lib/license/、模板 templates/、非目标包。
"""

from __future__ import annotations

import argparse
import base64
import compileall
import hashlib
import marshal
import py_compile
import random
import shutil
import struct
import subprocess
import sys
import zlib
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]

OBFUSCATE_PACKAGES = [
    "strategy",
    "trade",
    "key_monitor",
    "options",
    "hedge_plan",
]

_LOADER_TEMPLATE = '''# obfuscated module — do not edit
import base64, marshal, zlib
_B = {blob!r}
_K = {key!r}
def _d(b, k):
    raw = zlib.decompress(base64.b64decode(b))
    key = k.encode("utf-8") if isinstance(k, str) else k
    out = bytearray(len(raw))
    for i, c in enumerate(raw):
        out[i] = c ^ key[i % len(key)]
    return bytes(out)
exec(marshal.loads(_d(_B, _K)), globals())
'''


def _xor_compress(data: bytes, key: bytes) -> str:
    out = bytearray(len(data))
    for i, c in enumerate(data):
        out[i] = c ^ key[i % len(key)]
    return base64.b64encode(zlib.compress(bytes(out), 9)).decode("ascii")


def _obfuscate_py_file(src: Path, dest: Path, key: str) -> None:
    source = src.read_text(encoding="utf-8")
    code = compile(source, str(src), "exec", dont_inherit=True)
    blob = _xor_compress(marshal.dumps(code), key.encode("utf-8"))
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_text(_LOADER_TEMPLATE.format(blob=blob, key=key), encoding="utf-8")


def _copy_non_py(src_dir: Path, dest_dir: Path) -> None:
    for path in src_dir.rglob("*"):
        if path.is_dir():
            continue
        if path.suffix == ".py":
            continue
        rel = path.relative_to(src_dir)
        target = dest_dir / rel
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(path, target)


def marshal_obfuscate_package(src_pkg: Path, dest_pkg: Path, key: str) -> int:
    if dest_pkg.exists():
        shutil.rmtree(dest_pkg)
    dest_pkg.mkdir(parents=True, exist_ok=True)
    count = 0
    for path in src_pkg.rglob("*.py"):
        rel = path.relative_to(src_pkg)
        # keep __init__.py structure but still obfuscate
        _obfuscate_py_file(path, dest_pkg / rel, key)
        count += 1
    _copy_non_py(src_pkg, dest_pkg)
    return count


def _run(cmd: list[str], cwd: Path) -> None:
    print("+", " ".join(cmd))
    subprocess.check_call(cmd, cwd=str(cwd))


def try_pyarmor(packages: list[str], out: Path) -> bool:
    try:
        import pyarmor  # noqa: F401
    except ImportError:
        return False
    if out.exists():
        shutil.rmtree(out)
    out.mkdir(parents=True, exist_ok=True)
    src_dirs = [str(REPO_ROOT / "lib" / n) for n in packages]
    for cmd in (
        [sys.executable, "-m", "pyarmor.cli", "gen", "-O", str(out), *src_dirs],
        ["pyarmor", "gen", "-O", str(out), *src_dirs],
    ):
        try:
            _run(cmd, REPO_ROOT)
            return True
        except (subprocess.CalledProcessError, FileNotFoundError):
            continue
    return False


def main() -> int:
    parser = argparse.ArgumentParser(description="Obfuscate core lib packages")
    parser.add_argument("--output", default="dist/obfuscated_lib")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--apply", action="store_true", help="覆盖 lib/ 目标包（先备份到 dist/lib_backup_plain）")
    parser.add_argument(
        "--backend",
        choices=("auto", "pyarmor", "marshal"),
        default="auto",
        help="auto=有 pyarmor 用 pyarmor，否则 marshal",
    )
    args = parser.parse_args()

    lib = REPO_ROOT / "lib"
    out = REPO_ROOT / args.output
    packages = [n for n in OBFUSCATE_PACKAGES if (lib / n).is_dir()]
    print("packages:", ", ".join(packages))
    if args.dry_run:
        return 0

    backend = args.backend
    if backend == "auto":
        backend = "pyarmor"
        try:
            import pyarmor  # noqa: F401
        except ImportError:
            backend = "marshal"
            print("pyarmor 未安装，使用 marshal 后端")

    key = hashlib.sha256(b"crypto_monitor_user_release_v1").hexdigest()[:32]

    if backend == "pyarmor":
        ok = try_pyarmor(packages, out)
        if not ok:
            print("pyarmor 失败，回退 marshal", file=sys.stderr)
            backend = "marshal"

    if backend == "marshal":
        if out.exists():
            shutil.rmtree(out)
        out.mkdir(parents=True, exist_ok=True)
        total = 0
        for name in packages:
            n = marshal_obfuscate_package(lib / name, out / name, key)
            print(f"marshal {name}: {n} py files")
            total += n
        print(f"encrypted output: {out} ({total} files)")

    # smoke import from output path
    sys.path.insert(0, str(out.parent))
    # packages live in out/<name>, need lib-like layout for import? 
    # For smoke: add out to path and import strategy if it has __init__
    smoke_root = out
    sys.path.insert(0, str(smoke_root))
    try:
        import importlib

        # packages are top-level under out/
        m = importlib.import_module("trade.trade_policy_lib")
        print("smoke import trade.trade_policy_lib:", getattr(m, "__file__", m))
    except Exception as e:
        print("smoke import note:", e)

    if args.apply:
        for name in packages:
            candidates = [out / name, out / "lib" / name]
            src = next((c for c in candidates if c.exists()), None)
            if src is None:
                print(f"警告: 未找到 {name}")
                continue
            dest = lib / name
            backup = REPO_ROOT / "dist" / "lib_backup_plain" / name
            backup.parent.mkdir(parents=True, exist_ok=True)
            if dest.exists():
                if backup.exists():
                    shutil.rmtree(backup)
                shutil.copytree(dest, backup)
                shutil.rmtree(dest)
            shutil.copytree(src, dest)
            print(f"applied: lib/{name}")
        print("完成。请验证启动后再 push。")
    return 0


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