"""Diagnose the PyTorch / CUDA situation so the UI can explain *why* it's on CPU.

Pure (no Qt). ``gather()`` is the heavy part — it imports torch and shells out to
``nvidia-smi`` — so call it off the GUI thread. ``analyze()`` is fast formatting on
the gathered dict and figures out the most likely cause + concrete fix.
"""

from __future__ import annotations

import re
import shutil
import subprocess
import sys

# pip indexes for the CUDA builds (see README). cu121 needs a driver with CUDA >= 12.1.
CUDA_WHEELS = {
    "cu121": "https://download.pytorch.org/whl/cu121",
    "cu118": "https://download.pytorch.org/whl/cu118",
}

_NO_WINDOW = 0x08000000 if sys.platform == "win32" else 0  # CREATE_NO_WINDOW


def _run_nvidia_smi() -> dict:
    """Probe the NVIDIA driver/GPU via nvidia-smi. Never raises."""
    out: dict = {"found": False, "gpus": [], "driver_version": None, "cuda_driver": None}
    exe = shutil.which("nvidia-smi")
    if not exe and sys.platform == "win32":
        candidate = r"C:\Windows\System32\nvidia-smi.exe"
        exe = candidate if shutil.os.path.isfile(candidate) else None
    if not exe:
        return out
    # GPU names + driver version (robust CSV form).
    try:
        r = subprocess.run(
            [exe, "--query-gpu=name,driver_version", "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=10, creationflags=_NO_WINDOW,
        )
        if r.returncode == 0:
            out["found"] = True
            for line in r.stdout.strip().splitlines():
                parts = [p.strip() for p in line.split(",")]
                if parts and parts[0]:
                    out["gpus"].append(parts[0])
                    if len(parts) > 1 and parts[1]:
                        out["driver_version"] = parts[1]
    except Exception:
        return out
    # Max CUDA version the driver supports (only in the plain header).
    try:
        r2 = subprocess.run(
            [exe], capture_output=True, text=True, timeout=10, creationflags=_NO_WINDOW,
        )
        m = re.search(r"CUDA Version:\s*([\d.]+)", r2.stdout)
        if m:
            out["cuda_driver"] = m.group(1)
    except Exception:
        pass
    return out


def gather() -> dict:
    """Collect torch + NVIDIA facts. Never raises — missing torch/GPU are valid."""
    info: dict = {
        "installed": False,
        "version": None,        # torch.__version__ (e.g. "2.12.0+cpu")
        "built_cuda": None,     # torch.version.cuda (None on a CPU-only build)
        "cuda_available": False,
        "device_name": None,    # the active GPU's name, if any
        "import_error": None,
        # filled by nvidia-smi:
        "nvidia_smi": False,
        "gpus": [],
        "driver_version": None,
        "cuda_driver": None,
    }
    try:
        import torch
    except Exception as exc:
        info["import_error"] = str(exc)
    else:
        info["installed"] = True
        info["version"] = getattr(torch, "__version__", None)
        try:
            info["built_cuda"] = torch.version.cuda
        except Exception:
            info["built_cuda"] = None
        try:
            info["cuda_available"] = bool(torch.cuda.is_available())
        except Exception:
            info["cuda_available"] = False
        if info["cuda_available"]:
            try:
                info["device_name"] = torch.cuda.get_device_name(0)
            except Exception:
                info["device_name"] = None

    smi = _run_nvidia_smi()
    info["nvidia_smi"] = smi["found"]
    info["gpus"] = smi["gpus"]
    info["driver_version"] = smi["driver_version"]
    info["cuda_driver"] = smi["cuda_driver"]
    return info


def _ver_tuple(v: str | None) -> tuple[int, ...]:
    try:
        return tuple(int(x) for x in str(v).split(".")[:2])
    except (ValueError, AttributeError):
        return ()


def recommend_channel(info: dict) -> str:
    """Pick the pip CUDA wheel index that matches the driver (cu118 for older)."""
    cd = _ver_tuple(info.get("cuda_driver"))
    if cd and cd < (12, 1):
        return "cu118"
    return "cu121"


def install_command(channel: str = "cu121") -> str:
    """The pip commands to (re)install the chosen CUDA build.

    Targets the **running interpreter** (``sys.executable -m pip``) so the command
    hits the same venv that runs the app — not whatever ``pip`` is on PATH. (A common
    trap: running bare ``pip`` in a global shell while torch lives in the project venv.)
    """
    index = CUDA_WHEELS.get(channel, CUDA_WHEELS["cu121"])
    py = sys.executable or "python"
    q = f'"{py}"' if " " in py else py
    return (
        f"{q} -m pip uninstall -y torch torchvision\n"
        f"{q} -m pip install torch torchvision --index-url {index}"
    )


def analyze(info: dict) -> dict:
    """Turn the raw facts into {summary, details[list], steps, command}."""
    installed = info.get("installed")
    cuda = info.get("cuda_available")
    built = info.get("built_cuda")
    gpus = info.get("gpus") or []
    smi = info.get("nvidia_smi")
    driver = info.get("driver_version")
    cuda_driver = info.get("cuda_driver")
    channel = recommend_channel(info)
    command = install_command(channel)

    if not installed:
        build_str = "не установлен"
    elif built:
        build_str = f"CUDA {built}"
    else:
        build_str = "CPU-only (+cpu)"
    gpu_str = (
        ", ".join(gpus) if gpus
        else ("не обнаружена" if smi or info.get("nvidia_smi") is False else "nvidia-smi не найден")
    )
    if not gpus and not smi:
        gpu_str = "nvidia-smi не найден (нет драйвера NVIDIA?)"
    details = [
        f"PyTorch: {info.get('version') or 'не установлен'}",
        f"Сборка PyTorch: {build_str}",
        f"CUDA доступна в PyTorch: {'да' if cuda else 'нет'}",
        f"Видеокарта (nvidia-smi): {gpu_str}",
        f"Драйвер NVIDIA: {driver or '—'}",
        f"Макс. CUDA драйвера: {cuda_driver or '—'}",
        f"Интерпретатор (venv): {sys.executable}",
    ]
    if info.get("import_error"):
        details.append(f"Ошибка импорта torch: {info['import_error']}")

    if not installed:
        summary = "PyTorch не установлен — детекция YOLO и DeepMosaics идут на CPU."
        steps = ("Установите PyTorch (CUDA-сборку, если есть NVIDIA-видеокарта):\n\n"
                 + command + "\n\nДля YOLO также: pip install -e \".[yolo]\"")
    elif cuda:
        gpu = info.get("device_name") or (gpus[0] if gpus else "GPU")
        summary = f"Всё в порядке: PyTorch использует CUDA. Активный GPU: {gpu}."
        steps = "GPU уже задействован — ничего делать не нужно."
    elif not built:  # CPU-only torch build — the usual case
        if gpus:
            summary = (
                "Главная причина: установлена CPU-сборка PyTorch (+cpu) — она физически "
                f"не умеет в CUDA. Видеокарта ({gpus[0]}) и драйвер {driver or '?'} на месте, "
                "поэтому достаточно переустановить PyTorch со сборкой CUDA."
            )
            steps = ("Переустановите PyTorch под CUDA, затем перезапустите приложение:\n\n"
                     + command)
        else:
            summary = (
                "Установлена CPU-сборка PyTorch (+cpu), и видеокарта NVIDIA не обнаружена "
                "(nvidia-smi не отвечает). Либо нет NVIDIA GPU, либо не установлен драйвер."
            )
            steps = ("1. Проверьте видеокарту и драйвер: в консоли выполните  nvidia-smi\n"
                     "2. Если NVIDIA GPU есть — переустановите PyTorch под CUDA:\n\n"
                     + command +
                     "\n\nБез NVIDIA GPU всё работает на CPU — просто медленнее.")
    else:  # CUDA-enabled torch build, but CUDA still not available
        if not gpus:
            summary = (
                f"PyTorch собран под CUDA {built}, но видеокарта/драйвер NVIDIA не найдены. "
                "Скорее всего не установлен драйвер NVIDIA или нет GPU."
            )
            steps = "Установите свежий драйвер NVIDIA и перезапустите. Проверка: nvidia-smi."
        elif cuda_driver and _ver_tuple(cuda_driver) < _ver_tuple(built):
            summary = (
                f"Драйвер поддерживает CUDA {cuda_driver}, а PyTorch собран под CUDA {built} — "
                "версия драйвера слишком старая."
            )
            steps = ("Вариант A — обновите драйвер NVIDIA (рекомендуется).\n"
                     "Вариант B — поставьте PyTorch под CUDA вашего драйвера:\n\n" + command)
        else:
            summary = (
                f"PyTorch собран под CUDA {built}, GPU ({gpus[0]}) есть, но CUDA недоступна — "
                "вероятен конфликт версий или повреждённая установка."
            )
            steps = "Переустановите PyTorch под CUDA:\n\n" + command
    return {"summary": summary, "details": details, "steps": steps, "command": command}