"""Hardware detection and model provisioning helpers."""

from __future__ import annotations

import os
import platform
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path

from rc.config import RC_DIR
from rc.inference.catalog import (
    LLAMA_CPP_RUNTIME,
    MLX_RUNTIME,
    ModelSpec,
    choose_best_model,
    get_model_spec,
)

MODEL_CACHE_DIR = RC_DIR / "models"


@dataclass(frozen=True)
class HardwareProfile:
    accelerator: str
    gpu_name: str | None
    total_vram_gb: float | None
    total_memory_gb: float | None
    gpu_count: int
    platform_name: str


def detect_hardware() -> HardwareProfile:
    """Inspect the local machine and return a provider-relevant profile."""
    cuda_profile = _detect_with_torch() or _detect_with_nvidia_smi()
    if cuda_profile is not None:
        return cuda_profile

    if platform.system() == "Darwin":
        return HardwareProfile(
            accelerator="mps",
            gpu_name="Apple Silicon",
            total_vram_gb=None,
            total_memory_gb=_detect_system_memory_gb(),
            gpu_count=1,
            platform_name=platform.platform(),
        )

    return HardwareProfile(
        accelerator="cpu",
        gpu_name=None,
        total_vram_gb=None,
        total_memory_gb=_detect_system_memory_gb(),
        gpu_count=0,
        platform_name=platform.platform(),
    )


def auto_select_model() -> tuple[HardwareProfile, ModelSpec]:
    """Choose the best supported model for the detected hardware."""
    profile = detect_hardware()
    available_memory = profile.total_vram_gb or profile.total_memory_gb
    model = choose_best_model(available_memory)
    return profile, model


def download_model(repo_id: str) -> Path:
    """Download a supported model locally and return the file or directory path."""
    try:
        from huggingface_hub import hf_hub_download, snapshot_download
    except ImportError as exc:
        raise RuntimeError(
            "Model provisioning requires huggingface_hub. Reinstall the CLI."
        ) from exc

    spec = get_model_spec(repo_id)
    if spec is None:
        raise RuntimeError(f"Unsupported model '{repo_id}'.")

    model_dir = MODEL_CACHE_DIR / _safe_repo_name(spec.repo_id)
    model_dir.mkdir(parents=True, exist_ok=True)

    if spec.runtime == MLX_RUNTIME:
        snapshot_download(
            repo_id=spec.repo_id,
            local_dir=str(model_dir),
            local_dir_use_symlinks=False,
        )
        return model_dir

    if spec.filename is None:
        raise RuntimeError(f"Model '{repo_id}' is missing a downloadable filename.")

    model_path = hf_hub_download(
        repo_id=spec.repo_id,
        filename=spec.filename,
        local_dir=str(model_dir),
        local_dir_use_symlinks=False,
    )
    return Path(model_path)


def ensure_model_downloaded(repo_id: str, current_path: str = "") -> Path:
    """Return a local model path, downloading it if needed."""
    spec = get_model_spec(repo_id)
    if spec is None:
        raise RuntimeError(f"Unsupported model '{repo_id}'.")

    existing_path = Path(current_path).expanduser() if current_path else None
    if existing_path and existing_path.exists():
        if spec.runtime == MLX_RUNTIME and existing_path.is_dir() and any(existing_path.iterdir()):
            return existing_path
        if (
            spec.runtime == LLAMA_CPP_RUNTIME
            and existing_path.is_file()
            and existing_path.name == spec.filename
        ):
            return existing_path

    cache_dir = MODEL_CACHE_DIR / _safe_repo_name(spec.repo_id)
    if spec.runtime == MLX_RUNTIME and cache_dir.is_dir() and any(cache_dir.iterdir()):
        return cache_dir
    if spec.runtime == LLAMA_CPP_RUNTIME and spec.filename:
        model_path = cache_dir / spec.filename
        if model_path.exists():
            return model_path
    return download_model(repo_id)


def provider_runtime_supported(profile: HardwareProfile) -> bool:
    """The provider stack supports CPU/CUDA via llama.cpp and Apple Silicon via MLX."""
    return True


def provisioning_summary(profile: HardwareProfile, model: ModelSpec) -> str:
    """Human-readable summary for CLI output."""
    device_memory = (
        f"{profile.total_vram_gb:.1f} GB VRAM"
        if profile.total_vram_gb is not None
        else (
            f"{profile.total_memory_gb:.1f} GB system RAM"
            if profile.total_memory_gb is not None
            else "memory unknown"
        )
    )
    gpu_name = profile.gpu_name or "No dedicated GPU detected"
    return (
        f"Accelerator: {profile.accelerator}\n"
        f"GPU:         {gpu_name}\n"
        f"Capacity:    {device_memory}\n"
        f"Selected:    {model.label}\n"
        f"Runtime:     {model.runtime}\n"
        f"Repo:        {model.repo_id}\n"
        f"File:        {model.filename or '(MLX snapshot)'}"
    )


def _detect_with_torch() -> HardwareProfile | None:
    try:
        import torch
    except ImportError:
        return None

    if not torch.cuda.is_available():
        return None

    gpu_count = torch.cuda.device_count()
    names: list[str] = []
    total_vram_bytes = 0
    for idx in range(gpu_count):
        props = torch.cuda.get_device_properties(idx)
        names.append(props.name)
        total_vram_bytes += int(props.total_memory)

    return HardwareProfile(
        accelerator="cuda",
        gpu_name=", ".join(names),
        total_vram_gb=total_vram_bytes / (1024**3),
        total_memory_gb=_detect_system_memory_gb(),
        gpu_count=gpu_count,
        platform_name=platform.platform(),
    )


def _detect_with_nvidia_smi() -> HardwareProfile | None:
    if shutil.which("nvidia-smi") is None:
        return None

    try:
        result = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=name,memory.total",
                "--format=csv,noheader,nounits",
            ],
            check=True,
            capture_output=True,
            text=True,
        )
    except (OSError, subprocess.CalledProcessError):
        return None

    names: list[str] = []
    total_vram_mb = 0
    for line in result.stdout.strip().splitlines():
        if not line.strip():
            continue
        name, memory_mb = [part.strip() for part in line.split(",", 1)]
        names.append(name)
        total_vram_mb += int(memory_mb)

    if not names:
        return None

    return HardwareProfile(
        accelerator="cuda",
        gpu_name=", ".join(names),
        total_vram_gb=total_vram_mb / 1024,
        total_memory_gb=_detect_system_memory_gb(),
        gpu_count=len(names),
        platform_name=platform.platform(),
    )


def _detect_system_memory_gb() -> float | None:
    try:
        import psutil

        return psutil.virtual_memory().total / (1024**3)
    except ImportError:
        pass

    if hasattr(os, "sysconf"):
        try:
            page_size = os.sysconf("SC_PAGE_SIZE")
            page_count = os.sysconf("SC_PHYS_PAGES")
            return (page_size * page_count) / (1024**3)
        except (OSError, ValueError):
            return None

    return None


def _safe_repo_name(repo_id: str) -> str:
    return repo_id.replace("/", "--")
