"""
Auto-download model weights from Hugging Face Hub.

Downloads weights on first use per model, caching them locally.
Respects SLTK_WEIGHTS_DIR environment variable for custom cache location.

Usage:
    from sltk.weights.hub import ensure_weights

    # Download only what you need
    paths = ensure_weights("wilor")      # downloads WiLoR weights only
    paths = ensure_weights("segmentor")  # downloads segmentor weights only
"""

from __future__ import annotations

import logging
import os
from pathlib import Path

logger = logging.getLogger(__name__)

# HF repo containing all SLTK weights
HF_REPO_ID = "fiskenai/vltk"

# Default cache location
_DEFAULT_CACHE_DIR = Path.home() / ".cache" / "sltk" / "weights"

# Which HF repo files belong to each model
MODEL_FILES: dict[str, list[str]] = {
    "wilor": [
        "wilor/wilor_final.ckpt",
        "wilor/detector.pt",
        "wilor/model_config.yaml",
        "wilor/mano_data/MANO_LEFT.pkl",
        "wilor/mano_data/MANO_RIGHT.pkl",
        "wilor/mano_data/mano_mean_params.npz",
    ],
    "nlf": [
        "nlf/nlf_l_multi_0.3.2.torchscript",
        "nlf/SMPLX_NEUTRAL.npz",
        "nlf/yolov8x.pt",
    ],
    "signrep": [
        "signrep/ckpt.pt",
    ],
    "segmentor": [
        "segmentor/segmentor_v2.ckpt",
    ],
}

# Map weight registry names (from weights.py) to model groups
WEIGHT_TO_MODEL: dict[str, str] = {
    "wilor_checkpoint": "wilor",
    "wilor_detector": "wilor",
    "wilor_config": "wilor",
    "nlf_model": "nlf",
    "nlf_smplx": "nlf",
    "nlf_detector": "nlf",
    "signrep_checkpoint": "signrep",
    "segmentor": "segmentor",
    "segmentor_v2": "segmentor",
}


def get_cache_dir() -> Path:
    """Get the weights cache directory, respecting SLTK_WEIGHTS_DIR env var."""
    env_dir = os.environ.get("SLTK_WEIGHTS_DIR")
    if env_dir:
        return Path(env_dir)
    return _DEFAULT_CACHE_DIR


def ensure_weights(model: str) -> dict[str, Path]:
    """
    Ensure all weight files for a model are downloaded.

    Downloads from Hugging Face Hub on first use. Subsequent calls
    return cached paths immediately.

    Args:
        model: Model name — one of "wilor", "nlf", "signrep", "segmentor"

    Returns:
        Dict mapping relative file paths to their local absolute paths.

    Raises:
        ValueError: If model name is unknown.
        ImportError: If huggingface_hub is not installed.
    """
    if model not in MODEL_FILES:
        available = ", ".join(sorted(MODEL_FILES.keys()))
        raise ValueError(f"Unknown model '{model}'. Available: {available}")

    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        raise ImportError(
            "huggingface_hub is required for weight downloads. "
            "Install with: pip install huggingface_hub"
        )

    cache_dir = get_cache_dir()
    files = MODEL_FILES[model]
    paths: dict[str, Path] = {}

    for rel_path in files:
        local_path = cache_dir / rel_path

        if local_path.exists():
            paths[rel_path] = local_path
            continue

        logger.info("Downloading %s from %s...", rel_path, HF_REPO_ID)
        downloaded = hf_hub_download(  # type: ignore[call-overload]
            repo_id=HF_REPO_ID,
            filename=rel_path,
            local_dir=str(cache_dir),
            local_dir_use_symlinks=False,
        )
        paths[rel_path] = Path(downloaded)
        logger.info("Downloaded %s", rel_path)

    return paths


def ensure_weight_by_name(name: str) -> Path | None:
    """
    Ensure a specific weight is available by its registry name.

    This is the integration point with weights.py — call this when
    the bundled/env path doesn't exist to auto-download.

    Args:
        name: Weight registry name (e.g., "wilor_checkpoint")

    Returns:
        Path to the weight file, or None if the weight name is unknown.
    """
    model = WEIGHT_TO_MODEL.get(name)
    if not model:
        return None

    from sltk.extraction.weights import WEIGHT_PATHS_RELATIVE

    paths = ensure_weights(model)
    rel_path = WEIGHT_PATHS_RELATIVE.get(name)
    if rel_path and rel_path in paths:
        return paths[rel_path]

    return None
