"""Carga (y cachea) un modelo GLiREL en el device deseado."""

from __future__ import annotations

from typing import Any

# Cache global: (model_name, device) -> modelo cargado.
_MODEL_CACHE: dict[tuple[str, str], Any] = {}


def _resolve_device(device: str) -> str:
    """Resuelve `device='auto'` a `cuda` o `cpu` segun disponibilidad."""
    if device != "auto":
        return device
    try:
        import torch
    except ImportError:
        return "cpu"
    return "cuda" if torch.cuda.is_available() else "cpu"


def glirel_load_model(
    model_name: str = "jackboyla/glirel-large-v0",
    device: str = "auto",
) -> Any:
    """Carga un modelo GLiREL con cache por (model_name, device).

    La primera llamada descarga el modelo desde HuggingFace (~500 MB para
    `glirel-large-v0`). Llamadas sucesivas con los mismos parametros
    devuelven la instancia cacheada.

    Args:
        model_name: ID del modelo en HuggingFace Hub.
        device: 'auto' usa CUDA si esta disponible, o 'cpu'/'cuda'/'cuda:N'
            de forma explicita.

    Returns:
        Instancia del modelo GLiREL lista para `predict_relations`.

    Raises:
        ImportError: si la dependencia `glirel` no esta instalada.
            Solucion: `uv pip install glirel` o instalar el extra `nlp`
            del proyecto (`uv pip install -e '.[nlp]'`).
    """
    resolved_device = _resolve_device(device)
    cache_key = (model_name, resolved_device)
    cached = _MODEL_CACHE.get(cache_key)
    if cached is not None:
        return cached

    try:
        from glirel import GLiREL
    except ImportError as exc:
        raise ImportError(
            "glirel no esta instalado. Instalalo con "
            "`uv pip install glirel` o `uv pip install -e '.[nlp]'`."
        ) from exc

    model = GLiREL.from_pretrained(model_name)
    if hasattr(model, "to"):
        model.to(resolved_device)
    _MODEL_CACHE[cache_key] = model
    return model