"""
HuggingFace model downloader.

Handles:
  - Fetching repo metadata (files, model card, etc.)
  - Detecting model type (GGUF vs Transformer)
  - Downloading individual files or full snapshots
  - Quantisation selection for GGUF repos
"""

from __future__ import annotations

import os
import re
from pathlib import Path
from typing import Any

# huggingface_hub is a hard dependency
from huggingface_hub import HfApi, hf_hub_download, snapshot_download
from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError


# ── input normalisation ───────────────────────────────────────────────────────

# Unicode characters that look like a hyphen-minus but aren't ASCII 0x2D.
# Copy-pasting model names from web pages often substitutes these.
_UNICODE_DASHES = (
    "\u2010"  # HYPHEN
    "\u2011"  # NON-BREAKING HYPHEN
    "\u2012"  # FIGURE DASH
    "\u2013"  # EN DASH
    "\u2014"  # EM DASH
    "\u2015"  # HORIZONTAL BAR
    "\u00ad"  # SOFT HYPHEN
    "\uff0d"  # FULLWIDTH HYPHEN-MINUS
    "\u2212"  # MINUS SIGN
)


def normalize_repo_id(repo_id: str) -> str:
    """Replace Unicode lookalike dashes with ASCII hyphen-minus.

    Copy-pasting model names from web pages or docs sometimes introduces
    U+2011 NON-BREAKING HYPHEN and similar characters that pass HuggingFace's
    validate_repo_id check.
    """
    for ch in _UNICODE_DASHES:
        repo_id = repo_id.replace(ch, "-")
    return repo_id.strip()


# ── constants ─────────────────────────────────────────────────────────────────

# File extensions that mean "this is a GGUF repo"
GGUF_EXTENSIONS = {".gguf"}

# Transformer checkpoint extensions
TRANSFORMER_EXTENSIONS = {
    ".safetensors", ".bin", ".pt", ".pth",
    ".msgpack", ".flax",
}

# Files we never want to download (large, framework-specific extras)
SKIP_PATTERNS = re.compile(
    r"(\.git/|\.gitattributes|flax_model|tf_model|model\.onnx)",
    re.IGNORECASE,
)

# Quantisation quality ordering (higher index = higher quality / larger file)
QUANT_ORDER = [
    "IQ1_S", "IQ1_M", "IQ2_XXS", "IQ2_XS", "IQ2_S", "IQ2_M",
    "Q2_K", "IQ3_XXS", "IQ3_XS", "Q3_K_S", "IQ3_S", "IQ3_M",
    "Q3_K_M", "Q3_K_L", "IQ4_XS", "IQ4_NL", "Q4_0", "Q4_1",
    "Q4_K_S", "Q4_K_M", "Q5_0", "Q5_1", "Q5_K_S", "Q5_K_M",
    "Q6_K", "Q8_0", "F16", "BF16", "F32",
]


# ── repo inspection ────────────────────────────────────────────────────────────

class ModelInfo:
    """Lightweight wrapper around the HF repo info + file list."""

    def __init__(self, repo_id: str, api: HfApi, token: str | None):
        self.repo_id = repo_id
        self.token = token

        try:
            self._info = api.model_info(repo_id, token=token)
        except GatedRepoError as exc:
            raise ValueError(
                f"Model {repo_id!r} is gated. "
                "Accept the licence on https://huggingface.co and set HF_TOKEN."
            ) from exc
        except RepositoryNotFoundError as exc:
            raise ValueError(
                f"Model not found on HuggingFace: {repo_id!r}"
            ) from exc

        self._siblings = self._info.siblings or []

    # -- properties -----------------------------------------------------------

    @property
    def all_files(self) -> list[str]:
        """Return filenames of all files in the repo."""
        return [sibling.rfilename for sibling in self._siblings]

    @property
    def gguf_files(self) -> list[dict[str, Any]]:
        """Return list of {filename, size, quant} dicts for GGUF files."""
        result = []
        for sibling in self._siblings:
            if not sibling.rfilename.lower().endswith(".gguf"):
                continue
            quant = _parse_quant(sibling.rfilename)
            result.append({
                "filename": sibling.rfilename,
                "size": sibling.size or 0,
                "quant": quant,
            })
        # Sort by quality (ascending)
        result.sort(
            key=lambda x: (
                QUANT_ORDER.index(x["quant"]) if x["quant"] in QUANT_ORDER else -1
            )
        )
        return result

    @property
    def is_gguf(self) -> bool:
        """Return True if the repo contains GGUF files."""
        return any(fname.lower().endswith(".gguf") for fname in self.all_files)

    @property
    def is_transformer(self) -> bool:
        """Return True if the repo contains transformer checkpoint files."""
        return any(
            Path(fname).suffix.lower() in TRANSFORMER_EXTENSIONS
            for fname in self.all_files
        )

    @property
    def model_type(self) -> str:
        """Return 'gguf' or 'transformers' based on detected file types."""
        if self.is_gguf:
            return "gguf"
        return "transformers"

    @property
    def tags(self) -> list[str]:
        """Return model tags from the model card."""
        return list(getattr(self._info, "tags", []) or [])

    @property
    def card_data(self) -> dict[str, Any]:
        """Return model card data as a plain dict."""
        card_data_raw = getattr(self._info, "card_data", None)
        if card_data_raw is None:
            return {}
        # card_data may be a ModelCardData object
        if hasattr(card_data_raw, "__dict__"):
            return {
                key: val
                for key, val in card_data_raw.__dict__.items()
                if not key.startswith("_")
            }
        return dict(card_data_raw)

    @property
    def total_size(self) -> int:
        """Return total size of all files in bytes."""
        return sum(sibling.size or 0 for sibling in self._siblings)


# ── helpers ───────────────────────────────────────────────────────────────────

def _parse_quant(filename: str) -> str | None:
    """Extract quantisation name from a GGUF filename, e.g. Q4_K_M."""
    # Typical patterns: model-Q4_K_M.gguf, model.Q4_K_M.gguf
    match_obj = re.search(
        r"[._-]((?:IQ|Q|BF|F)\d+(?:_[A-Z0-9]+)*)\.gguf",
        filename,
        re.IGNORECASE,
    )
    if match_obj:
        return match_obj.group(1).upper()
    return None


def _hf_token() -> str | None:
    """Return HF token from environment or cached credential."""
    token = (
        os.environ.get("HF_TOKEN")
        or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    )
    if token:
        return token
    # Try huggingface_hub cached token
    try:
        from huggingface_hub import get_token  # pylint: disable=import-outside-toplevel
        return get_token()
    except ImportError:
        pass
    token_file = Path.home() / ".huggingface" / "token"
    if token_file.exists():
        return token_file.read_text(encoding="utf-8").strip() or None
    return None


# ── public API ────────────────────────────────────────────────────────────────

def search_models(query: str, limit: int = 8) -> list[dict[str, Any]]:
    """
    Search HuggingFace Hub for models matching *query*.

    Returns a list of dicts with keys: repo_id, pipeline_tag, downloads.
    Results are sorted by download count (most popular first).
    Returns an empty list on any failure.
    """
    api = HfApi()
    token = _hf_token()
    results: list[dict[str, Any]] = []
    try:
        for model in api.list_models(
            search=query,
            limit=limit,
            sort="downloads",
            direction=-1,
            token=token,
        ):
            results.append({
                "repo_id": model.id,
                "pipeline_tag": getattr(model, "pipeline_tag", None) or "—",
                "downloads": getattr(model, "downloads", 0) or 0,
            })
    except Exception:  # pylint: disable=broad-except
        pass
    return results


def fetch_model_info(repo_id: str) -> ModelInfo:
    """Fetch repo metadata from HuggingFace Hub."""
    api = HfApi()
    token = _hf_token()
    return ModelInfo(normalize_repo_id(repo_id), api, token)


def download_gguf_file(
    repo_id: str,
    filename: str,
    dest_dir: Path,
) -> Path:
    """
    Download a single GGUF file from a HuggingFace repo.

    Returns the local path to the downloaded file.
    """
    token = _hf_token()
    dest_dir.mkdir(parents=True, exist_ok=True)

    local_path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        local_dir=str(dest_dir),
        token=token,
    )
    return Path(local_path)


def download_transformer_model(
    repo_id: str,
    dest_dir: Path,
) -> Path:
    """
    Download all transformer model files via snapshot_download.

    Returns the local snapshot directory.
    """
    token = _hf_token()
    dest_dir.mkdir(parents=True, exist_ok=True)

    # Ignore large non-essential formats
    ignore_patterns = [
        "*.msgpack", "flax_model*", "tf_model*", "model.onnx",
        "*.ot", "rust_model.ot",
    ]

    local_dir = snapshot_download(
        repo_id=repo_id,
        local_dir=str(dest_dir),
        token=token,
        ignore_patterns=ignore_patterns,
    )
    return Path(local_dir)
