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

PACKAGE_NAME = "affinity-sched"

def main():
    # ==================== 参数解析 ====================
    parser = argparse.ArgumentParser(
        description="构建 affinity-sched whl 包",
        add_help=True
    )
    parser.add_argument("--version", "-v", required=True, help="指定 whl 包版本号（必填）")
    parser.add_argument("--output", "-o", help="指定输出目录（可选，默认：./dist）")
    
    args = parser.parse_args()
    VERSION = args.version
    OUTPUT_DIR = args.output

    # ==================== 路径定义 ====================
    SCRIPT_DIR = Path(__file__).parent.absolute()
    A_SCHED_DIR = SCRIPT_DIR.parent  # 项目根目录
    VERSION_FILE = A_SCHED_DIR / "VERSION"

    # 默认输出目录
    DEFAULT_DIST_DIR = A_SCHED_DIR / "dist"
    DIST_DIR = Path(OUTPUT_DIR) if OUTPUT_DIR else DEFAULT_DIST_DIR
    DIST_DIR.mkdir(parents=True, exist_ok=True)

    # ==================== 构建 whl ====================
    print(f"\n===== 构建 affinity-sched whl 包，版本: {VERSION} =====")

    VERSION_FILE.write_text(args.version.strip())
    
    cmd = [
        sys.executable, "-m", "pip", "wheel",
        "--no-deps",
        "--use-pep517",
        "--wheel-dir", str(DIST_DIR),
        "."
    ]

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

    whl = next(DIST_DIR.glob(f"{PACKAGE_NAME.replace('-', '_')}-{VERSION}-*.whl"), None)
    if not whl:
        print("错误：未找到构建好的 whl 包", file=sys.stderr)
        sys.exit(1)
    print(f"\n构建完成：{DIST_DIR}/{whl.name}")

    # 清理 build 目录
    if os.path.isdir("build"):
        shutil.rmtree("build")
        
    # 清理所有 .egg-info 目录
    for egg_dir in glob.glob("*.egg-info"):
        if os.path.isdir(egg_dir):
            shutil.rmtree(egg_dir)

if __name__ == "__main__":
    main()