#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
laser_cal 项目构建脚本
集成 CodeEnigma 混淆 + Nuitka 编译

支持三种编译模式：
    python build.py --module            # 编译为共享库 (.so / .pyd)
    python build.py --exe               # 编译为 standalone 可执行程序目录（推荐）
    python build.py --onefile           # 编译为单个可执行文件

其他选项：
    python build.py --skip-obfuscate    # 跳过 CodeEnigma 混淆步骤
    python build.py --no-license        # 不使用 license 文件授权（打包为无授权版本）
    python build.py --clean             # 清理构建缓存
"""

import os
import sys
import shutil
import subprocess
import argparse
import platform
import tempfile
from pathlib import Path

BASE_DIR = Path(__file__).parent.absolute()
DIST_DIR = BASE_DIR / "dist" / "program"
LICENSE_DIR = BASE_DIR / "dist" / "license"
OBFUSCATED_DIR = BASE_DIR / "dist" / "obfuscated"
BUILD_SRC_DIR = BASE_DIR / "dist" / "_build_src"

PROJECT_NAME = "laser_cal"

FILES_TO_OBFUSCATE = ["data_clean.py", "frequency_filter.py"]
LOCAL_MODULES = ["license_checker", "data_clean", "frequency_filter", "entry"]
PACKAGES_TO_INCLUDE = ["pandas", "numpy", "scipy", "cryptography"]

def log(msg, level="INFO"):
    print(f"[{level}] {msg}")

def run_cmd(cmd, cwd=None, env=None, allow_fail=False):
    log(f"执行: {' '.join(str(c) for c in cmd)}")
    result = subprocess.run([str(c) for c in cmd], cwd=str(cwd or BASE_DIR), env=env)
    if result.returncode != 0 and not allow_fail:
        log(f"命令失败，退出码: {result.returncode}", "ERROR")
        sys.exit(1)
    return result

def ensure_nuitka():
    try:
        subprocess.run([sys.executable, "-m", "nuitka", "--version"],
                       capture_output=True, check=True)
    except (FileNotFoundError, subprocess.CalledProcessError):
        log("安装 Nuitka...")
        run_cmd([sys.executable, "-m", "pip", "install",
                 "nuitka", "ordered-set", "zstandard", "--quiet"])

# ============================================================
# 步骤 1: CodeEnigma 混淆（修复版：混淆失败不中断构建）
# ============================================================

def step_obfuscate():
    log("=" * 60)
    log("步骤 1: CodeEnigma 代码混淆")
    log("=" * 60)

    # 确保 codeenigma 已安装
    try:
        subprocess.run(["codeenigma", "version"], capture_output=True, check=True)
    except (FileNotFoundError, subprocess.CalledProcessError):
        log("安装 codeenigma...")
        run_cmd([sys.executable, "-m", "pip", "install", "codeenigma", "--quiet"])

    if OBFUSCATED_DIR.exists():
        shutil.rmtree(OBFUSCATED_DIR)
    OBFUSCATED_DIR.mkdir(parents=True)

    with tempfile.TemporaryDirectory(dir=BASE_DIR / "dist", prefix="_tmp_mod_") as tmp_dir:
        tmp_path = Path(tmp_dir)
        for py_file in FILES_TO_OBFUSCATE:
            src = BASE_DIR / py_file
            if src.exists():
                shutil.copy2(src, tmp_path / py_file)
                log(f"  准备混淆: {py_file}")
            else:
                log(f"  文件不存在: {py_file}", "WARN")
        (tmp_path / "__init__.py").touch()

        # 执行混淆（允许失败）
        result = run_cmd(
            ["codeenigma", "obfuscate", str(tmp_path), "-o", str(OBFUSCATED_DIR)],
            allow_fail=True
        )

        if result.returncode != 0:
            log("CodeEnigma 报告错误，尝试自动恢复...", "WARN")

        # 1. 尝试从子目录恢复混淆后的 .py 文件
        for py_file in FILES_TO_OBFUSCATE:
            target = OBFUSCATED_DIR / py_file
            if not target.exists():
                found = list(OBFUSCATED_DIR.rglob(py_file))
                if found:
                    # 取第一个找到的文件（通常是 build 目录下的副本）
                    shutil.copy2(found[0], target)
                    log(f"  从 {found[0].relative_to(OBFUSCATED_DIR)} 恢复混淆文件: {py_file}")
                else:
                    log(f"  未找到混淆后的 {py_file}，稍后将回退到原始文件", "WARN")

        # 2. 恢复运行时库
        runtime_pyd = None
        search_dir = OBFUSCATED_DIR / "build"
        if search_dir.exists():
            if platform.system() == "Windows":
                candidates = list(search_dir.glob("**/*.pyd"))
            else:
                candidates = list(search_dir.glob("**/*.so"))
            candidates = [c for c in candidates if "codeenigma_runtime" in c.name]
            if candidates:
                runtime_pyd = candidates[0]
                log(f"  找到运行时库: {runtime_pyd}")

        runtime_dest_dir = OBFUSCATED_DIR / "codeenigma_runtime"
        runtime_dest_dir.mkdir(exist_ok=True)
        if runtime_pyd and runtime_pyd.exists():
            dest = runtime_dest_dir / runtime_pyd.name
            if not dest.exists():
                shutil.copy2(runtime_pyd, dest)
                log(f"  复制运行时库到: {dest}")
        else:
            log("  未找到编译后的运行时库，若程序依赖 codeenigma_runtime 将无法运行", "WARN")

        # 检查核心混淆文件最终是否可用
        core_ok = all((OBFUSCATED_DIR / f).exists() for f in FILES_TO_OBFUSCATE)
        if not core_ok:
            log("核心混淆文件最终仍缺失，构建将回退使用原始未混淆代码！", "ERROR")
            # 不退出，继续构建
        else:
            log("核心混淆文件已就绪，混淆成功")

# ============================================================
# 步骤 2: 准备编译目录（支持 --no-license，自动回退原始文件）
# ============================================================

def _generate_fake_license_module():
    sys.path.insert(0, str(BASE_DIR))
    try:
        import license_checker
    except ImportError:
        return "# Auto-generated fake license_checker\ndef check_license(*args, **kwargs): return True\n"

    public_names = [name for name in dir(license_checker) if not name.startswith('_')]
    if hasattr(license_checker, '__all__'):
        public_names = list(license_checker.__all__)

    lines = [
        "# -*- coding: utf-8 -*-",
        "# Auto-generated fake license_checker for no-license builds",
        ""
    ]
    for name in public_names:
        obj = getattr(license_checker, name, None)
        if callable(obj):
            lines.append(f"def {name}(*args, **kwargs):")
            lines.append(f"    return True")
            lines.append("")
        else:
            lines.append(f"{name} = None")
            lines.append("")
    sys.path.pop(0)
    return "\n".join(lines)

def step_prepare_build_dir(mode="module", no_license=False):
    log("=" * 60)
    log("步骤 2: 准备编译目录")
    log("=" * 60)

    if BUILD_SRC_DIR.exists():
        shutil.rmtree(BUILD_SRC_DIR)
    BUILD_SRC_DIR.mkdir(parents=True)

    if mode == "module":
        files_to_copy = ["entry.py", "license_checker.py"]
    else:
        files_to_copy = ["api_test.py", "license_checker.py", "entry.py"]

    for f in files_to_copy:
        src = BASE_DIR / f
        if not src.exists():
            continue
        if f == "license_checker.py" and no_license:
            log("  生成无授权模块: license_checker.py")
            with open(BUILD_SRC_DIR / f, "w", encoding="utf-8") as fout:
                fout.write(_generate_fake_license_module())
        else:
            shutil.copy2(src, BUILD_SRC_DIR / f)
            log(f"  复制: {f}")

    # 复制混淆后的文件（如果存在），否则回退原始文件
    for py_file in FILES_TO_OBFUSCATE:
        obf_file = OBFUSCATED_DIR / py_file
        if obf_file.exists():
            shutil.copy2(obf_file, BUILD_SRC_DIR / py_file)
            log(f"  复制混淆文件: {py_file}")
        else:
            orig_file = BASE_DIR / py_file
            if orig_file.exists():
                shutil.copy2(orig_file, BUILD_SRC_DIR / py_file)
                log(f"  混淆文件缺失，使用原始文件: {py_file}", "WARN")

    # 复制运行时库（如果有）
    runtime_dir = OBFUSCATED_DIR / "codeenigma_runtime"
    if runtime_dir.exists() and any(runtime_dir.iterdir()):
        shutil.copytree(runtime_dir, BUILD_SRC_DIR / "codeenigma_runtime")
        log("  复制 codeenigma_runtime")
    else:
        log("  未找到 codeenigma_runtime，跳过复制", "WARN")

    return BUILD_SRC_DIR

# ============================================================
# 步骤 3: Nuitka 编译
# ============================================================

def step_nuitka_module(build_src):
    log("=" * 60)
    log("步骤 3: Nuitka 编译为共享库 (.so / .pyd)")
    log("=" * 60)

    ensure_nuitka()

    cmd = [
        sys.executable, "-m", "nuitka",
        "--module",
        "--follow-import-to=license_checker",
        "--follow-import-to=data_clean",
        "--follow-import-to=frequency_filter",
        "--follow-import-to=entry",
        f"--output-dir={DIST_DIR}",
        "--python-flag=no_docstrings",
        "--python-flag=-O",
        "--remove-output",
        "--assume-yes-for-downloads",
    ]

    runtime_dir = build_src / "codeenigma_runtime"
    if runtime_dir.exists():
        for so_file in list(runtime_dir.glob("*.so")) + list(runtime_dir.glob("*.pyd")):
            cmd.append(f"--include-data-files={so_file}=codeenigma_runtime/{so_file.name}")

    cmd.append(str(build_src / "entry.py"))

    env = os.environ.copy()
    if platform.system() == "Linux":
        env["CC"] = "gcc"

    run_cmd(cmd, cwd=build_src, env=env)
    log(f"共享库已生成到: {DIST_DIR}")

def step_nuitka_exe(build_src, onefile=False):
    mode_name = "单文件可执行程序" if onefile else "standalone 可执行程序"
    log("=" * 60)
    log(f"步骤 3: Nuitka 编译为{mode_name}")
    log("=" * 60)

    ensure_nuitka()
    output_dir = DIST_DIR

    cmd = [sys.executable, "-m", "nuitka"]
    if onefile:
        cmd.append("--onefile")
    else:
        cmd.append("--standalone")

    cmd.extend([
        "--follow-imports",
        f"--output-dir={output_dir}",
        "--python-flag=no_docstrings",
        "--python-flag=-O",
        "--remove-output",
        "--assume-yes-for-downloads",
    ])

    for pkg in PACKAGES_TO_INCLUDE:
        cmd.append(f"--include-package={pkg}")

    cmd += [
        "--include-module=license_checker",
        "--include-module=data_clean",
        "--include-module=frequency_filter",
    ]

    runtime_dir = build_src / "codeenigma_runtime"
    if runtime_dir.exists():
        cmd.append(f"--include-data-dir={runtime_dir}=codeenigma_runtime")

    if platform.system() == "Windows":
        cmd.append(f"--output-filename={PROJECT_NAME}.exe")
    else:
        cmd.append(f"--output-filename={PROJECT_NAME}")

    cmd.append(str(build_src / "api_test.py"))

    env = os.environ.copy()
    if platform.system() == "Linux":
        env["CC"] = "gcc"

    run_cmd(cmd, cwd=build_src, env=env)

    if onefile:
        exe_name = PROJECT_NAME + ".exe" if platform.system() == "Windows" else PROJECT_NAME
        log(f"可执行文件已生成: {output_dir / exe_name}")
    else:
        log(f"可执行程序目录已生成: {output_dir / 'api_test.dist'}")

# ============================================================
# 步骤 4: 复制 License（根据 --no-license 跳过）
# ============================================================

def step_copy_license(mode="module", onefile=False, no_license=False):
    if no_license:
        log("=" * 60)
        log("步骤 4: 跳过 License 文件复制（无授权模式）")
        log("=" * 60)
        return

    log("=" * 60)
    log("步骤 4: 复制 License 文件")
    log("=" * 60)

    if not LICENSE_DIR.exists():
        log("  License 目录不存在，跳过", "WARN")
        return

    if mode == "module":
        target = DIST_DIR
    elif onefile:
        target = DIST_DIR
    else:
        target = DIST_DIR / "api_test.dist"

    target.mkdir(parents=True, exist_ok=True)
    for f in ["license.lic", "public_key.pem"]:
        src = LICENSE_DIR / f
        if src.exists():
            shutil.copy2(src, target / f)
            log(f"  复制: {f}")

# ============================================================
# 清理
# ============================================================

def clean():
    log("清理构建缓存...")
    for d in [DIST_DIR, OBFUSCATED_DIR, BUILD_SRC_DIR]:
        if d.exists():
            shutil.rmtree(d)
            log(f"  删除: {d}")

# ============================================================
# 主函数
# ============================================================

def main():
    parser = argparse.ArgumentParser(description="laser_cal 项目构建脚本 (CodeEnigma + Nuitka)",
                                     formatter_class=argparse.RawDescriptionHelpFormatter,
                                     epilog="""
编译模式说明：
  --module      编译为共享库 (.so/.pyd)，需要 Python 环境 + Node.js wrapper 调用
  --exe         编译为 standalone 可执行程序目录，Node.js 可直接 execFile 调用
  --onefile     编译为单个可执行文件，最便于分发

示例：
  python build.py --exe                       混淆 + 编译为可执行程序
  python build.py --exe --skip-obfuscate      仅编译为可执行程序
  python build.py --exe --no-license          混淆 + 编译为无授权可执行程序
  python build.py --module --skip-obfuscate   仅编译为共享库
  python build.py --onefile                   混淆 + 编译为单文件
  python build.py --clean                     清理构建缓存
        """)

    mode_group = parser.add_mutually_exclusive_group()
    mode_group.add_argument("--module", action="store_true")
    mode_group.add_argument("--exe", action="store_true")
    mode_group.add_argument("--onefile", action="store_true")

    parser.add_argument("--skip-obfuscate", action="store_true")
    parser.add_argument("--no-license", action="store_true")
    parser.add_argument("--clean", action="store_true")

    args = parser.parse_args()

    if args.clean:
        clean()
        return

    if not args.module and not args.exe and not args.onefile:
        args.module = True

    mode = "module" if args.module else ("exe" if args.exe else "onefile")

    log("=" * 60)
    log(f"laser_cal 构建脚本")
    log(f"  编译模式: {mode}")
    log(f"  代码混淆: {'是' if not args.skip_obfuscate else '否'}")
    log(f"  授权验证: {'否（无授权模式）' if args.no_license else '是（需 license 文件）'}")
    log(f"  操作系统: {platform.system()} ({platform.machine()})")
    log(f"  Python:   {sys.version.split()[0]}")
    log("=" * 60)

    # 混淆步骤（失败也不中断，会自动回退原始代码）
    if not args.skip_obfuscate:
        step_obfuscate()

    # 准备编译目录
    if not args.skip_obfuscate and OBFUSCATED_DIR.exists():
        build_src = step_prepare_build_dir(mode=mode, no_license=args.no_license)
    else:
        build_src = BASE_DIR

    # 编译
    if mode == "module":
        step_nuitka_module(build_src)
    elif mode == "exe":
        step_nuitka_exe(build_src, onefile=False)
    else:
        step_nuitka_exe(build_src, onefile=True)

    # 复制 License（无授权则跳过）
    step_copy_license(mode=mode, onefile=(mode == "onefile"), no_license=args.no_license)

    log("=" * 60)
    log("构建完成！")
    log(f"  输出目录: {DIST_DIR}")
    if args.no_license:
        log("  授权状态: 无授权版本（无需 license.lic / public_key.pem）")
    if mode == "module":
        log("  调用方式: Node.js -> laser_wrapper.py -> entry.pyd/.so")
    else:
        exe_name = PROJECT_NAME + (".exe" if platform.system() == "Windows" else "")
        log(f"  调用方式: Node.js -> child_process.execFile('{exe_name}', [api, data])")
    log("=" * 60)

if __name__ == "__main__":
    main()