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

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

from sagellm.ascend_repair import repair_ascend_runtime
from sagellm.cuda_repair import repair_cuda_runtime

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

VERIFIED_GITHUB_WHEEL_RELEASES: dict[str, tuple[str, str]] = {
    # Populate this only after the matching wheel set is published and smoke-tested.
}


@dataclass(frozen=True)
class BackendRepairSpec:
    """Configuration for a packaged backend repair flow."""

    runner: Callable[..., int]
    index_url: str | None = None


def resolve_backend_index_url(raw_index_url: str | None) -> str | None:
    """Resolve symbolic PyTorch index keys into concrete URLs."""
    if not raw_index_url:
        return None

    return PYTORCH_MIRRORS["official"].get(raw_index_url, raw_index_url)


def resolve_github_release_url(backend: str) -> str | None:
    """Return accelerated GitHub wheel release URL for verified backends only."""
    release = VERIFIED_GITHUB_WHEEL_RELEASES.get(backend)
    if release is None:
        return None
    pytorch_ver, platform = release
    return f"https://github.com/intellistream/sagellm-pytorch-wheels/releases/download/v{pytorch_ver}-{platform}/"


PACKAGED_REPAIR_SPECS = {
    "cuda": BackendRepairSpec(runner=repair_cuda_runtime, index_url="cu128"),
    "cuda11": BackendRepairSpec(runner=repair_cuda_runtime, index_url="cu118"),
    "cuda12": BackendRepairSpec(runner=repair_cuda_runtime, index_url="cu128"),
    "ascend": BackendRepairSpec(runner=repair_ascend_runtime),
}


def run_packaged_backend_repair(
    backend: str,
    *,
    python_executable: str,
) -> int | None:
    """Run a packaged repair flow for supported backends."""
    spec = PACKAGED_REPAIR_SPECS.get(backend)
    if spec is None:
        return None

    kwargs = {
        "python_executable": python_executable,
    }
    resolved_index_url = resolve_backend_index_url(spec.index_url)
    if resolved_index_url is not None:
        kwargs["torch_index_url"] = resolved_index_url

    return spec.runner(**kwargs)


# 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",
        "cu128": "https://download.pytorch.org/whl/cu128",
        "rocm5.7": "https://download.pytorch.org/whl/rocm5.7",
        "cpu": "https://download.pytorch.org/whl/cpu",
    },
}

PYTORCH_INSTALL_CONFIGS = {
    # === NVIDIA ===
    "cuda": {
        "name": "CUDA 12.8 (NVIDIA GPU)",
        "priority": 50,
        "index_url": "cu128",
        "packages": ["torch", "torchvision", "torchaudio"],
        "check": lambda: check_cuda_available(),
        "version_marker": "cu128",
        "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.8 (newer NVIDIA drivers)",
        "priority": 50,
        "index_url": "cu128",
        "packages": ["torch", "torchvision", "torchaudio"],
        "check": lambda: check_cuda_available(),
        "version_marker": "cu128",
        "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": ["numpy<2", "decorator", "scipy", "torch", "torch_npu"],
        "check": lambda: check_ascend_available(),
        "version_marker": None,
        "verify_script": "import importlib; import torch; import torch_npu; importlib.import_module('tbe'); 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__)",
    },
}
