"""设备检查。

项目硬性约定（见 CLAUDE.MD）：禁止在没有 Nvidia GPU 的设备上训练。
所有训练入口必须先调用 `require_cuda()`，检查不通过直接抛错退出。
"""

from __future__ import annotations


class NoNvidiaGpuError(RuntimeError):
    """当前设备不具备可用的 Nvidia GPU。"""


def cuda_devices() -> list[str]:
    """返回可用的 Nvidia GPU 名称列表；没有 torch 或没有 GPU 时返回空列表。"""
    try:
        import torch
    except ImportError:
        return []
    if not torch.cuda.is_available() or torch.version.cuda is None:
        return []
    return [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]


def require_cuda() -> list[str]:
    """训练前的守卫；无 Nvidia GPU 时抛出 `NoNvidiaGpuError`。"""
    try:
        import torch
    except ImportError as exc:
        raise NoNvidiaGpuError("未安装 PyTorch，无法确认 Nvidia GPU，训练已中止。") from exc

    if torch.version.cuda is None:
        raise NoNvidiaGpuError(
            f"当前 PyTorch ({torch.__version__}) 不是 CUDA 构建（可能是 CPU/ROCm/MPS 版本），"
            "本项目禁止在非 Nvidia GPU 设备上训练。"
        )
    if not torch.cuda.is_available() or torch.cuda.device_count() == 0:
        raise NoNvidiaGpuError(
            "未检测到可用的 Nvidia GPU，本项目禁止在无 GPU 的设备上训练。"
            "评估/推理请使用 eval.py 并显式指定 device=cpu。"
        )
    return cuda_devices()


def resolve_device(device: str = "cuda"):
    """把配置中的 device 字符串解析为 `torch.device`。

    仅用于评估/推理等允许 CPU 的场景；训练路径请先走 `require_cuda()`。
    """
    import torch

    if device.startswith("cuda") and not torch.cuda.is_available():
        raise NoNvidiaGpuError(f"请求了 {device}，但当前设备没有可用的 Nvidia GPU。")
    return torch.device(device)