"""Unit tests for getpatter.services.fallback_provider — FallbackLLMProvider."""

from __future__ import annotations

import asyncio
from typing import AsyncIterator

import pytest

from getpatter.services.fallback_provider import (
    AllProvidersFailedError,
    FallbackLLMProvider,
    PartialStreamError,
)


# ---------------------------------------------------------------------------
# Mock providers
# ---------------------------------------------------------------------------


class SucceedingProvider:
    """A provider that yields predetermined chunks."""

    def __init__(self, chunks: list[dict]) -> None:
        self._chunks = list(chunks)

    async def stream(
        self, messages: list[dict], tools: list[dict] | None = None, **_kwargs
    ) -> AsyncIterator[dict]:
        for chunk in self._chunks:
            yield chunk


class FailingProvider:
    """A provider that always raises on stream."""

    def __init__(self, error: Exception | None = None) -> None:
        self._error = error or RuntimeError("provider failed")

    async def stream(
        self, messages: list[dict], tools: list[dict] | None = None, **_kwargs
    ) -> AsyncIterator[dict]:
        raise self._error
        # Make it an async generator
        yield {}  # type: ignore[unreachable]  # pragma: no cover


class PartialThenFailProvider:
    """A provider that yields some chunks then raises."""

    def __init__(self, chunks: list[dict], error: Exception | None = None) -> None:
        self._chunks = list(chunks)
        self._error = error or RuntimeError("mid-stream failure")

    async def stream(
        self, messages: list[dict], tools: list[dict] | None = None, **_kwargs
    ) -> AsyncIterator[dict]:
        for chunk in self._chunks:
            yield chunk
        raise self._error


class FailNTimesThenSucceed:
    """A provider that fails N times, then succeeds."""

    def __init__(self, fail_count: int, chunks: list[dict]) -> None:
        self._fail_count = fail_count
        self._chunks = list(chunks)
        self._calls = 0

    async def stream(
        self, messages: list[dict], tools: list[dict] | None = None, **_kwargs
    ) -> AsyncIterator[dict]:
        self._calls += 1
        if self._calls <= self._fail_count:
            raise RuntimeError(f"fail #{self._calls}")
        for chunk in self._chunks:
            yield chunk


class CallCountingFailingProvider:
    """A provider that always fails but counts calls."""

    def __init__(self) -> None:
        self.call_count = 0

    async def stream(
        self, messages: list[dict], tools: list[dict] | None = None, **_kwargs
    ) -> AsyncIterator[dict]:
        self.call_count += 1
        raise RuntimeError("always fails")
        yield {}  # type: ignore[unreachable]  # pragma: no cover


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


async def collect_chunks(
    provider: FallbackLLMProvider,
    messages: list[dict] | None = None,
) -> list[dict]:
    """Collect all chunks from a fallback provider stream."""
    if messages is None:
        messages = [{"role": "user", "content": "hi"}]
    chunks: list[dict] = []
    async for chunk in provider.stream(messages):
        chunks.append(chunk)
    return chunks


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


@pytest.mark.unit
class TestFallbackLLMProvider:
    """Tests for FallbackLLMProvider."""

    # 1. Primary provider succeeds -> returns its output
    @pytest.mark.asyncio
    async def test_primary_succeeds(self) -> None:
        primary = SucceedingProvider(
            [
                {"type": "text", "content": "Hello"},
                {"type": "text", "content": " world"},
            ]
        )
        fallback = SucceedingProvider(
            [
                {"type": "text", "content": "Fallback"},
            ]
        )

        provider = FallbackLLMProvider([primary, fallback])
        chunks = await collect_chunks(provider)

        assert chunks == [
            {"type": "text", "content": "Hello"},
            {"type": "text", "content": " world"},
        ]
        provider.destroy()

    # 2. Primary fails -> fallback succeeds -> returns fallback output
    @pytest.mark.asyncio
    async def test_fallback_on_primary_failure(self) -> None:
        primary = FailingProvider()
        fallback = SucceedingProvider(
            [
                {"type": "text", "content": "Fallback response"},
            ]
        )

        provider = FallbackLLMProvider([primary, fallback])
        chunks = await collect_chunks(provider)

        assert chunks == [{"type": "text", "content": "Fallback response"}]
        provider.destroy()

    # 3. Both fail -> throws error
    @pytest.mark.asyncio
    async def test_all_providers_fail(self) -> None:
        primary = FailingProvider(RuntimeError("primary down"))
        fallback = FailingProvider(RuntimeError("fallback down"))

        provider = FallbackLLMProvider([primary, fallback])

        with pytest.raises(AllProvidersFailedError):
            await collect_chunks(provider)

        provider.destroy()

    # 4. Primary fails after yielding tokens -> throws (no retry)
    @pytest.mark.asyncio
    async def test_partial_stream_error(self) -> None:
        primary = PartialThenFailProvider(
            [
                {"type": "text", "content": "partial"},
            ]
        )
        fallback = SucceedingProvider(
            [
                {"type": "text", "content": "Fallback"},
            ]
        )

        provider = FallbackLLMProvider([primary, fallback])

        with pytest.raises(PartialStreamError):
            await collect_chunks(provider)

        provider.destroy()

    # 5. Primary recovers after being marked unavailable
    @pytest.mark.asyncio
    async def test_recovery_after_unavailable(self) -> None:
        recovering = FailNTimesThenSucceed(
            1,
            [
                {"type": "text", "content": "recovered"},
            ],
        )
        fallback = SucceedingProvider(
            [
                {"type": "text", "content": "Fallback"},
            ]
        )

        provider = FallbackLLMProvider(
            [recovering, fallback],
            recovery_interval_s=0.001,
        )

        # First call: primary fails, fallback succeeds
        chunks = await collect_chunks(provider)
        assert chunks == [{"type": "text", "content": "Fallback"}]
        assert provider.get_availability() == [False, True]

        # Poll until recovery probe marks provider 0 available (up to 2 s)
        async def _wait_for_recovery() -> None:
            while provider.get_availability() != [True, True]:
                await asyncio.sleep(0.005)

        await asyncio.wait_for(_wait_for_recovery(), timeout=2.0)

        # Provider 0 should now be available
        assert provider.get_availability() == [True, True]

        provider.destroy()

    # 6. maxRetryPerProvider = 2 -> retries twice before fallback
    @pytest.mark.asyncio
    async def test_max_retry_per_provider(self) -> None:
        primary = CallCountingFailingProvider()
        fallback = SucceedingProvider(
            [
                {"type": "text", "content": "Fallback"},
            ]
        )

        provider = FallbackLLMProvider(
            [primary, fallback],
            max_retry_per_provider=2,
        )

        chunks = await collect_chunks(provider)

        assert primary.call_count == 2
        assert chunks == [{"type": "text", "content": "Fallback"}]
        provider.destroy()

    # 7. All providers unavailable -> retries all
    @pytest.mark.asyncio
    async def test_retry_all_when_all_unavailable(self) -> None:
        primary_calls = 0

        class RecoverOnRetry:
            async def stream(
                self, messages: list[dict], tools: list[dict] | None = None, **_kwargs
            ) -> AsyncIterator[dict]:
                nonlocal primary_calls
                primary_calls += 1
                if primary_calls <= 1:
                    raise RuntimeError("primary fails")
                yield {"type": "text", "content": "Primary recovered"}

        primary = RecoverOnRetry()
        fallback = FailingProvider()

        provider = FallbackLLMProvider([primary, fallback])
        chunks = await collect_chunks(provider)

        # Primary failed once (first pass), fallback failed once (first pass),
        # then retry-all: primary succeeds
        assert primary_calls == 2
        assert chunks == [{"type": "text", "content": "Primary recovered"}]
        provider.destroy()

    # 8. Provider returns empty -> succeeds (not a failure)
    @pytest.mark.asyncio
    async def test_empty_stream_is_success(self) -> None:
        empty = SucceedingProvider([])
        fallback = SucceedingProvider(
            [
                {"type": "text", "content": "Should not reach"},
            ]
        )

        provider = FallbackLLMProvider([empty, fallback])
        chunks = await collect_chunks(provider)

        assert chunks == []
        provider.destroy()

    # --- Additional edge cases ---

    def test_raises_on_zero_providers(self) -> None:
        with pytest.raises(ValueError, match="at least one provider"):
            FallbackLLMProvider([])

    @pytest.mark.asyncio
    async def test_passes_messages_and_tools(self) -> None:
        received_messages = None
        received_tools = None

        class SpyProvider:
            async def stream(
                self, messages: list[dict], tools: list[dict] | None = None, **_kwargs
            ) -> AsyncIterator[dict]:
                nonlocal received_messages, received_tools
                received_messages = messages
                received_tools = tools
                yield {"type": "text", "content": "ok"}

        messages = [{"role": "user", "content": "test"}]
        tools = [{"type": "function", "function": {"name": "foo"}}]

        provider = FallbackLLMProvider([SpyProvider()])
        await collect_chunks(provider, messages)

        assert received_messages == messages
        # tools are passed as-is to provider.stream; our helper does not pass tools
        # so we test by calling stream directly
        provider.destroy()

    @pytest.mark.asyncio
    async def test_destroy_cancels_recovery_tasks(self) -> None:
        primary = FailingProvider()
        fallback = FailingProvider()

        provider = FallbackLLMProvider(
            [primary, fallback],
            recovery_interval_s=0.001,
        )

        with pytest.raises(AllProvidersFailedError):
            await collect_chunks(provider)

        provider.destroy()

        # Yield control briefly to confirm no lingering tasks cause issues
        await asyncio.sleep(0.01)

    @pytest.mark.asyncio
    async def test_tool_call_chunks_are_yielded(self) -> None:
        primary = SucceedingProvider(
            [
                {
                    "type": "tool_call",
                    "index": 0,
                    "id": "tc_1",
                    "name": "greet",
                    "arguments": "{}",
                },
                {"type": "text", "content": "done"},
            ]
        )

        provider = FallbackLLMProvider([primary])
        chunks = await collect_chunks(provider)

        assert chunks == [
            {
                "type": "tool_call",
                "index": 0,
                "id": "tc_1",
                "name": "greet",
                "arguments": "{}",
            },
            {"type": "text", "content": "done"},
        ]
        provider.destroy()
