"""Tests for the main rotator orchestrator."""

from __future__ import annotations

import pytest

from llm_rotator._types import LLMResponse, RoutingContext, Usage
from llm_rotator.config import RotatorConfig
from llm_rotator.exceptions import (
    AllAttemptsFailedError,
    KeyDeadError,
    ModelRateLimitError,
    ServerError,
)
from llm_rotator.rotator import LLMRotator
from tests.conftest import FakeLLMClient


def _resp(model: str = "gpt-4o", content: str = "ok", total: int = 20) -> LLMResponse:
    return LLMResponse(
        content=content,
        usage=Usage(
            prompt_tokens=total // 2,
            completion_tokens=total - total // 2,
            total_tokens=total,
        ),
        model=model,
        provider="fake",
        key_alias="unknown",
    )


def _config_one_provider(
    models: list[str] | None = None,
    model_groups: list[dict] | None = None,
    keys: list[dict] | None = None,
    retry: dict | None = None,
) -> RotatorConfig:
    provider = {
        "name": "openai",
        "client_type": "openai",
        "priority": 1,
        "keys": keys or [{"token": "sk-1", "alias": "key1"}],
    }
    if model_groups:
        provider["model_groups"] = model_groups
    else:
        provider["models"] = models or ["gpt-4o"]
    if retry:
        provider["server_error_retry"] = retry
    return RotatorConfig(providers=[provider])


def _config_two_providers() -> RotatorConfig:
    return RotatorConfig(
        providers=[
            {
                "name": "openai",
                "client_type": "openai",
                "priority": 1,
                "model_groups": [
                    {"name": "flagship", "tier": 1, "models": ["gpt-5", "gpt-4o"]},
                    {"name": "mini", "tier": 3, "models": ["gpt-4o-mini"]},
                ],
                "keys": [
                    {"token": "sk-1", "alias": "key1"},
                    {"token": "sk-2", "alias": "key2"},
                ],
            },
            {
                "name": "gemini",
                "client_type": "gemini",
                "priority": 2,
                "models": ["gemini-2.0-flash"],
                "keys": [{"token": "AIza-1", "alias": "gkey1"}],
            },
        ]
    )


def _config_tiered() -> RotatorConfig:
    return RotatorConfig(
        providers=[
            {
                "name": "openai",
                "client_type": "openai",
                "priority": 1,
                "model_groups": [
                    {"name": "flagship", "tier": 1, "models": ["gpt-5"]},
                    {"name": "standard", "tier": 2, "models": ["gpt-4o"]},
                    {"name": "mini", "tier": 3, "models": ["gpt-4o-mini"]},
                ],
                "keys": [{"token": "sk-1", "alias": "key1"}],
            },
        ]
    )


class TestRotatorHappyPath:
    async def test_happy_path(self):
        config = _config_one_provider()
        client = FakeLLMClient()
        client.enqueue(_resp(model="gpt-4o"))

        rotator = LLMRotator(config, clients={"openai": client})
        result = await rotator.complete(messages=[{"role": "user", "content": "hi"}])

        assert result.content == "ok"
        assert result.model == "gpt-4o"
        assert len(client.calls) == 1


class TestRotatorModelFirst:
    async def test_model_first_all_keys_before_downgrade(self):
        """On 429 for gpt-5/key1, try gpt-5/key2 before downgrading."""
        config = _config_two_providers()
        client = FakeLLMClient()
        # gpt-5/key1 → 429
        client.enqueue(
            ModelRateLimitError(key_alias="key1", model="gpt-5", retry_after=60)
        )
        # gpt-5/key2 → success
        client.enqueue(_resp(model="gpt-5", content="from key2"))

        rotator = LLMRotator(config, clients={"openai": client, "gemini": client})
        result = await rotator.complete(messages=[{"role": "user", "content": "hi"}])

        assert result.content == "from key2"
        assert len(client.calls) == 2
        assert client.calls[0]["api_key"] == "sk-1"
        assert client.calls[1]["api_key"] == "sk-2"

    async def test_downgrade_only_after_all_keys_exhausted(self):
        """Downgrade to gpt-4o only when all keys for gpt-5 are blocked."""
        config = _config_two_providers()
        client = FakeLLMClient()
        # gpt-5/key1 → 429
        client.enqueue(
            ModelRateLimitError(key_alias="key1", model="gpt-5", retry_after=60)
        )
        # gpt-5/key2 → 429
        client.enqueue(
            ModelRateLimitError(key_alias="key2", model="gpt-5", retry_after=60)
        )
        # gpt-4o/key1 → success
        client.enqueue(_resp(model="gpt-4o", content="downgraded"))

        rotator = LLMRotator(config, clients={"openai": client, "gemini": client})
        result = await rotator.complete(messages=[{"role": "user", "content": "hi"}])

        assert result.content == "downgraded"
        assert len(client.calls) == 3


class TestRotatorProviderFallback:
    async def test_provider_fallback(self):
        """If all OpenAI combos fail, fall back to Gemini."""
        config = _config_two_providers()
        openai_client = FakeLLMClient()
        gemini_client = FakeLLMClient()

        # All OpenAI combos fail (6 combos: 3 models x 2 keys)
        for _ in range(6):
            openai_client.enqueue(
                ModelRateLimitError(key_alias="x", model="x", retry_after=60)
            )
        # Gemini succeeds
        gemini_client.enqueue(_resp(model="gemini-2.0-flash", content="from gemini"))

        rotator = LLMRotator(
            config, clients={"openai": openai_client, "gemini": gemini_client}
        )
        result = await rotator.complete(messages=[{"role": "user", "content": "hi"}])
        assert result.content == "from gemini"


class TestRotatorServerErrorRetry:
    async def test_server_error_retry(self):
        """500 triggers one retry, then succeeds."""
        config = _config_one_provider(retry={"max_attempts": 1, "delay_seconds": 0.0})
        client = FakeLLMClient()
        client.enqueue(ServerError(provider="openai", status_code=500))
        client.enqueue(_resp(content="retry success"))

        rotator = LLMRotator(config, clients={"openai": client})
        result = await rotator.complete(messages=[{"role": "user", "content": "hi"}])

        assert result.content == "retry success"
        assert len(client.calls) == 2

    async def test_retry_disabled(self):
        """max_attempts=0: no retry, move to next candidate."""
        config = _config_one_provider(
            models=["gpt-4o", "gpt-5"],
            retry={"max_attempts": 0, "delay_seconds": 0.0},
        )
        client = FakeLLMClient()
        client.enqueue(ServerError(provider="openai", status_code=500))
        client.enqueue(_resp(model="gpt-5", content="next model"))

        rotator = LLMRotator(config, clients={"openai": client})
        result = await rotator.complete(messages=[{"role": "user", "content": "hi"}])

        assert result.content == "next model"


class TestRotatorAllFailed:
    async def test_all_failed_exception_group(self):
        """When all combos fail, AllAttemptsFailedError is raised."""
        config = _config_one_provider(
            keys=[{"token": "sk-1", "alias": "key1"}],
        )
        client = FakeLLMClient()
        client.enqueue(KeyDeadError(key_alias="key1", status_code=402))

        rotator = LLMRotator(config, clients={"openai": client})

        with pytest.raises(AllAttemptsFailedError) as exc_info:
            await rotator.complete(messages=[{"role": "user", "content": "hi"}])

        assert len(exc_info.value.exceptions) >= 1


class TestRotatorRoutingContext:
    async def test_routing_context_model_group(self):
        """model_group restricts to that specific group."""
        config = _config_two_providers()
        client = FakeLLMClient()
        client.enqueue(_resp(model="gpt-4o-mini", content="from mini"))

        rotator = LLMRotator(config, clients={"openai": client, "gemini": client})
        result = await rotator.complete(
            messages=[{"role": "user", "content": "hi"}],
            routing=RoutingContext(model_group="mini"),
        )

        assert result.content == "from mini"
        assert client.calls[0]["model"] == "gpt-4o-mini"

    async def test_routing_context_allowed_providers(self):
        """allowed_providers filters out other providers."""
        config = _config_two_providers()
        openai_client = FakeLLMClient()
        gemini_client = FakeLLMClient()
        gemini_client.enqueue(_resp(model="gemini-2.0-flash", content="gemini only"))

        rotator = LLMRotator(
            config, clients={"openai": openai_client, "gemini": gemini_client}
        )
        result = await rotator.complete(
            messages=[{"role": "user", "content": "hi"}],
            routing=RoutingContext(allowed_providers=["gemini"]),
        )

        assert result.content == "gemini only"
        assert len(openai_client.calls) == 0

    async def test_routing_context_disallow_downgrade(self):
        """allow_downgrade=False: only models from the specified group, no fallback to others."""
        config = _config_two_providers()
        client = FakeLLMClient()
        # All flagship candidates fail (gpt-5/key1, gpt-5/key2, gpt-4o/key1, gpt-4o/key2)
        for _ in range(4):
            client.enqueue(
                ModelRateLimitError(key_alias="x", model="x", retry_after=60)
            )

        rotator = LLMRotator(config, clients={"openai": client, "gemini": client})

        with pytest.raises(AllAttemptsFailedError):
            await rotator.complete(
                messages=[{"role": "user", "content": "hi"}],
                routing=RoutingContext(
                    model_group="flagship", allow_downgrade=False
                ),
            )

        # Should NOT have tried mini or gemini models
        models_tried = {c["model"] for c in client.calls}
        assert "gpt-4o-mini" not in models_tried
        assert "gemini-2.0-flash" not in models_tried


class TestRotatorTier:
    async def test_tier_ceiling(self):
        """tier=3 skips flagship (1) and standard (2), uses only mini (3)."""
        config = _config_tiered()
        client = FakeLLMClient()
        client.enqueue(_resp(model="gpt-4o-mini", content="mini only"))

        rotator = LLMRotator(config, clients={"openai": client})
        result = await rotator.complete(
            messages=[{"role": "user", "content": "hi"}],
            routing=RoutingContext(tier=3),
        )

        assert result.content == "mini only"
        assert client.calls[0]["model"] == "gpt-4o-mini"

    async def test_tier_fallback_down_only(self):
        """tier=2: if standard fails, fall to mini (3), NOT up to flagship (1)."""
        config = _config_tiered()
        client = FakeLLMClient()
        # standard/gpt-4o fails
        client.enqueue(
            ModelRateLimitError(key_alias="key1", model="gpt-4o", retry_after=60)
        )
        # mini/gpt-4o-mini succeeds
        client.enqueue(_resp(model="gpt-4o-mini", content="fell to mini"))

        rotator = LLMRotator(config, clients={"openai": client})
        result = await rotator.complete(
            messages=[{"role": "user", "content": "hi"}],
            routing=RoutingContext(tier=2),
        )

        assert result.content == "fell to mini"
        # Should NOT have tried flagship (tier=1)
        assert all(c["model"] != "gpt-5" for c in client.calls)

    async def test_tier_cross_provider(self):
        """tier filters groups across all providers."""
        config = RotatorConfig(
            providers=[
                {
                    "name": "openai",
                    "client_type": "openai",
                    "priority": 1,
                    "model_groups": [
                        {"name": "flagship", "tier": 1, "models": ["gpt-5"]},
                        {"name": "mini", "tier": 3, "models": ["gpt-4o-mini"]},
                    ],
                    "keys": [{"token": "sk-1", "alias": "key1"}],
                },
                {
                    "name": "gemini",
                    "client_type": "gemini",
                    "priority": 2,
                    "model_groups": [
                        {"name": "pro", "tier": 2, "models": ["gemini-2.0-pro"]},
                    ],
                    "keys": [{"token": "AIza-1", "alias": "gkey1"}],
                },
            ]
        )
        openai_client = FakeLLMClient()
        gemini_client = FakeLLMClient()
        gemini_client.enqueue(
            _resp(model="gemini-2.0-pro", content="gemini pro")
        )

        rotator = LLMRotator(
            config, clients={"openai": openai_client, "gemini": gemini_client}
        )
        await rotator.complete(
            messages=[{"role": "user", "content": "hi"}],
            routing=RoutingContext(tier=2),
        )

        # Flagship (tier=1) should be skipped
        assert all(c["model"] != "gpt-5" for c in openai_client.calls)

    async def test_tier_default_uses_all(self):
        """Without tier (default=1), all models are available starting from best."""
        config = _config_tiered()
        client = FakeLLMClient()
        client.enqueue(_resp(model="gpt-5", content="best model"))

        rotator = LLMRotator(config, clients={"openai": client})
        result = await rotator.complete(
            messages=[{"role": "user", "content": "hi"}],
        )

        assert result.content == "best model"
        assert client.calls[0]["model"] == "gpt-5"


class TestRotatorQuotaIntegration:
    async def test_quota_exceeded_skips_without_http(self):
        """If quota is exhausted, candidate is skipped without HTTP call."""
        config = RotatorConfig(
            providers=[
                {
                    "name": "openai",
                    "client_type": "openai",
                    "priority": 1,
                    "model_groups": [
                        {
                            "name": "flagship",
                            "tier": 1,
                            "models": ["gpt-5"],
                            "token_quota": {
                                "limit": 10,
                                "reset": "daily_utc",
                            },
                        },
                        {
                            "name": "mini",
                            "tier": 3,
                            "models": ["gpt-4o-mini"],
                        },
                    ],
                    "keys": [{"token": "sk-1", "alias": "key1"}],
                },
            ]
        )
        client = FakeLLMClient()
        # First call succeeds and uses up quota
        client.enqueue(_resp(model="gpt-5", content="first", total=20))
        # Second call should skip flagship (quota exceeded) and go to mini
        client.enqueue(_resp(model="gpt-4o-mini", content="mini fallback"))

        rotator = LLMRotator(config, clients={"openai": client})

        # First call — uses flagship
        result1 = await rotator.complete(
            messages=[{"role": "user", "content": "hi"}]
        )
        assert result1.model == "gpt-5"

        # Second call — flagship quota exhausted, falls to mini
        result2 = await rotator.complete(
            messages=[{"role": "user", "content": "hi"}]
        )
        assert result2.model == "gpt-4o-mini"
