"""
Weight path resolver for bundled model weights.

This module provides a centralized way to access bundled model weights,
with fallback support for environment variables, HF Hub auto-download,
and custom paths.

Resolution order:
1. Custom path (if provided)
2. Environment variable override
3. Bundled weights (sltk/weights/)
4. HF Hub cache (~/.cache/sltk/weights/)
5. Auto-download from HF Hub

Usage:
    from sltk.extraction.weights import get_weight_path, get_all_weights_status

    # Get path to a specific weight (auto-downloads if needed)
    wilor_path = get_weight_path("wilor_checkpoint")

    # Get status of all weights
    status = get_all_weights_status()
"""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import TypedDict

from sltk.exceptions import ModelNotFoundError

logger = logging.getLogger(__name__)

# Weights directory relative to this module
WEIGHTS_DIR = Path(__file__).parent.parent / "weights"

# Environment variable overrides
ENV_OVERRIDES = {
    "wilor_checkpoint": "SLTK_WILOR_CHECKPOINT",
    "wilor_detector": "SLTK_WILOR_DETECTOR",
    "wilor_config": "SLTK_WILOR_CONFIG",
    "nlf_model": "SLTK_NLF_MODEL",
    "nlf_smplx": "SLTK_NLF_SMPLX",
    "nlf_detector": "SLTK_NLF_DETECTOR",
    "signrep_checkpoint": "SLTK_SIGNREP_CHECKPOINT",
    "segmentor": "SLTK_SEGMENTOR_CHECKPOINT",
    "segmentor_v2": "SLTK_SEGMENTOR_V2_CHECKPOINT",
    "teaser_checkpoint": "SLTK_TEASER_CHECKPOINT",
    "teaser_face_detector": "SLTK_TEASER_FACE_DETECTOR",
    "flame_model": "SLTK_FLAME_MODEL",
    "flame_centroids": "SLTK_FLAME_CENTROIDS",
    "flame_indices": "SLTK_FLAME_INDICES",
    "smplfx_model": "SLTK_SMPLFX_MODEL",
    "smplfx_segmentation": "SLTK_SMPLFX_SEGMENTATION",
    "smplfx_flame_masks": "SLTK_SMPLFX_FLAME_MASKS",
    "smplfx_flame_corr": "SLTK_SMPLFX_FLAME_CORR",
    "depth_anything_vits": "SLTK_DEPTH_ANYTHING_VITS",
    "depth_anything_vitl": "SLTK_DEPTH_ANYTHING_VITL",
}

# Relative paths within any weights directory (bundled or cache)
WEIGHT_PATHS_RELATIVE = {
    "wilor_checkpoint": "wilor/wilor_final.ckpt",
    "wilor_detector": "wilor/detector.pt",
    "wilor_config": "wilor/model_config.yaml",
    "nlf_model": "nlf/nlf_l_multi_0.3.2.torchscript",
    "nlf_smplx": "nlf/SMPLX_NEUTRAL.npz",
    "nlf_detector": "nlf/yolov8x.pt",
    "signrep_checkpoint": "signrep/ckpt.pt",
    "segmentor": "segmentor/segmentor_v2.ckpt",
    "segmentor_v2": "segmentor/segmentor_v2.ckpt",
}

# Default paths within the bundled weights directory
WEIGHT_PATHS = {name: WEIGHTS_DIR / rel_path for name, rel_path in WEIGHT_PATHS_RELATIVE.items()}

# Descriptions for UI display
WEIGHT_DESCRIPTIONS = {
    "wilor_checkpoint": "WiLoR hand pose model checkpoint",
    "wilor_detector": "YOLO hand detector for WiLoR",
    "wilor_config": "WiLoR model configuration",
    "nlf_model": "NLF SMPLX pose estimation model (TorchScript)",
    "nlf_smplx": "SMPLX body model neutral shape",
    "nlf_detector": "YOLOv8 body detector for NLF",
    "signrep_checkpoint": "SignRep video embedding model checkpoint",
    "segmentor": "Sign segmentation model checkpoint",
    "segmentor_v2": "Sign segmentation model v2 (WiLoR-only)",
    "teaser_checkpoint": "TEASER FLAME face encoder checkpoint",
    "teaser_face_detector": "YOLOv8 face detector for TEASER",
    "flame_model": "FLAME 2020 generic face model (pickle)",
    "flame_centroids": "K-means centroids for expression classification (20 clusters)",
    "flame_indices": "Indices of interest (lower-face vertices) for classification",
    # SMPLFX
    "smplfx_model": "SMPLX neutral body model for SMPLFX fitting",
    "smplfx_segmentation": "SMPLX vertex segmentation JSON for collision detection",
    "smplfx_flame_masks": "FLAME vertex masks (PKL) for SMPLFX face fitting",
    "smplfx_flame_corr": "SMPLX-to-FLAME correspondence mapping (NPY)",
    # Depth Anything
    "depth_anything_vits": "Video Depth Anything ViT-S checkpoint",
    "depth_anything_vitl": "Video Depth Anything ViT-L checkpoint",
}


class WeightStatus(TypedDict):
    """Status information for a weight file."""

    path: str
    exists: bool
    size_mb: float
    source: str  # 'bundled', 'env', or 'missing'
    description: str


def get_weight_path(name: str, custom_path: str | None = None) -> Path | None:
    """
    Get path to a bundled or configured weight file.

    Resolution order:
    1. Custom path (if provided and exists)
    2. Environment variable override
    3. Bundled weights directory

    Args:
        name: Weight identifier (e.g., 'wilor_checkpoint', 'nlf_detector')
        custom_path: Optional user-specified path to use instead

    Returns:
        Path to the weight file if found, None otherwise

    Example:
        >>> path = get_weight_path("wilor_checkpoint")
        >>> if path:
        ...     model.load_weights(path)
    """
    # Check custom path first
    if custom_path:
        custom = Path(custom_path)
        if custom.exists():
            return custom

    # Check environment variable
    env_var = ENV_OVERRIDES.get(name)
    if env_var:
        env_path = os.environ.get(env_var)
        if env_path:
            env_path_obj = Path(env_path)
            if env_path_obj.exists():
                return env_path_obj

    # Check bundled weights
    bundled = WEIGHT_PATHS.get(name)
    if bundled and bundled.exists():
        return bundled

    # Check HF cache directory
    try:
        from sltk.weights.hub import WEIGHT_TO_MODEL, get_cache_dir
    except (ImportError, ModuleNotFoundError):
        return None

    rel_path = WEIGHT_PATHS_RELATIVE.get(name)
    if rel_path:
        cached = get_cache_dir() / rel_path
        if cached.exists():
            return cached

    # Auto-download from HF Hub if this is a known model weight
    model = WEIGHT_TO_MODEL.get(name)
    if model and rel_path:
        try:
            from sltk.weights.hub import ensure_weights

            logger.info("Weight '%s' not found locally, downloading from HF Hub...", name)
            paths = ensure_weights(model)
            if rel_path in paths:
                return paths[rel_path]
        except Exception as e:
            logger.warning("Failed to download weight '%s': %s", name, e)

    return None


def get_all_weights_status() -> dict[str, WeightStatus]:
    """
    Get status of all known weight files.

    Returns:
        Dictionary mapping weight names to their status including
        path, existence, size, source, and description.

    Example:
        >>> status = get_all_weights_status()
        >>> for name, info in status.items():
        ...     print(f"{name}: {'Found' if info['exists'] else 'Missing'}")
    """
    result = {}

    for name, bundled_path in WEIGHT_PATHS.items():
        # Check sources in order of priority
        env_var = ENV_OVERRIDES.get(name)
        env_path = os.environ.get(env_var) if env_var else None

        if env_path:
            path = Path(env_path)
            source = "env"
        else:
            path = bundled_path
            source = "bundled"

        exists = path.exists()

        try:
            size_mb = path.stat().st_size / (1024 * 1024) if exists else 0.0
        except (OSError, PermissionError):
            size_mb = 0.0

        result[name] = WeightStatus(
            path=str(path),
            exists=exists,
            size_mb=round(size_mb, 2),
            source=source if exists else "missing",
            description=WEIGHT_DESCRIPTIONS.get(name, ""),
        )

    return result


def ensure_weights_available(required: list[str]) -> dict[str, Path]:
    """
    Ensure all required weights are available.

    Args:
        required: List of weight names that must be available

    Returns:
        Dictionary mapping weight names to their paths

    Raises:
        FileNotFoundError: If any required weight is missing

    Example:
        >>> paths = ensure_weights_available(["wilor_checkpoint", "wilor_detector"])
        >>> extractor.load_model(checkpoint=paths["wilor_checkpoint"])
    """
    paths = {}
    missing = []

    for name in required:
        path = get_weight_path(name)
        if path:
            paths[name] = path
        else:
            missing.append(name)

    if missing:
        missing_str = ", ".join(missing)
        raise ModelNotFoundError(
            missing_str,
            path=str(WEIGHTS_DIR),
        )

    return paths


def get_weights_dir() -> Path:
    """Get the bundled weights directory path."""
    return WEIGHTS_DIR
