"""Tests for streaming in rotator."""

from __future__ import annotations

from collections.abc import AsyncIterator

import pytest

from llm_rotator._types import LLMResponse, RoutingContext, StreamChunk, Usage
from llm_rotator.clients import AbstractLLMClient
from llm_rotator.config import RotatorConfig
from llm_rotator.exceptions import (
    AllAttemptsFailedError,
    ModelRateLimitError,
    ServerError,
)
from llm_rotator.rotator import LLMRotator


def _simple_config() -> RotatorConfig:
    return RotatorConfig(
        providers=[
            {
                "name": "openai",
                "client_type": "openai",
                "priority": 1,
                "models": ["gpt-4o", "gpt-5"],
                "keys": [
                    {"token": "sk-1", "alias": "key1"},
                    {"token": "sk-2", "alias": "key2"},
                ],
            }
        ]
    )


class FakeStreamClient(AbstractLLMClient):
    """Client that returns configurable stream results."""

    def __init__(self) -> None:
        self._stream_queue: list[list[StreamChunk] | Exception] = []
        self.stream_calls: list[dict] = []

    def enqueue_stream(self, chunks: list[StreamChunk]) -> None:
        self._stream_queue.append(chunks)

    def enqueue_stream_error(self, error: Exception) -> None:
        self._stream_queue.append(error)

    async def generate(self, messages, model, api_key, base_url=None) -> LLMResponse:
        raise NotImplementedError

    async def stream(
        self, messages, model, api_key, base_url=None
    ) -> AsyncIterator[StreamChunk]:
        self.stream_calls.append({
            "model": model,
            "api_key": api_key,
        })
        if not self._stream_queue:
            yield StreamChunk(delta="default", done=True)
            return

        result = self._stream_queue.pop(0)
        if isinstance(result, Exception):
            raise result

        for chunk in result:
            yield chunk


class TestRotatorStreamHappyPath:
    async def test_stream_basic(self):
        config = _simple_config()
        client = FakeStreamClient()
        client.enqueue_stream([
            StreamChunk(delta="Hello"),
            StreamChunk(delta=" world"),
            StreamChunk(
                delta="",
                usage=Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7),
                done=True,
            ),
        ])

        rotator = LLMRotator(config, clients={"openai": client})
        chunks = []
        async for chunk in rotator.stream(
            messages=[{"role": "user", "content": "hi"}]
        ):
            chunks.append(chunk)

        assert len(chunks) == 3
        assert chunks[0].delta == "Hello"
        assert chunks[1].delta == " world"
        assert chunks[2].done is True


class TestRotatorStreamFallback:
    async def test_mid_stream_error_retries_next_candidate(self):
        """If stream fails mid-way, retry from scratch on next candidate."""
        config = _simple_config()
        client = FakeStreamClient()

        # First candidate: error (before any chunks)
        client.enqueue_stream_error(
            ModelRateLimitError(key_alias="key1", model="gpt-4o", retry_after=60)
        )
        # Second candidate: success
        client.enqueue_stream([
            StreamChunk(delta="from key2"),
            StreamChunk(delta="", done=True),
        ])

        rotator = LLMRotator(config, clients={"openai": client})
        chunks = []
        async for chunk in rotator.stream(
            messages=[{"role": "user", "content": "hi"}]
        ):
            chunks.append(chunk)

        assert chunks[0].delta == "from key2"
        assert len(client.stream_calls) == 2
        assert client.stream_calls[0]["api_key"] == "sk-1"
        assert client.stream_calls[1]["api_key"] == "sk-2"

    async def test_server_error_retry_then_fallback(self):
        """500 triggers retry, then falls to next candidate."""
        config = RotatorConfig(
            providers=[
                {
                    "name": "openai",
                    "client_type": "openai",
                    "priority": 1,
                    "models": ["gpt-4o"],
                    "keys": [
                        {"token": "sk-1", "alias": "key1"},
                        {"token": "sk-2", "alias": "key2"},
                    ],
                    "server_error_retry": {
                        "max_attempts": 1,
                        "delay_seconds": 0.0,
                    },
                }
            ]
        )
        client = FakeStreamClient()

        # key1 attempt 1: 500
        client.enqueue_stream_error(
            ServerError(provider="openai", status_code=500)
        )
        # key1 attempt 2 (retry): 500 again
        client.enqueue_stream_error(
            ServerError(provider="openai", status_code=500)
        )
        # key2: success
        client.enqueue_stream([
            StreamChunk(delta="ok from key2", done=True),
        ])

        rotator = LLMRotator(config, clients={"openai": client})
        chunks = []
        async for chunk in rotator.stream(
            messages=[{"role": "user", "content": "hi"}]
        ):
            chunks.append(chunk)

        assert chunks[0].delta == "ok from key2"
        assert len(client.stream_calls) == 3  # 2 retries + 1 success


class TestRotatorStreamAllFailed:
    async def test_all_failed_raises(self):
        config = RotatorConfig(
            providers=[
                {
                    "name": "openai",
                    "client_type": "openai",
                    "priority": 1,
                    "models": ["gpt-4o"],
                    "keys": [{"token": "sk-1", "alias": "key1"}],
                    "server_error_retry": {
                        "max_attempts": 0,
                        "delay_seconds": 0.0,
                    },
                }
            ]
        )
        client = FakeStreamClient()
        client.enqueue_stream_error(
            ServerError(provider="openai", status_code=500)
        )

        rotator = LLMRotator(config, clients={"openai": client})
        with pytest.raises(AllAttemptsFailedError):
            async for _ in rotator.stream(
                messages=[{"role": "user", "content": "hi"}]
            ):
                pass


class TestRotatorStreamWithRouting:
    async def test_stream_respects_tier(self):
        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"}],
                }
            ]
        )
        client = FakeStreamClient()
        client.enqueue_stream([
            StreamChunk(delta="mini response", done=True),
        ])

        rotator = LLMRotator(config, clients={"openai": client})
        chunks = []
        async for chunk in rotator.stream(
            messages=[{"role": "user", "content": "hi"}],
            routing=RoutingContext(tier=3),
        ):
            chunks.append(chunk)

        assert chunks[0].delta == "mini response"
        assert client.stream_calls[0]["model"] == "gpt-4o-mini"
