"""Tests for ModelRouter with fallback scenarios."""

from typing import Any, AsyncIterator
from unittest.mock import MagicMock

import pytest
from openai import APIConnectionError, InternalServerError, RateLimitError
from openai.types.chat import ChatCompletion, ChatCompletionMessage
from openai.types.chat.chat_completion import Choice
from openai.types.completion_usage import CompletionUsage

from swarm.providers.base import ModelProvider
from swarm.providers.factory import ProviderFactory
from swarm.providers.router import ModelRouter
from swarm.providers.types import CompletionRequest


def create_mock_completion(content: str = "Hello") -> ChatCompletion:
    """Create a mock ChatCompletion object."""
    return ChatCompletion(
        id="chatcmpl-test",
        choices=[
            Choice(
                finish_reason="stop",
                index=0,
                message=ChatCompletionMessage(
                    content=content,
                    role="assistant",
                ),
            )
        ],
        created=1234567890,
        model="test-model",
        object="chat.completion",
        usage=CompletionUsage(
            completion_tokens=10,
            prompt_tokens=5,
            total_tokens=15,
        ),
    )


class MockProvider(ModelProvider):
    """Mock provider for testing."""

    def __init__(self, response: ChatCompletion = None, error: Exception = None):
        self._response = response or create_mock_completion()
        self._error = error
        self.call_count = 0

    def create_completion(
        self, request: CompletionRequest, model: str
    ) -> ChatCompletion:
        self.call_count += 1
        if self._error:
            raise self._error
        return self._response

    async def create_completion_async(
        self, request: CompletionRequest, model: str
    ) -> ChatCompletion:
        self.call_count += 1
        if self._error:
            raise self._error
        return self._response

    async def stream_completion(
        self, request: CompletionRequest, model: str
    ) -> AsyncIterator[Any]:
        self.call_count += 1
        if self._error:
            raise self._error
        # Yield mock chunks
        yield {"choices": [{"delta": {"content": "Hello"}}]}
        yield {"choices": [{"delta": {"content": " world"}}]}


class MockFactory(ProviderFactory):
    """Mock factory that returns configured providers."""

    def __init__(self, providers: dict[str, ModelProvider]):
        self._providers = providers
        self._cache = {}

    def get_provider_for_model(self, model: str) -> ModelProvider:
        if "/" not in model:
            raise ValueError(f"Model must include provider prefix: '{model}'")
        prefix = model.split("/")[0]
        if prefix not in self._providers:
            raise ValueError(f"No provider for prefix: {prefix}")
        return self._providers[prefix]


class TestModelRouterSingleModel:
    """Test ModelRouter with single model."""

    def test_create_completion_success(self):
        """Test successful completion with single model."""
        provider = MockProvider(create_mock_completion("Success"))
        factory = MockFactory({"openai": provider})
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o"],
        )

        result = router.create_completion(request)
        assert result.choices[0].message.content == "Success"
        assert provider.call_count == 1

    @pytest.mark.asyncio
    async def test_create_completion_async_success(self):
        """Test successful async completion."""
        provider = MockProvider(create_mock_completion("Async success"))
        factory = MockFactory({"openai": provider})
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o"],
        )

        result = await router.create_completion_async(request)
        assert result.choices[0].message.content == "Async success"
        assert provider.call_count == 1

    @pytest.mark.asyncio
    async def test_stream_completion_success(self):
        """Test successful streaming completion."""
        provider = MockProvider()
        factory = MockFactory({"openai": provider})
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o"],
            stream=True,
        )

        chunks = []
        async for chunk in router.stream_completion(request):
            chunks.append(chunk)

        assert len(chunks) == 2
        assert provider.call_count == 1


class TestModelRouterFallback:
    """Test ModelRouter fallback behavior."""

    def test_fallback_on_connection_error(self):
        """Test fallback when first model has connection error."""
        error_provider = MockProvider(error=APIConnectionError(request=MagicMock()))
        success_provider = MockProvider(create_mock_completion("Fallback success"))
        factory = MockFactory(
            {
                "openai": error_provider,
                "anthropic": success_provider,
            }
        )
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o", "anthropic/claude-3"],
        )

        result = router.create_completion(request)
        assert result.choices[0].message.content == "Fallback success"
        assert error_provider.call_count == 1
        assert success_provider.call_count == 1

    def test_fallback_on_rate_limit(self):
        """Test fallback when first model is rate limited."""
        error_provider = MockProvider(
            error=RateLimitError(
                message="Rate limited",
                response=MagicMock(status_code=429),
                body=None,
            )
        )
        success_provider = MockProvider(create_mock_completion("Rate limit fallback"))
        factory = MockFactory(
            {
                "openai": error_provider,
                "anthropic": success_provider,
            }
        )
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o", "anthropic/claude-3"],
        )

        result = router.create_completion(request)
        assert result.choices[0].message.content == "Rate limit fallback"

    def test_fallback_on_internal_server_error(self):
        """Test fallback when first model returns 500."""
        error_provider = MockProvider(
            error=InternalServerError(
                message="Internal error",
                response=MagicMock(status_code=500),
                body=None,
            )
        )
        success_provider = MockProvider(create_mock_completion("500 fallback"))
        factory = MockFactory(
            {
                "openai": error_provider,
                "anthropic": success_provider,
            }
        )
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o", "anthropic/claude-3"],
        )

        result = router.create_completion(request)
        assert result.choices[0].message.content == "500 fallback"

    @pytest.mark.asyncio
    async def test_async_fallback_on_connection_error(self):
        """Test async fallback when first model fails."""
        error_provider = MockProvider(error=APIConnectionError(request=MagicMock()))
        success_provider = MockProvider(create_mock_completion("Async fallback"))
        factory = MockFactory(
            {
                "openai": error_provider,
                "anthropic": success_provider,
            }
        )
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o", "anthropic/claude-3"],
        )

        result = await router.create_completion_async(request)
        assert result.choices[0].message.content == "Async fallback"

    @pytest.mark.asyncio
    async def test_stream_fallback_on_error(self):
        """Test streaming fallback when first model fails."""
        error_provider = MockProvider(error=APIConnectionError(request=MagicMock()))
        success_provider = MockProvider()
        factory = MockFactory(
            {
                "openai": error_provider,
                "anthropic": success_provider,
            }
        )
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o", "anthropic/claude-3"],
            stream=True,
        )

        chunks = []
        async for chunk in router.stream_completion(request):
            chunks.append(chunk)

        assert len(chunks) == 2
        assert error_provider.call_count == 1
        assert success_provider.call_count == 1


class TestModelRouterAllFail:
    """Test ModelRouter when all models fail."""

    def test_all_models_fail_raises_runtime_error(self):
        """Test RuntimeError when all models fail."""
        error1 = MockProvider(error=APIConnectionError(request=MagicMock()))
        error2 = MockProvider(error=APIConnectionError(request=MagicMock()))
        factory = MockFactory(
            {
                "openai": error1,
                "anthropic": error2,
            }
        )
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o", "anthropic/claude-3"],
        )

        with pytest.raises(RuntimeError) as exc_info:
            router.create_completion(request)

        assert "All models failed" in str(exc_info.value)
        assert error1.call_count == 1
        assert error2.call_count == 1

    @pytest.mark.asyncio
    async def test_all_models_fail_async(self):
        """Test async RuntimeError when all models fail."""
        error1 = MockProvider(error=APIConnectionError(request=MagicMock()))
        error2 = MockProvider(
            error=RateLimitError(
                message="Rate limited",
                response=MagicMock(status_code=429),
                body=None,
            )
        )
        factory = MockFactory(
            {
                "openai": error1,
                "anthropic": error2,
            }
        )
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o", "anthropic/claude-3"],
        )

        with pytest.raises(RuntimeError) as exc_info:
            await router.create_completion_async(request)

        assert "All models failed" in str(exc_info.value)

    @pytest.mark.asyncio
    async def test_all_models_fail_stream(self):
        """Test streaming RuntimeError when all models fail."""
        error1 = MockProvider(error=APIConnectionError(request=MagicMock()))
        error2 = MockProvider(error=APIConnectionError(request=MagicMock()))
        factory = MockFactory(
            {
                "openai": error1,
                "anthropic": error2,
            }
        )
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o", "anthropic/claude-3"],
            stream=True,
        )

        with pytest.raises(RuntimeError) as exc_info:
            async for _ in router.stream_completion(request):
                pass

        assert "All models failed" in str(exc_info.value)


class TestModelRouterMultipleFallbacks:
    """Test ModelRouter with multiple fallback models."""

    def test_triple_fallback(self):
        """Test fallback through three models."""
        error1 = MockProvider(error=APIConnectionError(request=MagicMock()))
        error2 = MockProvider(
            error=RateLimitError(
                message="Rate limited",
                response=MagicMock(status_code=429),
                body=None,
            )
        )
        success = MockProvider(create_mock_completion("Third time's a charm"))
        factory = MockFactory(
            {
                "openai": error1,
                "anthropic": error2,
                "google": success,
            }
        )
        router = ModelRouter(factory=factory)

        request = CompletionRequest(
            messages=[{"role": "user", "content": "Hi"}],
            model=["openai/gpt-4o", "anthropic/claude-3", "google/gemini-pro"],
        )

        result = router.create_completion(request)
        assert result.choices[0].message.content == "Third time's a charm"
        assert error1.call_count == 1
        assert error2.call_count == 1
        assert success.call_count == 1
