"""LLMHosts model registry — maps friendly names to GGUF files on HuggingFace."""

from __future__ import annotations

from pathlib import Path
from typing import Any

_MODELS_DIR = str(Path.home() / ".llmhosts" / "models")

MODELS: dict[str, dict[str, Any]] = {
    "qwen2.5:0.5b": {
        "repo": "Qwen/Qwen2.5-0.5B-Instruct-GGUF",
        "file": "qwen2.5-0.5b-instruct-q4_k_m.gguf",
        "tokenizer_repo": "Qwen/Qwen2.5-0.5B-Instruct",
        "size_gb": 0.4,
        "min_vram_gb": 1,
        "architecture": "llama",
        "context_length": 32768,
    },
    "qwen2.5:3b": {
        "repo": "Qwen/Qwen2.5-3B-Instruct-GGUF",
        "file": "qwen2.5-3b-instruct-q4_k_m.gguf",
        "tokenizer_repo": "Qwen/Qwen2.5-3B-Instruct",
        "size_gb": 2.0,
        "min_vram_gb": 3,
        "architecture": "llama",
        "context_length": 32768,
    },
    "qwen2.5:7b": {
        "repo": "Qwen/Qwen2.5-7B-Instruct-GGUF",
        "file": "qwen2.5-7b-instruct-q4_k_m.gguf",
        "tokenizer_repo": "Qwen/Qwen2.5-7B-Instruct",
        "size_gb": 4.4,
        "min_vram_gb": 6,
        "architecture": "llama",
        "context_length": 32768,
    },
    "qwen2.5:14b": {
        "repo": "Qwen/Qwen2.5-14B-Instruct-GGUF",
        "file": "qwen2.5-14b-instruct-q4_k_m.gguf",
        "tokenizer_repo": "Qwen/Qwen2.5-14B-Instruct",
        "size_gb": 8.7,
        "min_vram_gb": 10,
        "architecture": "llama",
        "context_length": 32768,
    },
    "llama3.1:8b": {
        "repo": "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF",
        "file": "Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
        "tokenizer_repo": "meta-llama/Llama-3.1-8B-Instruct",
        "size_gb": 4.9,
        "min_vram_gb": 6,
        "architecture": "llama",
        "context_length": 131072,
    },
    "llama3.3:70b": {
        "repo": "bartowski/Llama-3.3-70B-Instruct-GGUF",
        "file": "Llama-3.3-70B-Instruct-Q4_K_M.gguf",
        "tokenizer_repo": "meta-llama/Llama-3.3-70B-Instruct",
        "size_gb": 42.0,
        "min_vram_gb": 44,
        "architecture": "llama",
        "context_length": 131072,
    },
    "mistral:7b": {
        "repo": "TheBloke/Mistral-7B-Instruct-v0.3-GGUF",
        "file": "mistral-7b-instruct-v0.3.Q4_K_M.gguf",
        "tokenizer_repo": "mistralai/Mistral-7B-Instruct-v0.3",
        "size_gb": 4.4,
        "min_vram_gb": 6,
        "architecture": "llama",
        "context_length": 32768,
    },
    "deepseek-r1:7b": {
        "repo": "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF",
        "file": "DeepSeek-R1-Distill-Qwen-7B-Q4_K_M.gguf",
        "tokenizer_repo": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
        "size_gb": 4.7,
        "min_vram_gb": 6,
        "architecture": "llama",
        "context_length": 32768,
    },
}


def get_model(name: str) -> dict[str, Any] | None:
    """Get model info by friendly name."""
    return MODELS.get(name)


def list_models() -> list[str]:
    """List all available model names."""
    return list(MODELS.keys())


def recommend_for_vram(vram_mb: int) -> dict[str, Any] | None:
    """Recommend the largest model that fits in available VRAM."""
    vram_gb = vram_mb / 1024
    candidates = sorted(MODELS.items(), key=lambda x: x[1]["size_gb"], reverse=True)
    for name, info in candidates:
        if info["min_vram_gb"] <= vram_gb:
            return {**info, "name": name}
    # Fall back to smallest model for CPU-only or very low VRAM
    smallest = min(MODELS.items(), key=lambda x: x[1]["size_gb"])
    return {**smallest[1], "name": smallest[0]}


def model_path(name: str) -> str:
    """Return the expected local path for a model's GGUF file."""
    info = MODELS.get(name)
    if info is None:
        raise ValueError(f"Unknown model: {name}")
    return str(Path(_MODELS_DIR) / info["file"])


def models_dir() -> str:
    """Return the models storage directory."""
    return _MODELS_DIR
