"""
model_registry.py
-----------------
Explicit env-wired model registry. No model loads without a registered spec,
and no spec resolves without a valid env var (hub fallback must be opted into explicitly).

Usage at startup:
    from .model_loaders import register_all
    register_all()          # fails loudly if any env var is missing

Usage in pipeline code:
    from .model_registry import get_model, ModelKey
    whisper = get_model(ModelKey.WHISPER)
"""
from __future__ import annotations

import logging
import os
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Dict

logger = logging.getLogger(__name__)


class ModelKey(str, Enum):
    WHISPER  = "whisper"
    LED      = "led"        # Longformer-Encoder-Decoder summarizer
    FLAN     = "flan"       # FLAN-T5 refiner
    KEYBERT  = "keybert"    # keyword extraction


@dataclass
class ModelSpec:
    key:               ModelKey
    env_var:           str          # env var whose value is the local model directory
    fallback_hub_id:   str          # HuggingFace hub ID — only used if allow_hub_fallback=True
    loader:            Callable     # loader(path: str) → model instance
    allow_hub_fallback: bool = False  # must be explicitly opted into; never True in prod


# Single module-level registry dict. Not a global config object — just state.
_REGISTRY: Dict[ModelKey, Any] = {}


def _resolve_path(spec: ModelSpec) -> str:
    """
    Resolve the model path from the environment.
    Raises EnvironmentError loudly rather than silently falling back.
    """
    path = os.environ.get(spec.env_var, "").strip()
    if path and os.path.isdir(path):
        return path
    if path and not os.path.isdir(path):
        raise EnvironmentError(
            f"[ModelRegistry] {spec.env_var}='{path}' is set but is not a valid directory."
        )
    # env var not set at all
    if spec.allow_hub_fallback:
        logger.warning(
            "[ModelRegistry] %s not set; falling back to HuggingFace hub: %s. "
            "Set allow_hub_fallback=False before going to production.",
            spec.env_var,
            spec.fallback_hub_id,
        )
        return spec.fallback_hub_id
    raise EnvironmentError(
        f"[ModelRegistry] {spec.env_var} is not set or does not point to a valid directory.\n"
        f"  Expected: a local path to '{spec.key}' model files.\n"
        f"  Hub fallback ({spec.fallback_hub_id}) is disabled for this model.\n"
        f"  Set the env var before starting the application."
    )


def register_model(spec: ModelSpec) -> None:
    """Load and register a model from its spec. Idempotent — safe to call multiple times."""
    if spec.key in _REGISTRY:
        logger.debug("[ModelRegistry] %s already registered, skipping.", spec.key)
        return
    path = _resolve_path(spec)
    logger.info("[ModelRegistry] Loading %s from %s ...", spec.key, path)
    _REGISTRY[spec.key] = spec.loader(path)
    logger.info("[ModelRegistry] %s ready.", spec.key)


def get_model(key: ModelKey) -> Any:
    """
    Retrieve a registered model. Raises KeyError if not yet registered.
    This is intentional — callers should not silently get None.
    """
    if key not in _REGISTRY:
        raise KeyError(
            f"[ModelRegistry] Model '{key}' is not registered. "
            f"Call register_model() or register_all() at startup before using the pipeline."
        )
    return _REGISTRY[key]


def is_registered(key: ModelKey) -> bool:
    return key in _REGISTRY


def unregister_model(key: ModelKey) -> None:
    """Remove a model from the registry (useful in tests to force reload)."""
    _REGISTRY.pop(key, None)
