#!/usr/bin/env python3
import argparse
import sys
import shutil
import subprocess
from pathlib import Path

def main():
    parser = argparse.ArgumentParser(
        description="构建 affinity-sched run 包",
        add_help=True
    )
    parser.add_argument("--version", "-v", required=True, help="run 包版本号")
    parser.add_argument("--output", "-o", help="run 包输出目录（默认：./dist）")

    args = parser.parse_args()
    VERSION = args.version
    OUTPUT_DIR = args.output

    # ==================== 路径定义 ====================
    SCRIPT_DIR = Path(__file__).parent.absolute()
    ROOT_DIR = SCRIPT_DIR.parent

    # 工具路径
    MAKESELF_SH = ROOT_DIR / "makeself" / "makeself.sh"

    # help说明
    MAKESELF_HEADER = ROOT_DIR / "makeself" / "makeself-header.sh"
    HELP_FILTER = SCRIPT_DIR / "help.conf"

    INSTALL_PY    = SCRIPT_DIR / "install.py"
    UNINSTALL_PY  = SCRIPT_DIR / "uninstall.py"
    UNINSTALL_SH  = SCRIPT_DIR / "uninstall.sh"
    MAIN_SH       = SCRIPT_DIR / "main.sh"

    # 输出目录
    DIST_DIR = Path(OUTPUT_DIR)
    DIST_DIR.mkdir(parents=True, exist_ok=True)

    # 打包临时目录
    PKG_DIR = DIST_DIR / "run_package"

    #  包名
    PACKAGE_NAME = "mindstudio-boost"
    ARCH_LIST = ["aarch64", "x86_64"]

    # ==================== 检查依赖 ====================
    print("===== 检查环境 =====")

    if not MAKESELF_SH.exists():
        print(f"错误：未找到 makeself.sh：{MAKESELF_SH}", file=sys.stderr)
        sys.exit(1)
    if not INSTALL_PY.exists():
        print(f"错误：未找到 install.py：{INSTALL_PY}", file=sys.stderr)
        sys.exit(1)
    if not UNINSTALL_PY.exists():
        print(f"错误：未找到 uninstall.py：{UNINSTALL_PY}", file=sys.stderr)
        sys.exit(1)
    if not MAIN_SH.exists():
        print(f"错误：未找到 main.sh：{MAIN_SH}", file=sys.stderr)
        sys.exit(1)
    if not UNINSTALL_SH.exists():
        print(f"错误：未找到 uninstall.sh：{UNINSTALL_SH}", file=sys.stderr)
        sys.exit(1)

    # 查找 whl
    A_SCHED = "affinity-sched"
    whl_list = list(DIST_DIR.glob(f"{A_SCHED.replace('-', '_')}-*.whl"))
    if len(whl_list) == 0:
        print(f"错误：未找到 whl 包：{DIST_DIR}", file=sys.stderr)
        sys.exit(1)
    A_SCHED_WHL = whl_list[0]
    print(f"使用 whl 包：{A_SCHED_WHL.name}")

    # ==================== 准备打包文件 ====================
    print("===== 准备打包文件 =====")

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

    # 复制所有必需文件
    shutil.copy2(A_SCHED_WHL, PKG_DIR)
    shutil.copy2(INSTALL_PY, PKG_DIR)
    shutil.copy2(UNINSTALL_PY, PKG_DIR)
    shutil.copy2(UNINSTALL_SH, PKG_DIR)
    shutil.copy2(MAIN_SH, PKG_DIR)

    main_sh_in_pkg = PKG_DIR / "main.sh"
    main_sh_in_pkg.chmod(main_sh_in_pkg.stat().st_mode | 0o755)

    uninstall_sh_in_pkg = PKG_DIR / "uninstall.sh"
    uninstall_sh_in_pkg.chmod(uninstall_sh_in_pkg.stat().st_mode | 0o755)

    (PKG_DIR / "VERSION").write_text(VERSION)

    # ==================== 生成 run 包 ====================
    label = f"{PACKAGE_NAME} {VERSION}"

    for ARCH in ARCH_LIST:
        RUN_OUTPUT = DIST_DIR / f"{PACKAGE_NAME}_{VERSION}_{ARCH}.run"
        print(f"===== 生成 run 包 ({ARCH}) =====")

        cmd = [
            str(MAKESELF_SH),
            "--header", str(MAKESELF_HEADER),
            "--help-header", str(HELP_FILTER),
            "--tar-quietly",
            "--complevel", "4",
            "--nomd5",
            "--sha256",
            "--chown",
            str(PKG_DIR),
            str(RUN_OUTPUT),
            label,
            "./main.sh"
        ]

        try:
            subprocess.run(cmd, check=True)
        except subprocess.CalledProcessError:
            print("错误：run 包构建失败", file=sys.stderr)
            sys.exit(1)

        print(f"构建完成：{RUN_OUTPUT}")

    # 清理
    shutil.rmtree(PKG_DIR)

if __name__ == "__main__":
    main()