"""Torch environment isolation utilities for sageLLM (#43).

Provides helpers to:
1. Detect which torch variant is currently installed (CUDA / Ascend /
   Kunlun / Haiguang / Muxi / CPU-only)
2. Validate that the correct variant matches the detected hardware
3. Generate per-backend install commands so that CI and quickstart scripts
   can stay declarative

The key principle is **fail-fast**: if torch is imported but its variant
does not match the available hardware, sageLLM should raise a clear error
rather than silently fall back to CPU.

Usage
-----
From Python::

    from sagellm.torch_env import TorchEnvInfo, detect_torch_env

    info = detect_torch_env()
    print(info.variant)       # "cuda", "ascend", "cpu", ...
    print(info.torch_version) # "2.1.0+cu121"

From CLI::

    sagellm env --torch-info
"""

from __future__ import annotations

import importlib
import importlib.util
import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

logger = logging.getLogger(__name__)


class TorchVariant(str, Enum):
    """Known torch build variants.

    Attributes:
        CUDA: Standard NVIDIA CUDA build (``+cuXXX`` suffix).
        ROCM: AMD ROCm build.
        ASCEND: Huawei Ascend (CANN / torch_npu).
        KUNLUN: Baidu Kunlun (torch_xmlir).
        HAIGUANG: Haiguang DCU (torch-dcu / hip).
        MUXI: MThreads Muxi (torch-musa).
        CPU: CPU-only build.
        UNKNOWN: torch is present but variant cannot be determined.
        NOT_INSTALLED: torch is not installed.
    """

    CUDA = "cuda"
    ROCM = "rocm"
    ASCEND = "ascend"
    KUNLUN = "kunlun"
    HAIGUANG = "haiguang"
    MUXI = "muxi"
    CPU = "cpu"
    UNKNOWN = "unknown"
    NOT_INSTALLED = "not_installed"


@dataclass
class TorchEnvInfo:
    """Summary of the current torch environment.

    Attributes:
        variant: Detected :class:`TorchVariant`.
        torch_version: Version string from ``torch.__version__``
            (e.g. ``"2.1.0+cu121"``).
        device_count: Number of accelerator devices visible to torch.
        device_names: Human-readable names of visible devices.
        extra: Additional metadata from the specific backend.
    """

    variant: TorchVariant = TorchVariant.NOT_INSTALLED
    torch_version: str = ""
    device_count: int = 0
    device_names: list[str] = field(default_factory=list)
    extra: dict[str, Any] = field(default_factory=dict)

    def is_gpu_available(self) -> bool:
        """Return ``True`` if at least one accelerator is available."""
        return self.device_count > 0

    def is_compatible_with(self, backend: str) -> bool:
        """Check whether this torch env is compatible with a sageLLM backend.

        Args:
            backend: Backend name: ``"cpu"``, ``"cuda"``, ``"ascend"``,
                ``"kunlun"``, ``"haiguang"``, ``"muxi"``.

        Returns:
            ``True`` when compatible; ``False`` when a mismatch is detected.
        """
        backend = backend.lower()
        if backend == "cpu":
            return True  # CPU works with any torch variant
        mapping = {
            "cuda": TorchVariant.CUDA,
            "ascend": TorchVariant.ASCEND,
            "kunlun": TorchVariant.KUNLUN,
            "haiguang": TorchVariant.HAIGUANG,
            "muxi": TorchVariant.MUXI,
        }
        required = mapping.get(backend)
        if required is None:
            return False
        return self.variant == required


def detect_torch_env() -> TorchEnvInfo:
    """Detect the current torch environment.

    Imports torch (if available) and inspects which variant is installed.
    This function should **not** raise — it returns a :class:`TorchEnvInfo`
    with ``variant=NOT_INSTALLED`` when torch is missing.

    Returns:
        :class:`TorchEnvInfo` describing the detected environment.
    """
    if importlib.util.find_spec("torch") is None:
        return TorchEnvInfo(variant=TorchVariant.NOT_INSTALLED)

    try:
        import torch  # type: ignore[import]
    except ImportError:
        return TorchEnvInfo(variant=TorchVariant.NOT_INSTALLED)

    version: str = getattr(torch, "__version__", "")
    info = TorchEnvInfo(torch_version=version)

    # 1. Check Ascend (torch_npu)
    if importlib.util.find_spec("torch_npu") is not None:
        try:
            import torch_npu  # type: ignore[import]

            info.variant = TorchVariant.ASCEND
            npu_count = torch_npu.npu.device_count() if hasattr(torch_npu, "npu") else 0
            info.device_count = npu_count
            info.device_names = [torch_npu.npu.get_device_name(i) for i in range(npu_count)]
            return info
        except Exception as exc:
            logger.debug("torch_npu import/check failed: %s", exc)

    # 2. Check Kunlun (torch_xmlir)
    if importlib.util.find_spec("torch_xmlir") is not None:
        info.variant = TorchVariant.KUNLUN
        return info

    # 3. Check MThreads Muxi (torch_musa)
    if importlib.util.find_spec("torch_musa") is not None:
        info.variant = TorchVariant.MUXI
        return info

    # 4. Check CUDA
    if hasattr(torch, "cuda") and torch.cuda.is_available():
        info.variant = TorchVariant.CUDA
        info.device_count = torch.cuda.device_count()
        info.device_names = [torch.cuda.get_device_name(i) for i in range(info.device_count)]
        # Detect ROCm (AMD)
        if "+rocm" in version or (hasattr(torch.version, "hip") and torch.version.hip is not None):
            info.variant = TorchVariant.ROCM
        # Detect Haiguang DCU (HIP device but not AMD ROCm)
        elif hasattr(torch, "_C") and "+dcu" in version:
            info.variant = TorchVariant.HAIGUANG
        return info

    # 5. CPU-only
    info.variant = TorchVariant.CPU
    return info


def validate_torch_for_backend(backend: str) -> None:
    """Validate that the installed torch variant matches the requested backend.

    This implements the **fail-fast** principle: if torch is present but
    incompatible, raise a clear ``EnvironmentError`` rather than silently
    falling back.

    Args:
        backend: Backend name (``"cuda"``, ``"ascend"``, ``"cpu"``, …).

    Raises:
        EnvironmentError: If the torch variant is incompatible, or if torch
            is required but not installed.
    """
    info = detect_torch_env()

    if backend == "cpu":
        # CPU backend works with any torch (or without torch)
        return

    if info.variant == TorchVariant.NOT_INSTALLED:
        raise OSError(
            f"torch is required for backend '{backend}' but is not installed.\n"
            f"Install it with: sagellm install {backend}"
        )

    if not info.is_compatible_with(backend):
        raise OSError(
            f"Backend '{backend}' requires torch variant "
            f"'{backend}' but found '{info.variant.value}' "
            f"(torch {info.torch_version!r}).\n"
            f"Re-install with: sagellm install {backend}"
        )


def get_install_command(backend: str, prefer_mirror: bool = False) -> str:
    """Return the recommended pip install command for the given backend.

    Args:
        backend: One of ``"cuda"``, ``"cuda11"``, ``"cuda12"``, ``"ascend"``,
            ``"kunlun"``, ``"haiguang"``, ``"muxi"``, ``"cpu"``.
        prefer_mirror: When ``True``, prefer a China-accessible mirror URL
            where available.

    Returns:
        A pip command string suitable for shell execution.
    """
    _INSTALL_CMDS: dict[str, str] = {
        "cpu": (
            "pip install torch torchvision torchaudio "
            "--index-url https://download.pytorch.org/whl/cpu"
        ),
        "cuda": (
            "pip install torch torchvision torchaudio "
            "--index-url https://download.pytorch.org/whl/cu121"
        ),
        "cuda11": (
            "pip install torch torchvision torchaudio "
            "--index-url https://download.pytorch.org/whl/cu118"
        ),
        "cuda12": (
            "pip install torch torchvision torchaudio "
            "--index-url https://download.pytorch.org/whl/cu121"
        ),
        "ascend": (
            "# Ascend requires the CANN toolkit pre-installed.\n"
            "# Then install torch_npu from the official Ascend channel:\n"
            "pip install torch==2.1.0 torch_npu==2.1.0r0.post3 "
            "--index-url https://download.pytorch.org/whl/cpu\n"
            "# See: https://www.hiascend.com/document/detail/en/Pytorch/"
        ),
        "kunlun": (
            "# Kunlun requires the Kunlun SDK pre-installed.\n"
            "pip install torch_xmlir  # from Kunlun official channel"
        ),
        "haiguang": (
            "# Haiguang DCU requires the DTK toolkit pre-installed.\n"
            "pip install torch torch_dcu  # from Haiguang official channel"
        ),
        "muxi": (
            "# MThreads Muxi requires the MUSA SDK pre-installed.\n"
            "pip install torch_musa  # from MThreads official channel"
        ),
    }
    cmd = _INSTALL_CMDS.get(backend.lower())
    if cmd is None:
        raise ValueError(f"Unknown backend '{backend}'. Valid options: {sorted(_INSTALL_CMDS)}")
    return cmd
