"""HuggingFace Hub integration for COMBO models."""

import json
import logging
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

MODEL_REGISTRY_REPO = "clarin-pl/combo-nlp-models"
MODEL_REGISTRY_FILE = "combo-nlp-models.json"

# Cache for the language→model mapping
_language_model_map: Optional[dict[str, str]] = None
# Cache for the raw registry data (needed for LAMBO model lookups)
_registry_data: Optional[dict] = None


def get_language_model_map(force_refresh: bool = False) -> dict[str, str]:
    """Fetch language→model mapping from HuggingFace Hub.

    The mapping is fetched once and cached in memory. It maps language names
    (e.g. "Polish") to HuggingFace model IDs (e.g. "clarin-pl/combo-polish-herbert-base").

    Args:
        force_refresh: Force re-fetching the mapping.

    Returns:
        Dictionary mapping language names to model IDs.
    """
    global _language_model_map, _registry_data
    if _language_model_map is not None and not force_refresh:
        return _language_model_map

    try:
        from huggingface_hub import hf_hub_download

        path = hf_hub_download(
            repo_id=MODEL_REGISTRY_REPO,
            filename=MODEL_REGISTRY_FILE,
            repo_type="model",
        )
        with open(path) as f:
            data = json.load(f)

        _registry_data = data

        # Build language→model mapping
        # Expected format: list of objects or dict with language keys
        if isinstance(data, list):
            _language_model_map = {
                entry.get("language", entry.get("name", "")): entry.get("model", entry.get("model_id", ""))
                for entry in data
                if entry.get("language", entry.get("name")) and entry.get("model", entry.get("model_id"))
            }
        elif isinstance(data, dict):
            # Could be {language: model_id} or {language: {model: ...}}
            mapping = {}
            for key, value in data.items():
                if isinstance(value, str):
                    mapping[key] = value
                elif isinstance(value, dict):
                    mapping[key] = value.get("default", value.get("model", value.get("model_id", "")))
            _language_model_map = mapping
        else:
            _language_model_map = {}

        logger.info(f"Loaded {len(_language_model_map)} language→model mappings from HF Hub")
        return _language_model_map

    except Exception as e:
        logger.warning(f"Failed to fetch language model map from HF Hub: {e}")
        _language_model_map = {}
        return _language_model_map


def resolve_model(language_or_path: str) -> str:
    """Resolve a language name to a HuggingFace model ID.

    If the input contains "/" it is treated as a HuggingFace model path and
    returned as-is. Otherwise, it is looked up in the language→model mapping.

    Args:
        language_or_path: Language name (e.g. "Polish") or HF model ID.

    Returns:
        HuggingFace model ID.

    Raises:
        ValueError: If the language is not found in the mapping.
    """
    if "/" in language_or_path:
        return language_or_path

    mapping = get_language_model_map()

    # Try exact match first
    if language_or_path in mapping:
        return mapping[language_or_path]

    # Try case-insensitive match
    lower = language_or_path.lower()
    for lang, model_id in mapping.items():
        if lang.lower() == lower:
            return model_id

    from .languages import LanguageNotTrainedError

    available = ", ".join(sorted(mapping.keys()))
    raise LanguageNotTrainedError(
        f"Language {language_or_path!r} does not have a trained COMBO model yet. "
        f"Languages with trained models: {available or '(none - could not fetch model registry)'}. "
        f"You can also pass a full HuggingFace model path with "
        f"COMBO.from_pretrained('clarin-pl/combo-polish-herbert-base')."
    )


def download_model(model_id: str, cache_dir: Optional[str] = None) -> Path:
    """Download a COMBO model from HuggingFace Hub.

    Each model repository is expected to contain:
    - config.json — model configuration
    - pytorch_model.bin — model weights
    - encoders/morpho_encoder.json
    - encoders/deprel_encoder.json
    - encoders/char_vocab.json

    Args:
        model_id: HuggingFace model ID (e.g. "clarin-pl/combo-polish-herbert-base").
        cache_dir: Optional cache directory.

    Returns:
        Path to the downloaded model directory.
    """
    from huggingface_hub import snapshot_download

    local_dir = snapshot_download(
        repo_id=model_id,
        cache_dir=cache_dir,
    )
    return Path(local_dir)


def resolve_lambo_model(combo_model_id: str) -> Optional[str]:
    """Look up the LAMBO model name for a given COMBO model ID.

    Uses the registry data from combo-nlp-models.json. The expected format is::

        {
          "Polish": {
            "default": "clarin-pl/combo-nlp-herbert-base-polish-pdb-ud2.17",
            "models": {
              "clarin-pl/combo-nlp-herbert-base-polish-pdb-ud2.17": {
                "lambo": "LAMBO_213-UD_Polish-PDB"
              }
            }
          }
        }

    Args:
        combo_model_id: HuggingFace COMBO model ID.

    Returns:
        LAMBO model name string, or None if not found.
    """
    # Ensure registry is loaded
    get_language_model_map()

    if not isinstance(_registry_data, dict):
        return None

    for lang_entry in _registry_data.values():
        if not isinstance(lang_entry, dict):
            continue
        models = lang_entry.get("models", {})
        if combo_model_id in models:
            model_info = models[combo_model_id]
            if isinstance(model_info, dict):
                return model_info.get("lambo")

    return None
