"""PyTorch installation configurations for different backends."""

from .ascend import check_ascend_available
from .cuda import check_cuda_available
from .domestic import check_haiguang_available, check_kunlun_available, check_muxi_available

# Backend priority for auto-detection (higher = preferred)
BACKEND_PRIORITY = {
    "ascend": 100,  # 华为昇腾
    "kunlun": 90,  # 百度昆仑
    "haiguang": 85,  # 海光 DCU
    "muxi": 80,  # 沐曦
    "cuda": 50,  # NVIDIA
    "cpu": 10,  # CPU 回退
}

# PyTorch Mirror Sites (国内镜像源)
# 注意：清华源和阿里源不镜像 PyTorch 官方 wheels
# PyTorch CUDA 版本必须从 download.pytorch.org 下载
PYTORCH_MIRRORS = {
    "official": {
        "cu121": "https://download.pytorch.org/whl/cu121",
        "cu118": "https://download.pytorch.org/whl/cu118",
        "rocm5.7": "https://download.pytorch.org/whl/rocm5.7",
        "cpu": "https://download.pytorch.org/whl/cpu",
    },
}

PYTORCH_INSTALL_CONFIGS = {
    # === NVIDIA ===
    "cuda": {
        "name": "CUDA (NVIDIA GPU)",
        "priority": 50,
        "index_url": "cu121",  # 改为 key，运行时从 PYTORCH_MIRRORS 查找
        "packages": ["torch", "torchvision", "torchaudio"],
        "check": lambda: check_cuda_available(),
        "version_marker": "cu121",
        "verify_script": "import torch; print(torch.__version__); assert torch.cuda.is_available(), 'CUDA not available'",
    },
    "cuda11": {
        "name": "CUDA 11.8 (older NVIDIA drivers)",
        "priority": 50,
        "index_url": "https://download.pytorch.org/whl/cu118",
        "packages": ["torch", "torchvision", "torchaudio"],
        "check": lambda: check_cuda_available(),
        "version_marker": "cu118",
        "verify_script": "import torch; print(torch.__version__); assert torch.cuda.is_available(), 'CUDA not available'",
    },
    "cuda12": {
        "name": "CUDA 12.1 (newer NVIDIA drivers)",
        "priority": 50,
        "index_url": "https://download.pytorch.org/whl/cu121",
        "packages": ["torch", "torchvision", "torchaudio"],
        "check": lambda: check_cuda_available(),
        "version_marker": "cu121",
        "verify_script": "import torch; print(torch.__version__); assert torch.cuda.is_available(), 'CUDA not available'",
    },
    # === 华为昇腾 Ascend ===
    "ascend": {
        "name": "Ascend NPU (华为昇腾)",
        "priority": 100,
        "index_url": None,  # Special: requires CANN toolkit
        "packages": ["torch", "torch_npu"],
        "check": lambda: check_ascend_available(),
        "version_marker": None,
        "verify_script": "import torch; import torch_npu; print(torch.__version__); assert torch_npu.npu.is_available(), 'NPU not available'",
        "setup_guide": "https://www.hiascend.com/",
    },
    # === 百度昆仑 Kunlun ===
    "kunlun": {
        "name": "Kunlun XPU (百度昆仑)",
        "priority": 90,
        "index_url": None,  # Special: requires Kunlun SDK
        "packages": ["torch", "torch_xmlir"],  # or paddlepaddle for some cases
        "check": lambda: check_kunlun_available(),
        "version_marker": None,
        "verify_script": "import torch; import torch_xmlir; print(torch.__version__)",
        "setup_guide": "https://www.kunlunxin.com/",
    },
    # === 海光 DCU (Hygon) ===
    "haiguang": {
        "name": "Haiguang DCU (海光)",
        "priority": 85,
        "index_url": "https://download.pytorch.org/whl/rocm5.7",
        "packages": ["torch", "torchvision", "torchaudio"],
        "check": lambda: check_haiguang_available(),
        "version_marker": "rocm",
        "verify_script": "import torch; print(torch.__version__); assert torch.cuda.is_available(), 'DCU not available'",  # DCU uses CUDA API
        "setup_guide": "https://www.hygon.cn/",
    },
    # === 沐曦 Muxi ===
    "muxi": {
        "name": "Muxi GPU (沐曦)",
        "priority": 80,
        "index_url": None,  # Special: requires MUSA toolkit
        "packages": ["torch", "torch_muxi"],
        "check": lambda: check_muxi_available(),
        "version_marker": None,
        "verify_script": "import torch; import torch_muxi; print(torch.__version__); assert torch.muxi.is_available(), 'Muxi not available'",
        "setup_guide": "https://www.muxi.com/",
    },
    # === CPU ===
    "cpu": {
        "name": "CPU only",
        "priority": 10,
        "index_url": "https://download.pytorch.org/whl/cpu",
        "packages": ["torch", "torchvision", "torchaudio"],
        "check": lambda: (True, "✅ CPU ready"),
        "version_marker": "cpu",
        "verify_script": "import torch; print(torch.__version__)",
    },
}
