"""Model registry: the small instruct models this demo steers.

On ZeroGPU every model is loaded at import time and moved to "cuda" (emulated until a GPU is
attached), as ZeroGPU requires. Locally the same models run on the CPU in float32, which is how the
recorded results in data/recorded/ were measured.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass
from typing import Any

import torch

from steering.core import load_model, num_layers

log = logging.getLogger(__name__)


@dataclass(frozen=True)
class ModelSpec:
    key: str
    label: str
    model_id: str


# The first entry is the default: the model the scenarios were tuned on and the takeaways describe.
SPECS = [
    ModelSpec("qwen25-05b", "Qwen2.5-0.5B-Instruct", "Qwen/Qwen2.5-0.5B-Instruct"),
    ModelSpec("smollm2-135m", "SmolLM2-135M-Instruct", "HuggingFaceTB/SmolLM2-135M-Instruct"),
]
BY_KEY = {spec.key: spec for spec in SPECS}
BY_ID = {spec.model_id: spec for spec in SPECS}
DEFAULT = SPECS[0]


@dataclass
class Loaded:
    spec: ModelSpec
    tok: Any
    model: Any
    depth: int


REGISTRY: dict[str, Loaded] = {}
LOAD_ERRORS: list[str] = []
DEVICE = torch.device("cpu")


def on_zerogpu() -> bool:
    return bool(os.environ.get("SPACES_ZERO_GPU") or os.environ.get("SPACE_ID"))


def pick_device() -> torch.device:
    if on_zerogpu():
        return torch.device("cuda")
    return torch.device(os.environ.get("SD_DEVICE", "cpu"))


def resolve(name: str | None) -> ModelSpec:
    """A model key or a Hugging Face id -> its spec. Anything else -> the default model.

    A result always names the model that produced it, so a request for an unknown model that
    falls back to the default is visible rather than misleading.
    """
    name = (name or "").strip()
    return BY_KEY.get(name) or BY_ID.get(name) or DEFAULT


def init(keys: list[str] | None = None) -> None:
    """Load every model, one at a time. A model that fails to load is skipped and reported."""
    global DEVICE
    DEVICE = pick_device()
    wanted = keys or [k for k in os.environ.get("SD_MODELS", "").split(",") if k]
    for spec in SPECS:
        if wanted and spec.key not in wanted:
            continue
        try:
            log.info("loading %s on %s", spec.model_id, DEVICE)
            tok, model = load_model(spec.model_id, device=str(DEVICE))
            REGISTRY[spec.key] = Loaded(spec, tok, model, num_layers(model))
        except Exception as exc:  # network error, missing weights, ...
            log.exception("could not load %s", spec.model_id)
            LOAD_ERRORS.append(f"{spec.label} ({type(exc).__name__}: {str(exc).splitlines()[0][:160]})")