"""Curated list of models supported by the network with VRAM requirements."""

from __future__ import annotations

import json
from dataclasses import dataclass
from importlib.resources import files
from pathlib import Path


@dataclass
class ModelPreset:
    name: str
    hf_repo: str
    num_layers: int
    # VRAM in GB at each quantization level (None = too large to quantize this way)
    vram_nf4: float
    vram_int8: float
    vram_bf16: float
    description: str
    metadata: dict | None = None
    aliases: tuple[str, ...] = ()

    def vram_for_quant(self, quant: str) -> float:
        """Return VRAM requirement in GB for the given quantization."""
        q = quant.lower().replace("bfloat16", "bf16")
        if q == "nf4":
            return self.vram_nf4
        if q in ("int8", "int8"):
            return self.vram_int8
        if q in ("bf16", "bfloat16"):
            return self.vram_bf16
        raise ValueError(f"unknown quantization: {quant!r}")

    def fits_vram(self, available_gb: float, quant: str) -> bool:
        return self.vram_for_quant(quant) <= available_gb

    def recommended_quant(self, available_gb: float) -> str | None:
        """Return the highest-quality quantization that fits available VRAM, or None."""
        if self.vram_bf16 <= available_gb:
            return "bf16"
        if self.vram_int8 <= available_gb:
            return "int8"
        if self.vram_nf4 <= available_gb:
            return "nf4"
        return None


def _load_model_metadata() -> dict[str, dict]:
    try:
        raw = files("meshnet_node").joinpath("model_metadata.json").read_text()
        data = json.loads(raw)
    except Exception:
        return {}
    models = data.get("models", {})
    if not isinstance(models, dict):
        return {}
    return {
        str(repo): metadata
        for repo, metadata in models.items()
        if isinstance(metadata, dict)
    }


_MODEL_METADATA = _load_model_metadata()


def _local_model_path(hf_repo: str, cache_dir: Path | None) -> Path | None:
    if cache_dir is None:
        return None
    if (cache_dir / "config.json").exists():
        return cache_dir
    candidate = cache_dir / hf_repo.split("/")[-1]
    if (candidate / "config.json").exists():
        return candidate
    return None


CURATED_MODELS: list[ModelPreset] = [
    ModelPreset(
        name="Qwen2.5-0.5B-Instruct",
        hf_repo="Qwen/Qwen2.5-0.5B-Instruct",
        num_layers=24,
        vram_nf4=0.4,
        vram_int8=0.6,
        vram_bf16=1.0,
        description="Smallest no-gating model — great for testing, ~1 GB",
    ),
    ModelPreset(
        name="Qwen2.5-1.5B-Instruct",
        hf_repo="Qwen/Qwen2.5-1.5B-Instruct",
        num_layers=28,
        vram_nf4=1.0,
        vram_int8=1.8,
        vram_bf16=3.2,
        description="Fast no-gating model — good quality, ~3 GB",
    ),
    ModelPreset(
        name="Llama-3-70B-Instruct",
        hf_repo="meta-llama/Meta-Llama-3-70B-Instruct",
        num_layers=80,
        vram_nf4=18.0,
        vram_int8=40.0,
        vram_bf16=140.0,
        description="Meta's flagship 70B instruction model",
    ),
    ModelPreset(
        name="Qwen2.5-72B-Instruct",
        hf_repo="Qwen/Qwen2.5-72B-Instruct",
        num_layers=80,
        vram_nf4=19.0,
        vram_int8=41.0,
        vram_bf16=145.0,
        description="Alibaba's 72B multilingual instruction model",
    ),
    ModelPreset(
        name="Mixtral-8x7B-Instruct",
        hf_repo="mistralai/Mixtral-8x7B-Instruct-v0.1",
        num_layers=32,
        vram_nf4=7.0,
        vram_int8=14.0,
        vram_bf16=27.0,
        description="Mistral's sparse MoE — fast and efficient",
    ),
    ModelPreset(
        name="Llama-3-8B-Instruct",
        hf_repo="meta-llama/Meta-Llama-3-8B-Instruct",
        num_layers=32,  # gated repo — requires HF login
        vram_nf4=4.5,
        vram_int8=8.5,
        vram_bf16=16.0,
        description="Meta's compact 8B model — good for low-VRAM nodes",
    ),
    ModelPreset(
        name="Phi-3-medium-128k",
        hf_repo="microsoft/Phi-3-medium-128k-instruct",
        num_layers=40,
        vram_nf4=4.0,
        vram_int8=8.0,
        vram_bf16=15.0,
        description="Microsoft's efficient 14B model with 128k context",
    ),
    ModelPreset(
        name="Gemma-2-27B-IT",
        hf_repo="google/gemma-2-27b-it",
        num_layers=46,
        vram_nf4=10.0,
        vram_int8=20.0,
        vram_bf16=54.0,
        description="Google's 27B instruction-tuned model",
    ),
    ModelPreset(
        name="DeepSeek-V2-Lite-Chat",
        hf_repo="deepseek-ai/DeepSeek-V2-Lite-Chat",
        num_layers=27,
        vram_nf4=5.0,
        vram_int8=9.0,
        vram_bf16=16.0,
        description="DeepSeek's efficient MoE — strong coding + reasoning",
    ),
    ModelPreset(
        name="Kimi-K2.7-Code",
        hf_repo="unsloth/Kimi-K2.7-Code",
        num_layers=61,
        vram_nf4=500.0,
        vram_int8=1000.0,
        vram_bf16=2000.0,
        description="Large coding-focused MoE model",
        metadata=_MODEL_METADATA.get("unsloth/Kimi-K2.7-Code"),
    ),
    ModelPreset(
        name="Qwen3.6-27B",
        hf_repo="Qwen/Qwen3.6-27B",
        num_layers=64,
        vram_nf4=13.0,
        vram_int8=26.0,
        vram_bf16=52.0,
        description="Qwen 27B hybrid linear-attention model",
        aliases=("qwen3.6-27b",),
    ),
]


def resolve_model_alias(value: str) -> ModelPreset | None:
    """Resolve a curated name, repository, or alias case-insensitively."""
    normalized = value.strip().casefold()
    for model in CURATED_MODELS:
        candidates = (model.name, model.hf_repo, *model.aliases)
        if normalized in {candidate.casefold() for candidate in candidates}:
            return model
    return None


def layers_from_config(cfg) -> int | None:
    """Extract the transformer layer count from a HuggingFace config object.

    Composite configs (vision-language, some MoE) nest the decoder settings in
    ``text_config`` — e.g. Qwen3.5-MoE has no top-level ``num_hidden_layers``.
    """
    candidates = [cfg]
    get_text_config = getattr(cfg, "get_text_config", None)
    if callable(get_text_config):
        try:
            candidates.append(get_text_config())
        except Exception:
            pass
    nested = getattr(cfg, "text_config", None)
    if nested is not None:
        candidates.append(nested)
    for candidate in candidates:
        for attr in ("num_hidden_layers", "num_layers", "n_layer", "n_layers"):
            value = getattr(candidate, attr, None)
            if value is not None:
                return int(value)
    return None


def detect_num_layers(hf_repo: str) -> int | None:
    """Return num_hidden_layers from HuggingFace config.json (downloads ~1 KB only)."""
    # Check curated list first (no network call)
    for m in CURATED_MODELS:
        if m.hf_repo == hf_repo:
            return m.num_layers
    try:
        from transformers import AutoConfig  # type: ignore[import]
        cfg = AutoConfig.from_pretrained(hf_repo)
        return layers_from_config(cfg)
    except Exception:
        return None


def model_metadata_for(
    hf_repo: str,
    num_layers: int | None = None,
    cache_dir: Path | None = None,
) -> dict:
    """Return operator-facing model metadata for a HuggingFace repo."""
    for model in CURATED_MODELS:
        if model.hf_repo == hf_repo:
            metadata = dict(model.metadata or {})
            metadata.setdefault("num_layers", model.num_layers)
            return metadata

    metadata: dict = {}
    if num_layers is not None:
        metadata["num_layers"] = num_layers
    try:
        from transformers import AutoConfig  # type: ignore[import]

        local_model = _local_model_path(hf_repo, cache_dir)
        load_source = str(local_model) if local_model is not None else hf_repo
        cfg = AutoConfig.from_pretrained(
            load_source,
            cache_dir=str(cache_dir) if cache_dir is not None and local_model is None else None,
        )
        # Composite configs (VLM/MoE) nest decoder fields in text_config.
        text_cfg = getattr(cfg, "text_config", None) or cfg
        for attr, key in (
            ("model_type", "architecture"),
            ("num_hidden_layers", "num_layers"),
            ("hidden_size", "hidden_size"),
            ("num_attention_heads", "attention_heads"),
            ("vocab_size", "vocabulary_size"),
            ("max_position_embeddings", "context_length"),
        ):
            value = getattr(cfg, attr, None)
            if value is None:
                value = getattr(text_cfg, attr, None)
            if value is not None:
                metadata[key] = value
        if "num_layers" not in metadata:
            layers = layers_from_config(cfg)
            if layers is not None:
                metadata["num_layers"] = layers
    except Exception:
        pass
    return metadata


def browse_hf_hub(top_n: int = 20) -> list[dict]:
    """Fetch top downloaded text-generation models from HuggingFace Hub."""
    try:
        from huggingface_hub import list_models  # type: ignore[import]

        models = list(
            list_models(
                pipeline_tag="text-generation",
                library="transformers",
                sort="downloads",
                direction=-1,
                limit=top_n,
            )
        )
        return [
            {
                "repo": m.id,
                "downloads": getattr(m, "downloads", 0) or 0,
            }
            for m in models
        ]
    except Exception as exc:
        raise RuntimeError(f"HuggingFace Hub lookup failed: {exc}") from exc
