"""Registry: deklarative Slots/Adapter (YAML) plus Laufzeitzustand (Store)."""

from __future__ import annotations

from pathlib import Path
from typing import Any, Literal

from pydantic import BaseModel, Field

from .store import Store
from .util import load_yaml

SlotStatus = Literal["ok", "degraded", "unavailable"]


class SlotSpec(BaseModel):
    name: str
    tier: int
    device: Literal["gpu", "cpu", "external"] = "gpu"
    engine: str = "vllm"
    node: str | None = None
    litellm_model: str
    base_model: str
    max_context: int
    effective_context: int | None = None
    supports_lora: bool = False
    fallbacks: list[str] = Field(default_factory=list)
    kind: Literal["chat", "embedding", "rerank"] = "chat"


class AdapterSpec(BaseModel):
    name: str
    base_model: str
    cluster: str
    version: int = 1
    status: Literal["active", "candidate", "retired"] = "active"


class Slot(BaseModel):
    """Effektive Sicht: Spezifikation + Laufzeitzustand."""

    spec: SlotSpec
    status: SlotStatus = "ok"
    reason: str | None = None
    set_by: str | None = None
    calibrated_context: int | None = None
    base_model_override: str | None = None
    canary_model: str | None = None
    canary_share: float = 0.0
    canary_cluster: str | None = None

    @property
    def name(self) -> str:
        return self.spec.name

    @property
    def base_model(self) -> str:
        return self.base_model_override or self.spec.base_model

    @property
    def effective_context(self) -> int:
        for value in (self.calibrated_context, self.spec.effective_context):
            if value:
                return min(value, self.spec.max_context)
        return self.spec.max_context

    @property
    def available(self) -> bool:
        return self.status == "ok"


class Registry:
    def __init__(self, slots: dict[str, SlotSpec], adapters: list[AdapterSpec], store: Store):
        self._slots = slots
        self._adapters = {a.name: a for a in adapters}
        self.store = store

    @classmethod
    def from_file(cls, path: str | Path, store: Store) -> Registry:
        data = load_yaml(path)
        return cls.from_dict(data, store)

    @classmethod
    def from_dict(cls, data: dict[str, Any], store: Store) -> Registry:
        slots = {name: SlotSpec(name=name, **spec) for name, spec in (data.get("slots") or {}).items()}
        adapters = [AdapterSpec(**a) for a in data.get("adapters") or []]
        for spec in slots.values():
            for fb in spec.fallbacks:
                if fb not in slots:
                    raise ValueError(f"Slot {spec.name}: unbekannter Fallback {fb}")
        return cls(slots, adapters, store)

    # --- Slots ------------------------------------------------------------------------------
    def slot_names(self) -> list[str]:
        return list(self._slots)

    def has_slot(self, name: str) -> bool:
        return name in self._slots

    def slot(self, name: str) -> Slot:
        spec = self._slots[name]
        state = self.store.get_slot_state(name) or {}
        return Slot(
            spec=spec,
            status=state.get("status") or "ok",
            reason=state.get("reason"),
            set_by=state.get("set_by"),
            calibrated_context=state.get("effective_context"),
            base_model_override=state.get("base_model"),
            canary_model=state.get("canary_model"),
            canary_share=state.get("canary_share") or 0.0,
            canary_cluster=state.get("canary_cluster"),
        )

    def slots(self) -> list[Slot]:
        return [self.slot(n) for n in self._slots]

    def strong_slot(self) -> Slot:
        chat = [s for s in self.slots() if s.spec.kind == "chat" and s.spec.tier == 0]
        if not chat:
            raise LookupError("Kein Slot mit tier 0 konfiguriert")
        available = [s for s in chat if s.available]
        return (available or chat)[0]

    def set_status(self, name: str, status: SlotStatus, reason: str = "", set_by: str = "human") -> None:
        if name not in self._slots:
            raise KeyError(name)
        self.store.set_slot_state(name, status=status, reason=reason, set_by=set_by)

    def set_calibrated_context(self, name: str, tokens: int | None) -> None:
        self.store.set_slot_state(name, effective_context=tokens)

    def set_base_model(self, name: str, base_model: str | None) -> None:
        self.store.set_slot_state(name, base_model=base_model)

    def set_canary(self, name: str, model: str | None, share: float = 0.0, cluster: str | None = None) -> None:
        """Leitet `share` des Stufe-1-Verkehrs (optional nur eines Clusters) auf ein Kandidatenmodell."""
        self.store.set_slot_state(name, canary_model=model, canary_share=share if model else 0.0,
                                  canary_cluster=cluster if model else None)

    # --- Adapter ----------------------------------------------------------------------------
    def adapters(self) -> list[AdapterSpec]:
        merged = {n: a.model_copy() for n, a in self._adapters.items()}
        for row in self.store.all_adapter_states():
            if row["name"] in merged:
                if row.get("status"):
                    merged[row["name"]].status = row["status"]
            else:
                merged[row["name"]] = AdapterSpec(
                    name=row["name"],
                    base_model=row["base_model"],
                    cluster=row["cluster"],
                    version=row.get("version") or 1,
                    status=row.get("status") or "candidate",
                )
        return list(merged.values())

    def active_adapter(self, base_model: str, cluster: str) -> AdapterSpec | None:
        found = [
            a for a in self.adapters()
            if a.base_model == base_model and a.cluster == cluster and a.status == "active"
        ]
        return max(found, key=lambda a: a.version) if found else None

    def set_adapter(self, adapter: AdapterSpec) -> None:
        self.store.set_adapter_state(
            adapter.name,
            base_model=adapter.base_model,
            cluster=adapter.cluster,
            version=adapter.version,
            status=adapter.status,
        )

    def export(self) -> dict[str, Any]:
        """Effektiver Stand als YAML-fähiges Dict (für Rückführung nach Git)."""
        slots = {}
        for s in self.slots():
            data = s.spec.model_dump(exclude={"name"})
            data["base_model"] = s.base_model
            if s.calibrated_context:
                data["effective_context"] = s.calibrated_context
            slots[s.name] = data
        return {"slots": slots, "adapters": [a.model_dump() for a in self.adapters()]}