"""Tests for the model registry system."""

import pytest
import numpy as np
from unittest.mock import MagicMock

from roohai.registry import _Registry
from roohai.base import STTModel, TTSModel, LLMModel


class ConcreteSTT(STTModel):
    def __init__(self, **kwargs):
        self.kwargs = kwargs
        self._loaded = False

    def load(self):
        self._loaded = True

    def transcribe(self, audio: np.ndarray, sample_rate: int = 16000) -> str:
        return "hello"

    def unload(self):
        self._loaded = False

    @property
    def is_loaded(self) -> bool:
        return self._loaded


class ConcreteTTS(TTSModel):
    def __init__(self, **kwargs):
        self.kwargs = kwargs
        self._loaded = False

    def load(self):
        self._loaded = True

    def synthesize(self, text: str) -> tuple[np.ndarray, int]:
        return np.zeros(100, dtype=np.float32), 24000

    def unload(self):
        self._loaded = False

    @property
    def is_loaded(self) -> bool:
        return self._loaded


class ConcreteLLM(LLMModel):
    def __init__(self, **kwargs):
        self.kwargs = kwargs
        self._loaded = False

    def load(self):
        self._loaded = True

    async def chat(self, message, history=None, system_prompt=None):
        return "response"

    def unload(self):
        self._loaded = False

    @property
    def is_loaded(self) -> bool:
        return self._loaded


class TestModelRegistry:
    def test_register_and_create_stt(self):
        reg = _Registry()
        reg.register_stt("test-stt", ConcreteSTT, model_id="test")
        model = reg.create_stt("test-stt")
        assert isinstance(model, ConcreteSTT)
        assert model.kwargs["model_id"] == "test"

    def test_register_and_create_tts(self):
        reg = _Registry()
        reg.register_tts("test-tts", ConcreteTTS, vocoder="test-vocoder")
        model = reg.create_tts("test-tts")
        assert isinstance(model, ConcreteTTS)
        assert model.kwargs["vocoder"] == "test-vocoder"

    def test_register_and_create_llm(self):
        reg = _Registry()
        reg.register_llm("test-llm", ConcreteLLM)
        model = reg.create_llm("test-llm")
        assert isinstance(model, ConcreteLLM)

    def test_list_stt(self):
        reg = _Registry()
        reg.register_stt("a", ConcreteSTT)
        reg.register_stt("b", ConcreteSTT)
        assert sorted(reg.list_stt()) == ["a", "b"]

    def test_list_tts(self):
        reg = _Registry()
        reg.register_tts("x", ConcreteTTS)
        assert reg.list_tts() == ["x"]

    def test_list_llm(self):
        reg = _Registry()
        reg.register_llm("llm1", ConcreteLLM)
        assert reg.list_llm() == ["llm1"]

    def test_create_unknown_model_raises(self):
        reg = _Registry()
        with pytest.raises(KeyError):
            reg.create_stt("nonexistent")

    def test_override_kwargs(self):
        reg = _Registry()
        reg.register_stt("test", ConcreteSTT, model_id="default")
        model = reg.create_stt("test", model_id="override")
        assert model.kwargs["model_id"] == "override"
