# SPDX-License-Identifier: AGPL-3.0-only
# Copyright (c) 2026 sol pbc

"""CUDA GPU discovery and backend selection for the bundled local provider."""

from __future__ import annotations

import logging
import re
import shutil
import subprocess
from dataclasses import dataclass
from enum import StrEnum
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from solstone.think.providers.local_install import CudaServerPin

LOG = logging.getLogger(__name__)

_PROBE_TIMEOUT_S = 10
_ARCH_RE = re.compile(r"\bsm_\d+a?\b")
_CUDA_VERSION_RE = re.compile(r"CUDA Version\s*:?\s*(\d+)(?:\.\d+)?")
MEMORY_SOURCE_UNAVAILABLE = "unavailable"
MEMORY_SOURCE_NVIDIA_VRAM = "nvidia memory.total"
MEMORY_SOURCE_SYSTEM_AVAILABLE = "system MemAvailable (unified memory)"
# The CUDA arch set embedded in the pinned llama.cpp CUDA server image, and the
# minimum driver CUDA major version it requires. The bundled-provider CUDA pin
# reads these as its sole source of truth.
# TODO(AC10): confirm via cuobjdump --list-elf libggml-cuda.so on hardware.
CUDA_EMBEDDED_ARCH_SET: frozenset[str] = frozenset(
    {"sm_86", "sm_89", "sm_120a", "sm_121a"}
)
CUDA_MIN_DRIVER_VERSION = 13


class LocalCudaError(RuntimeError):
    """CUDA local-provider failure with a recovery reason code."""

    def __init__(self, reason_code: str, message: str) -> None:
        super().__init__(message)
        self.reason_code = reason_code


class ArtifactTrust(StrEnum):
    TRUSTED = "trusted"
    ABSENT = "absent"
    UNAVAILABLE = "unavailable"


@dataclass(frozen=True)
class NvidiaProbe:
    index: int | None
    compute_cap: str | None
    driver_cuda_version: int | None
    vram_mib: int | None
    tiering_memory_mib: int | None
    memory_source: str
    detected: bool


@dataclass(frozen=True)
class BackendChoice:
    backend: str
    reason: str


def _base_arch(arch: str) -> str:
    return arch[:-1] if arch.endswith("a") else arch


def _compute_cap_to_arch(cap: str) -> str | None:
    parts = cap.strip().split(".")
    if len(parts) != 2 or not all(part.isdigit() for part in parts):
        return None
    return f"sm_{parts[0]}{parts[1]}"


def _undetected() -> NvidiaProbe:
    return NvidiaProbe(
        index=None,
        compute_cap=None,
        driver_cuda_version=None,
        vram_mib=None,
        tiering_memory_mib=None,
        memory_source=MEMORY_SOURCE_UNAVAILABLE,
        detected=False,
    )


def _parse_nvidia_smi_row(
    line: str,
) -> tuple[int, str, str | None, int | None] | None:
    fields = [field.strip() for field in line.split(",")]
    if len(fields) != 5:
        return None
    try:
        index = int(fields[0])
    except ValueError:
        return None

    gpu_name = fields[1]
    compute_cap = _compute_cap_to_arch(fields[2])
    try:
        vram_mib = int(fields[4])
    except ValueError:
        vram_mib = None
    return index, gpu_name, compute_cap, vram_mib


def has_unified_memory_name(gpu_name: str) -> bool:
    return "GB10" in gpu_name.upper()


def _parse_memavailable_mib(meminfo_text: str) -> int | None:
    for line in meminfo_text.splitlines():
        key, sep, value = line.partition(":")
        if sep != ":" or key != "MemAvailable":
            continue
        parts = value.strip().split()
        if not parts:
            return None
        try:
            value_kib = int(parts[0])
        except ValueError:
            return None
        if len(parts) > 1 and parts[1] != "kB":
            return None
        return value_kib // 1024
    return None


def _read_linux_memavailable_mib(
    meminfo_path: Path = Path("/proc/meminfo"),
) -> int | None:
    try:
        return _parse_memavailable_mib(meminfo_path.read_text(encoding="utf-8"))
    except OSError as exc:
        LOG.warning("Linux MemAvailable probe could not read %s: %s", meminfo_path, exc)
        return None


def _probe_driver_cuda_version() -> int | None:
    # TODO(AC10): driver_version→CUDA-version threshold fallback (candidate floor
    # 580.65.06) is unverified; None here fails safe to Vulkan.
    try:
        completed = subprocess.run(
            ["nvidia-smi", "--version"],
            capture_output=True,
            text=True,
            timeout=_PROBE_TIMEOUT_S,
            check=False,
        )
    except subprocess.TimeoutExpired:
        LOG.warning(
            "NVIDIA driver CUDA version probe timed out after %.0fs", _PROBE_TIMEOUT_S
        )
        return None
    except OSError as exc:
        LOG.warning("NVIDIA driver CUDA version probe could not start: %s", exc)
        return None

    if completed.returncode != 0:
        LOG.warning(
            "NVIDIA driver CUDA version probe exited with status %s",
            completed.returncode,
        )
        return None

    match = _CUDA_VERSION_RE.search(completed.stdout)
    if match is None:
        LOG.warning("NVIDIA driver CUDA version probe returned unparseable output")
        return None
    return int(match.group(1))


def detect_nvidia_unified_memory() -> bool:
    if shutil.which("nvidia-smi") is None:
        return False

    try:
        completed = subprocess.run(
            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
            capture_output=True,
            text=True,
            timeout=_PROBE_TIMEOUT_S,
            check=False,
        )
    except subprocess.TimeoutExpired:
        LOG.warning(
            "NVIDIA unified-memory probe timed out after %.0fs", _PROBE_TIMEOUT_S
        )
        return False
    except OSError as exc:
        LOG.warning("NVIDIA unified-memory probe could not start: %s", exc)
        return False

    if completed.returncode != 0:
        LOG.warning(
            "NVIDIA unified-memory probe exited with status %s", completed.returncode
        )
        return False

    first_row = next(
        (line.strip() for line in completed.stdout.splitlines() if line.strip()),
        None,
    )
    if first_row is None:
        LOG.warning("NVIDIA unified-memory probe returned no rows")
        return False
    return has_unified_memory_name(first_row)


def probe_nvidia_gpu() -> NvidiaProbe:
    if shutil.which("nvidia-smi") is None:
        return _undetected()

    try:
        completed = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=index,name,compute_cap,driver_version,memory.total",
                "--format=csv,noheader,nounits",
            ],
            capture_output=True,
            text=True,
            timeout=_PROBE_TIMEOUT_S,
            check=False,
        )
    except subprocess.TimeoutExpired:
        LOG.warning("NVIDIA GPU probe timed out after %.0fs", _PROBE_TIMEOUT_S)
        return _undetected()
    except OSError as exc:
        LOG.warning("NVIDIA GPU probe could not start: %s", exc)
        return _undetected()

    if completed.returncode != 0:
        LOG.warning("NVIDIA GPU probe exited with status %s", completed.returncode)
        return _undetected()

    first_row = next(
        (line.strip() for line in completed.stdout.splitlines() if line.strip()),
        None,
    )
    if first_row is None:
        LOG.warning("NVIDIA GPU probe returned no rows")
        return _undetected()

    parsed = _parse_nvidia_smi_row(first_row)
    if parsed is None:
        LOG.warning("NVIDIA GPU probe returned invalid CSV")
        return _undetected()

    index, gpu_name, compute_cap, vram_mib = parsed
    tiering_memory_mib = vram_mib
    memory_source = MEMORY_SOURCE_NVIDIA_VRAM
    if vram_mib is None and has_unified_memory_name(gpu_name):
        tiering_memory_mib = _read_linux_memavailable_mib()
        memory_source = (
            MEMORY_SOURCE_SYSTEM_AVAILABLE
            if tiering_memory_mib is not None
            else MEMORY_SOURCE_UNAVAILABLE
        )
    elif vram_mib is None:
        memory_source = MEMORY_SOURCE_UNAVAILABLE
    return NvidiaProbe(
        index=index,
        compute_cap=compute_cap,
        driver_cuda_version=_probe_driver_cuda_version(),
        vram_mib=vram_mib,
        tiering_memory_mib=tiering_memory_mib,
        memory_source=memory_source,
        detected=True,
    )


def select_local_backend(
    probe: NvidiaProbe,
    arch_set: frozenset[str],
    cuda_version: int,
    trust: ArtifactTrust,
    *,
    persisted_installed_cuda: bool,
) -> BackendChoice:
    hardware_rejection = _hardware_backend_rejection(probe, arch_set, cuda_version)
    if hardware_rejection is not None:
        return hardware_rejection

    assert probe.compute_cap is not None
    assert probe.driver_cuda_version is not None
    cuda_reason = (
        f"compute_cap {probe.compute_cap} covered; "
        f"driver CUDA {probe.driver_cuda_version} >= {cuda_version}"
    )
    if trust == ArtifactTrust.TRUSTED or (
        trust == ArtifactTrust.UNAVAILABLE and persisted_installed_cuda
    ):
        return BackendChoice("cuda", cuda_reason)

    return BackendChoice(
        "vulkan",
        (f"{cuda_reason}; no trusted CUDA runtime artifact present"),
    )


def _hardware_backend_rejection(
    probe: NvidiaProbe,
    arch_set: frozenset[str],
    cuda_version: int,
) -> BackendChoice | None:
    if not probe.detected:
        return BackendChoice("vulkan", "no NVIDIA GPU detected")
    if probe.compute_cap is None:
        return BackendChoice("vulkan", "NVIDIA compute capability unreadable")

    base_arches = {_base_arch(arch) for arch in arch_set}
    if _base_arch(probe.compute_cap) not in base_arches:
        return BackendChoice(
            "vulkan",
            f"compute_cap {probe.compute_cap} not in CUDA image arch set",
        )

    if probe.driver_cuda_version is None:
        return BackendChoice("vulkan", "driver CUDA version unreadable")
    if probe.driver_cuda_version < cuda_version:
        return BackendChoice(
            "vulkan",
            f"driver CUDA {probe.driver_cuda_version} < required {cuda_version}",
        )
    return None


def resolve_local_backend(pin: CudaServerPin) -> BackendChoice:
    probe = probe_nvidia_gpu()
    hardware_rejection = _hardware_backend_rejection(
        probe,
        pin.embedded_arch_set,
        pin.cuda_version,
    )
    if hardware_rejection is not None:
        return hardware_rejection

    from solstone.think.providers import local_install

    trust = local_install.probe_cuda_runtime_artifact_trust(pin)
    persisted_installed_cuda = local_install.has_persisted_installed_cuda_target()
    return select_local_backend(
        probe,
        pin.embedded_arch_set,
        pin.cuda_version,
        trust,
        persisted_installed_cuda=persisted_installed_cuda,
    )


def parse_embedded_arch_set(cuobjdump_list_elf_text: str) -> frozenset[str]:
    # TODO(AC10): live cuobjdump --list-elf on the extracted libggml-cuda.so
    # runs on hardware; in-lode uses a synthetic fixture.
    return frozenset(_ARCH_RE.findall(cuobjdump_list_elf_text))


def verify_cuda_pin_arch_set(text: str, declared: frozenset[str]) -> None:
    actual = parse_embedded_arch_set(text)
    if actual != declared:
        raise LocalCudaError(
            "arch_set_mismatch",
            f"CUDA embedded arch set mismatch: declared={sorted(declared)}, actual={sorted(actual)}",
        )


__all__ = [
    "ArtifactTrust",
    "BackendChoice",
    "CUDA_EMBEDDED_ARCH_SET",
    "CUDA_MIN_DRIVER_VERSION",
    "LocalCudaError",
    "MEMORY_SOURCE_NVIDIA_VRAM",
    "MEMORY_SOURCE_SYSTEM_AVAILABLE",
    "MEMORY_SOURCE_UNAVAILABLE",
    "NvidiaProbe",
    "detect_nvidia_unified_memory",
    "has_unified_memory_name",
    "parse_embedded_arch_set",
    "probe_nvidia_gpu",
    "resolve_local_backend",
    "select_local_backend",
    "verify_cuda_pin_arch_set",
]
