"""Tests for the curated model-capability table (#81)."""

from __future__ import annotations

from backend.model_capabilities import (
    CAPABILITY_KEYS,
    get_capabilities_for_models,
    get_model_capabilities,
    model_supports_vision,
)


class TestCapabilityShape:
    def test_every_result_has_all_keys(self):
        for provider, model in [
            ("deepseek", "deepseek-chat"),
            ("gemini", "gemini-2.0-flash"),
            ("qwencloud", "qwen-vl-max"),
            ("openrouter", "unknown-model"),
        ]:
            caps = get_model_capabilities(provider, model)
            assert set(caps.keys()) == set(CAPABILITY_KEYS)
            assert all(isinstance(v, bool) for v in caps.values())


class TestVisionDetection:
    def test_vision_models_detected(self):
        assert model_supports_vision("qwencloud", "qwen-vl-max")
        assert model_supports_vision("openrouter", "openai/gpt-4o")
        assert model_supports_vision("gemini", "gemini-2.0-flash")
        assert model_supports_vision("mistral", "pixtral-large-latest")

    def test_text_models_not_vision(self):
        assert not model_supports_vision("deepseek", "deepseek-chat")
        assert not model_supports_vision("xiaomi", "mimo-v2.5-pro")

    def test_vision_implies_chat(self):
        caps = get_model_capabilities("qwencloud", "qwen-vl-plus")
        assert caps["vision"] is True
        assert caps["chat"] is True


class TestSpecialModalities:
    def test_embeddings(self):
        caps = get_model_capabilities("openai", "text-embedding-3-small")
        assert caps["embeddings"] is True
        assert caps["chat"] is False

    def test_rerank(self):
        caps = get_model_capabilities("nvidia", "nvidia/llama-3.2-nv-rerankqa-1b-v2")
        assert caps["rerank"] is True

    def test_audio_transcription(self):
        caps = get_model_capabilities("xiaomi", "mimo-v2.5-asr")
        assert caps["audio_transcription"] is True

    def test_audio_speech(self):
        caps = get_model_capabilities("xiaomi", "mimo-v2.5-tts")
        assert caps["audio_speech"] is True

    def test_image_generation(self):
        caps = get_model_capabilities("openai", "dall-e-3")
        assert caps["images"] is True


class TestProviderDefaults:
    def test_unknown_model_uses_provider_default(self):
        caps = get_model_capabilities("deepseek", "some-new-model")
        assert caps["chat"] is True
        assert caps["vision"] is False

    def test_unknown_provider_defaults_to_chat(self):
        caps = get_model_capabilities("nope", "mystery")
        assert caps["chat"] is True

    def test_empty_model_uses_provider_default(self):
        assert get_model_capabilities("gemini", "")["vision"] is True


class TestBatch:
    def test_capabilities_for_models_map(self):
        result = get_capabilities_for_models(
            "qwencloud", ["qwen-max", "qwen-vl-max"]
        )
        assert result["qwen-max"]["vision"] is False
        assert result["qwen-vl-max"]["vision"] is True