"""Unit tests for LLMRotator.embed() — rotation, circuit breaker, quota, fallback."""

from __future__ import annotations

from unittest.mock import AsyncMock

import pytest

from llm_rotator._types import EmbeddingResponse, EmbeddingUsage
from llm_rotator.clients import AbstractLLMClient
from llm_rotator.config import RotatorConfig
from llm_rotator.exceptions import (
    AllAttemptsFailedError,
    EmbeddingNotSupportedError,
    ServerError,
)
from llm_rotator.rotator import LLMRotator


def _make_embedding_response(
    provider: str = "openai", model: str = "text-embedding-3-small"
) -> EmbeddingResponse:
    return EmbeddingResponse(
        embeddings=[[0.1, 0.2, 0.3]],
        usage=EmbeddingUsage(prompt_tokens=5, total_tokens=5),
        model=model,
        provider=provider,
        key_alias="test...",
    )


def _two_provider_config() -> RotatorConfig:
    return RotatorConfig(
        providers=[
            {
                "name": "provider_a",
                "client_type": "openai",
                "priority": 1,
                "base_url": "https://api-a.example.com/v1",
                "models": ["text-embedding-3-small"],
                "keys": [{"token": "sk-a1", "alias": "a_key1"}],
            },
            {
                "name": "provider_b",
                "client_type": "gemini",
                "priority": 2,
                "models": ["text-embedding-004"],
                "keys": [{"token": "gk-b1", "alias": "b_key1"}],
            },
        ]
    )


class TestRotatorEmbedHappyPath:
    async def test_embed_success(self):
        mock_client = AsyncMock(spec=AbstractLLMClient)
        mock_client.embed.return_value = _make_embedding_response()

        config = RotatorConfig(
            providers=[
                {
                    "name": "p1",
                    "client_type": "openai",
                    "priority": 1,
                    "models": ["text-embedding-3-small"],
                    "keys": [{"token": "sk-1", "alias": "key1"}],
                }
            ]
        )
        rotator = LLMRotator(config, clients={"openai": mock_client})
        result = await rotator.embed(input=["hello"])

        assert isinstance(result, EmbeddingResponse)
        assert result.embeddings == [[0.1, 0.2, 0.3]]
        mock_client.embed.assert_called_once()

    async def test_embed_string_input_normalized(self):
        """Single string input is auto-wrapped in a list."""
        mock_client = AsyncMock(spec=AbstractLLMClient)
        mock_client.embed.return_value = _make_embedding_response()

        config = RotatorConfig(
            providers=[
                {
                    "name": "p1",
                    "client_type": "openai",
                    "priority": 1,
                    "models": ["emb-model"],
                    "keys": [{"token": "sk-1", "alias": "key1"}],
                }
            ]
        )
        rotator = LLMRotator(config, clients={"openai": mock_client})
        await rotator.embed(input="single string")

        call_kwargs = mock_client.embed.call_args
        assert call_kwargs.kwargs["input"] == ["single string"]


class TestRotatorEmbedFallback:
    async def test_fallback_on_server_error(self):
        """First provider fails with ServerError, falls back to second."""
        client_a = AsyncMock(spec=AbstractLLMClient)
        client_a.embed.side_effect = ServerError(provider="openai", status_code=500)

        client_b = AsyncMock(spec=AbstractLLMClient)
        client_b.embed.return_value = _make_embedding_response(
            provider="gemini", model="text-embedding-004"
        )

        config = _two_provider_config()
        rotator = LLMRotator(config, clients={"openai": client_a, "gemini": client_b})
        result = await rotator.embed(input=["test"])

        assert result.provider == "gemini"
        assert result.model == "text-embedding-004"

    async def test_fallback_on_embedding_not_supported(self):
        """Provider that doesn't support embeddings is skipped."""
        client_a = AsyncMock(spec=AbstractLLMClient)
        client_a.embed.side_effect = EmbeddingNotSupportedError("openai")

        client_b = AsyncMock(spec=AbstractLLMClient)
        client_b.embed.return_value = _make_embedding_response(
            provider="gemini", model="text-embedding-004"
        )

        config = _two_provider_config()
        rotator = LLMRotator(config, clients={"openai": client_a, "gemini": client_b})
        result = await rotator.embed(input=["test"])

        assert result.provider == "gemini"


class TestRotatorEmbedAllFail:
    async def test_all_fail_raises(self):
        mock_client = AsyncMock(spec=AbstractLLMClient)
        mock_client.embed.side_effect = ServerError(provider="openai", status_code=500)

        config = RotatorConfig(
            providers=[
                {
                    "name": "p1",
                    "client_type": "openai",
                    "priority": 1,
                    "models": ["emb-model"],
                    "keys": [{"token": "sk-1", "alias": "key1"}],
                }
            ]
        )
        rotator = LLMRotator(config, clients={"openai": mock_client})

        with pytest.raises(AllAttemptsFailedError):
            await rotator.embed(input=["test"])


class TestRotatorEmbedQuota:
    async def test_quota_exhaustion_fallback(self):
        """When token quota is exceeded, embed falls back to next group."""
        mock_client = AsyncMock(spec=AbstractLLMClient)

        call_count = 0

        async def embed_side_effect(**kwargs):
            nonlocal call_count
            call_count += 1
            return EmbeddingResponse(
                embeddings=[[0.1]],
                usage=EmbeddingUsage(prompt_tokens=15, total_tokens=15),
                model=kwargs["model"],
                provider="openai",
                key_alias="test...",
            )

        mock_client.embed.side_effect = embed_side_effect

        config = RotatorConfig(
            providers=[
                {
                    "name": "p1",
                    "client_type": "openai",
                    "priority": 1,
                    "model_groups": [
                        {
                            "name": "premium",
                            "tier": 1,
                            "models": ["text-embedding-3-large"],
                            "token_quota": {"limit": 20, "reset": "daily_utc"},
                        },
                        {
                            "name": "basic",
                            "tier": 2,
                            "models": ["text-embedding-3-small"],
                        },
                    ],
                    "keys": [{"token": "sk-1", "alias": "key1"}],
                }
            ]
        )
        rotator = LLMRotator(config, clients={"openai": mock_client})

        r1 = await rotator.embed(input=["test1"])
        assert r1.model == "text-embedding-3-large"

        r2 = await rotator.embed(input=["test2"])
        assert r2.model == "text-embedding-3-large"

        # After 2 calls: 30 tokens used, quota is 20 → falls to basic
        r3 = await rotator.embed(input=["test3"])
        assert r3.model == "text-embedding-3-small"
