"""NVIDIA GPU facts for launch decisions (best-effort, via nvidia-smi).

sglang-omni model pipelines carry GPU-architecture constraints the app
must respect when it builds a server spec (e.g. ZONOS2's default pipeline
quantizes its MoE experts to FP8 — a Triton kernel that only compiles on
compute capability 8.9+). The GPU's compute capability is read from
``nvidia-smi`` rather than torch so the app venv needs no CUDA stack, and
every answer here is cached: the hardware cannot change mid-process.

Everything is best-effort: when nvidia-smi is absent, errors out, or
reports something unparsable the callers get None and keep upstream
defaults instead of second-guessing the environment.
"""

import shutil
import subprocess
from functools import lru_cache
from typing import Optional, Tuple


@lru_cache(maxsize=1)
def _query() -> Optional[Tuple[str, str]]:
    """(name, "M.m" compute cap) for GPU 0, or None when unanswerable."""
    if shutil.which("nvidia-smi") is None:
        return None
    try:
        proc = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,compute_cap",
             "--format=csv,noheader,nounits", "-i", "0"],
            capture_output=True, text=True, timeout=10)
    except (OSError, subprocess.SubprocessError):
        return None
    first = proc.stdout.strip().splitlines()[:1]
    if proc.returncode != 0 or not first:
        return None
    fields = [field.strip() for field in first[0].split(",")]
    if len(fields) < 2:
        return None
    return fields[0], fields[1]


def compute_capability() -> Optional[Tuple[int, int]]:
    """The first NVIDIA GPU's (major, minor) compute capability, or None.

    GPU 0 is what the managed sglang-omni pipeline stages bind to. An
    unparsable capability string counts as unanswerable."""
    answer = _query()
    if answer is None:
        return None
    try:
        major, minor = (int(part) for part in answer[1].split("."))
    except ValueError:
        return None
    return major, minor


def describe() -> Optional[str]:
    """A human-readable "NAME (compute capability M.m)" line, or None."""
    answer = _query()
    if answer is None:
        return None
    return f"{answer[0]} (compute capability {answer[1]})"