"""根据目标 CUDA 版本创建独立的 GPU uv 环境。"""

import argparse
import shutil
import subprocess
import sys
from pathlib import Path

PADDLE_INDEXES = {
    "cu118": "https://www.paddlepaddle.org.cn/packages/stable/cu118/",
    "cu126": "https://www.paddlepaddle.org.cn/packages/stable/cu126/",
}


def main() -> int:
    parser = argparse.ArgumentParser(
        description="安装 PaddleOCR-VL-1.6 GPU 子项目依赖",
    )
    parser.add_argument(
        "--cuda",
        choices=sorted(PADDLE_INDEXES),
        required=True,
        help="目标机器的 CUDA Wheel 类型；必须依据 PaddlePaddle 官方兼容表选择",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="只显示命令，不创建环境",
    )
    parser.add_argument(
        "--allow-no-gpu",
        action="store_true",
        help="允许在未检测到 nvidia-smi 时创建环境（仅用于准备/CI，不代表可运行）",
    )
    args = parser.parse_args()

    uv = shutil.which("uv")
    if not uv:
        print("[ERROR] 未找到 uv，请先安装 uv。", file=sys.stderr)
        return 1

    nvidia_smi = shutil.which("nvidia-smi")
    if not args.dry_run and not nvidia_smi and not args.allow_no_gpu:
        print(
            "[ERROR] 未检测到 nvidia-smi，拒绝在无 NVIDIA GPU 的机器安装 CUDA 依赖。\n"
            "如仅准备环境，请显式添加 --allow-no-gpu。",
            file=sys.stderr,
        )
        return 2

    project_dir = Path(__file__).resolve().parent
    index_url = PADDLE_INDEXES[args.cuda]
    command = [
        uv,
        "sync",
        "--project",
        str(project_dir),
        "--index",
        index_url,
    ]

    print(f"目标 CUDA Wheel: {args.cuda}")
    print(f"PaddlePaddle 索引: {index_url}")
    print("执行命令:")
    print("  " + " ".join(command))

    if args.dry_run:
        return 0

    ready_marker = project_dir / ".gpu-ready"
    if ready_marker.exists():
        ready_marker.unlink()

    completed = subprocess.run(command, check=False)
    if completed.returncode != 0:
        print(
            "[ERROR] 依赖安装失败。请检查 GPU、驱动、Python 和 PaddlePaddle Wheel 兼容性。",
            file=sys.stderr,
        )
        return completed.returncode

    ready_marker.write_text(
        f"cuda={args.cuda}\nindex={index_url}\n",
        encoding="utf-8",
    )
    print("\n[OK] GPU 子项目环境已创建。下一步从仓库根目录运行：")
    print("  python ocr.py verify --device gpu")
    return 0


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