"""Model registry for mapping names to model classes with default kwargs."""

from __future__ import annotations

from typing import Any

from roohai.base import STTModel, TTSModel, LLMModel, VADModel


class _Registry:
    def __init__(self) -> None:
        # Each entry: (cls, kwargs, meta_overrides)
        self._stt: dict[str, tuple[type[STTModel], dict[str, Any], dict[str, Any]]] = {}
        self._tts: dict[str, tuple[type[TTSModel], dict[str, Any], dict[str, Any]]] = {}
        self._llm: dict[str, tuple[type[LLMModel], dict[str, Any], dict[str, Any]]] = {}
        self._vad: dict[str, tuple[type[VADModel], dict[str, Any], dict[str, Any]]] = {}

    def register_stt(self, name: str, cls: type[STTModel], meta: dict[str, Any] | None = None, **kwargs: Any) -> None:
        self._stt[name] = (cls, kwargs, meta or {})

    def register_tts(self, name: str, cls: type[TTSModel], meta: dict[str, Any] | None = None, **kwargs: Any) -> None:
        self._tts[name] = (cls, kwargs, meta or {})

    def register_llm(self, name: str, cls: type[LLMModel], meta: dict[str, Any] | None = None, **kwargs: Any) -> None:
        self._llm[name] = (cls, kwargs, meta or {})

    def register_vad(self, name: str, cls: type[VADModel], meta: dict[str, Any] | None = None, **kwargs: Any) -> None:
        self._vad[name] = (cls, kwargs, meta or {})

    def create_stt(self, name: str, **overrides: Any) -> STTModel:
        cls, kwargs, _meta = self._stt[name]
        merged = {**kwargs, **overrides}
        return cls(**merged)

    def create_tts(self, name: str, **overrides: Any) -> TTSModel:
        cls, kwargs, _meta = self._tts[name]
        merged = {**kwargs, **overrides}
        return cls(**merged)

    def create_llm(self, name: str, **overrides: Any) -> LLMModel:
        cls, kwargs, _meta = self._llm[name]
        merged = {**kwargs, **overrides}
        return cls(**merged)

    def create_vad(self, name: str, **overrides: Any) -> VADModel:
        cls, kwargs, _meta = self._vad[name]
        merged = {**kwargs, **overrides}
        return cls(**merged)

    def list_stt(self) -> list[str]:
        return list(self._stt.keys())

    def list_tts(self) -> list[str]:
        return list(self._tts.keys())

    def list_llm(self) -> list[str]:
        return list(self._llm.keys())

    def list_vad(self) -> list[str]:
        return list(self._vad.keys())

    def catalog(self) -> dict[str, list[dict[str, Any]]]:
        """Return metadata for all registered models, grouped by type.

        Per-registration meta overrides are merged on top of the class META,
        so each variant can have its own display_name, description, etc.
        Registration kwargs (e.g. model_id) update config_fields defaults
        so the wizard shows the correct value for each variant.
        """
        result: dict[str, list[dict[str, Any]]] = {}
        for model_type, store in [
            ("stt", self._stt),
            ("tts", self._tts),
            ("llm", self._llm),
            ("vad", self._vad),
        ]:
            entries = []
            for name, (cls, kwargs, meta_overrides) in store.items():
                entry = {"name": name, **cls.META, **meta_overrides}

                # Update config_fields defaults from registration kwargs
                # so variants show correct values (e.g. Sonnet shows its model_id)
                if kwargs and "config_fields" in entry:
                    entry["config_fields"] = [
                        {**f, "default": kwargs[f["name"]]}
                        if f["name"] in kwargs else f
                        for f in entry["config_fields"]
                    ]

                entries.append(entry)
            if entries:
                result[model_type] = entries
        return result


registry = _Registry()
