"""
model_loaders.py
----------------
One loader function per model. Each loader takes a resolved path string and returns
a ready-to-use model instance. Loaders are called exactly once per process at startup.

Call register_all() in your app entry point:
    from .model_loaders import register_all
    register_all()

Or register selectively for partial deployments / testing:
    from .model_loaders import register_selective
    from .model_registry import ModelKey
    register_selective(ModelKey.WHISPER, ModelKey.FLAN)

Required environment variables:
    WHISPER_MODEL_PATH   — local dir containing whisper model files
    WHISPER_MODEL_SIZE   — model size string: tiny|base|small|medium|large (default: base)
    LED_MODEL_PATH       — local dir containing LED/Longformer model files
    FLAN_MODEL_PATH      — local dir containing FLAN-T5 model files
    KEYBERT_MODEL_PATH   — local dir containing sentence-transformers model files
"""
from __future__ import annotations

import logging
import os
from typing import Any

from .model_registry import ModelKey, ModelSpec, register_model

logger = logging.getLogger(__name__)


# ── Loader functions ──────────────────────────────────────────────────────────
# Each returns a ready-to-use object. The registry stores this object and hands
# it out via get_model(key). Types are kept opaque — callers import nothing
# from transformers/whisper/keybert directly.

def _load_whisper(path: str) -> Any:
    import whisper
    size = os.environ.get("WHISPER_MODEL_SIZE", "base")
    logger.info("[Loader] whisper size=%s from %s", size, path)
    return whisper.load_model(size, download_root=path)


def _load_led(path: str) -> dict:
    """Returns dict with 'tokenizer' and 'model' keys (LED is stateful, keep together)."""
    from transformers import LEDForConditionalGeneration, LEDTokenizer
    logger.info("[Loader] LED from %s", path)
    return {
        "tokenizer": LEDTokenizer.from_pretrained(path),
        "model":     LEDForConditionalGeneration.from_pretrained(path),
    }


def _load_flan(path: str) -> Any:
    """Returns a HuggingFace pipeline for text2text-generation."""
    import torch
    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
    logger.info("[Loader] FLAN-T5 from %s", path)
    tok = AutoTokenizer.from_pretrained(path)
    mdl = AutoModelForSeq2SeqLM.from_pretrained(path)
    device = 0 if torch.cuda.is_available() else -1
    return pipeline("text2text-generation", model=mdl, tokenizer=tok, device=device)


def _load_keybert(path: str) -> Any:
    """Returns a KeyBERT instance backed by a local SentenceTransformer."""
    from keybert import KeyBERT
    from sentence_transformers import SentenceTransformer
    logger.info("[Loader] KeyBERT/SentenceTransformer from %s", path)
    st = SentenceTransformer(path)
    return KeyBERT(model=st)


# ── Spec table ────────────────────────────────────────────────────────────────
# allow_hub_fallback=False everywhere. Set to True only in dev/testing environments
# by patching the spec or by setting the env var to a valid hub ID.

_SPECS: list[ModelSpec] = [
    ModelSpec(
        key=ModelKey.WHISPER,
        env_var="WHISPER_MODEL_PATH",
        fallback_hub_id="openai/whisper-base",
        loader=_load_whisper,
        allow_hub_fallback=False,
    ),
    ModelSpec(
        key=ModelKey.LED,
        env_var="LED_MODEL_PATH",
        fallback_hub_id="allenai/led-base-16384",
        loader=_load_led,
        allow_hub_fallback=False,
    ),
    ModelSpec(
        key=ModelKey.FLAN,
        env_var="FLAN_MODEL_PATH",
        fallback_hub_id="google/flan-t5-xl",
        loader=_load_flan,
        allow_hub_fallback=False,
    ),
    ModelSpec(
        key=ModelKey.KEYBERT,
        env_var="KEYBERT_MODEL_PATH",
        fallback_hub_id="sentence-transformers/all-MiniLM-L6-v2",
        loader=_load_keybert,
        allow_hub_fallback=False,
    ),
]

_SPEC_MAP: dict[ModelKey, ModelSpec] = {s.key: s for s in _SPECS}


# ── Public registration API ───────────────────────────────────────────────────

def register_all() -> None:
    """
    Register all models. Call once at application startup.
    Fails loudly if any required env var is missing.
    """
    for spec in _SPECS:
        register_model(spec)


def register_selective(*keys: ModelKey) -> None:
    """
    Register only specific models.
    Useful for: partial deployments, testing a single stage, CI without GPU.
    """
    for key in keys:
        if key not in _SPEC_MAP:
            raise ValueError(f"[model_loaders] Unknown ModelKey: {key}")
        register_model(_SPEC_MAP[key])
