from unittest.mock import MagicMock

import pytest

from rhodecode.apps.ai_agents.models.claude import ClaudeService
from rhodecode.apps.ai_agents.models.gemini import GeminiService
from rhodecode.apps.ai_agents.models.gpt import GPTService


class TestModelVersionFiltering:
    def test_filter_GPT_only_text_versions(self):
        service = GPTService(MagicMock())
        input_versions = [
            ("o3", "o3"),
            ("o4-mini", "o4 mini"),
            ("sora-2-pro", "sora-2 pro"),
            ("sora-2", "sora-2"),
            ("gpt-audio-mini-2025-10-06", "gpt audio mini 2025-10-06"),
            ("dall-e-3", "dall-e 3"),
            ("dall-e-2", "dall-e 2"),
            ("gpt-realtime", "gpt realtime"),
            ("gpt-realtime-2025-08-28", "gpt realtime 2025-08-28"),
            ("gpt-4o-audio-preview-2025-06-03", "gpt-4o audio preview 2025-06-03"),
            ("gpt-image-1", "gpt image-1"),
            ("gpt-image-1-mini", "gpt image-1 mini"),
            ("gpt-5-codex", "gpt-5 codex"),
            ("gpt-4o-mini-search-preview-2025-03-11", "gpt-4o mini search preview 2025-03-11"),
            ("gpt-4o-mini-search-preview", "gpt-4o mini search preview"),
            ("babbage-002", "babbage 002"),
            ("omni-moderation-latest", "omni moderation latest"),
            ("omni-moderation-2024-09-26", "omni moderation 2024-09-26"),
            ("davinci-002", "davinci 002"),
            ("gpt-4o-transcribe", "gpt-4o transcribe"),
            ("gpt-4o-mini-transcribe", "gpt-4o mini transcribe"),
            ("gpt-4o-mini-tts", "gpt-4o mini tts"),
            ("whisper-1", "whisper-1"),
            ("text-embedding-ada-002", "text embedding ada 002"),
        ]

        res = service._filter_api_models(input_versions)
        expected_output = [
            ("o3", "o3"),
            ("o4-mini", "o4 mini"),
        ]

        assert res == expected_output

    def test_filter_Gemini_only_text_versions(self):
        service = GeminiService(MagicMock())
        input_versions = [
            ("gemini-2.5-pro-preview-03-25", "Gemini 2.5 Pro Preview 03-25"),
            ("veo-2.0-generate-001", "Veo 2"),
            ("veo-3.0-generate-preview", "Veo 3"),
            ("gemini-2.5-flash-preview-native-audio-dialog", "Gemini 2.5 Flash Preview Native Audio Dialog"),
            ("lyria-realtime-exp", "Lyria Realtime Experimental"),
            ("aqa", "Model that performs Attributed Question Answering."),
            ("learnlm-2.0-flash-experimental", "LearnLM 2.0 Flash Experimental"),
            ("imagen-3.0-generate-002", "Imagen 3.0"),
            ("imagen-4.0-generate-preview-06-06", "Imagen 4 (Preview)"),
            ("gemini-2.5-flash-image-preview", "Nano Banana"),
            ("gemini-2.5-flash-image", "Nano Banana"),
            ("gemini-embedding-exp-03-07", "Gemini Embedding Experimental 03-07"),
            ("gemini-embedding-exp", "Gemini Embedding Experimental"),
            ("gemini-2.5-computer-use-preview-10-2025", "Gemini 2.5 Computer Use Preview 10-2025"),
            ("gemini-robotics-er-1.5-preview", "Gemini Robotics-ER 1.5 Preview"),
        ]

        res = service._filter_api_models(input_versions)
        expected_output = [
            ("gemini-2.5-pro-preview-03-25", "Gemini 2.5 Pro Preview 03-25"),
        ]

        assert res == expected_output

    def test_filter_Claude_only_text_versions(self):
        service = ClaudeService(MagicMock())
        input_versions = [
            ("claude-sonnet-4-5-20250929", "Claude Sonnet 4.5"),
            ("claude-opus-4-1-20250805", "Claude Opus 4.1"),
            ("claude-opus-4-20250514", "Claude Opus 4"),
            ("claude-sonnet-4-20250514", "Claude Sonnet 4"),
            ("claude-3-7-sonnet-20250219", "Claude Sonnet 3.7"),
            ("claude-3-5-haiku-20241022", "Claude Haiku 3.5"),
            ("claude-3-haiku-20240307", "Claude Haiku 3"),
        ]

        res = service._filter_api_models(input_versions)
        expected_output = [
            ("claude-sonnet-4-5-20250929", "Claude Sonnet 4.5"),
            ("claude-opus-4-1-20250805", "Claude Opus 4.1"),
            ("claude-opus-4-20250514", "Claude Opus 4"),
            ("claude-sonnet-4-20250514", "Claude Sonnet 4"),
            ("claude-3-7-sonnet-20250219", "Claude Sonnet 3.7"),
            ("claude-3-5-haiku-20241022", "Claude Haiku 3.5"),
            ("claude-3-haiku-20240307", "Claude Haiku 3"),
        ]

        assert res == expected_output
