"""Tests for LLM retry logic."""

from __future__ import annotations

import httpx
import pytest

from koan.llm.retry import _is_retryable, llm_retry

# -- Tests: _is_retryable classification -----------------------------------

class TestIsRetryable:
    def test_429_is_retryable(self) -> None:
        err = Exception("rate limit")
        err.status_code = 429  # type: ignore[attr-defined]
        assert _is_retryable(err) is True

    def test_500_is_retryable(self) -> None:
        err = Exception("server error")
        err.status_code = 500  # type: ignore[attr-defined]
        assert _is_retryable(err) is True

    def test_502_is_retryable(self) -> None:
        err = Exception("bad gateway")
        err.status_code = 502  # type: ignore[attr-defined]
        assert _is_retryable(err) is True

    def test_503_is_retryable(self) -> None:
        err = Exception("service unavailable")
        err.status_code = 503  # type: ignore[attr-defined]
        assert _is_retryable(err) is True

    def test_529_is_retryable(self) -> None:
        err = Exception("overloaded")
        err.status_code = 529  # type: ignore[attr-defined]
        assert _is_retryable(err) is True

    def test_400_not_retryable(self) -> None:
        err = Exception("bad request")
        err.status_code = 400  # type: ignore[attr-defined]
        assert _is_retryable(err) is False

    def test_401_not_retryable(self) -> None:
        err = Exception("unauthorized")
        err.status_code = 401  # type: ignore[attr-defined]
        assert _is_retryable(err) is False

    def test_403_not_retryable(self) -> None:
        err = Exception("forbidden")
        err.status_code = 403  # type: ignore[attr-defined]
        assert _is_retryable(err) is False

    def test_404_not_retryable(self) -> None:
        err = Exception("not found")
        err.status_code = 404  # type: ignore[attr-defined]
        assert _is_retryable(err) is False

    def test_connect_error_is_retryable(self) -> None:
        assert _is_retryable(httpx.ConnectError("refused")) is True

    def test_read_timeout_is_retryable(self) -> None:
        assert _is_retryable(httpx.ReadTimeout("timed out")) is True

    def test_pool_timeout_is_retryable(self) -> None:
        assert _is_retryable(httpx.PoolTimeout("pool exhausted")) is True

    def test_generic_exception_not_retryable(self) -> None:
        assert _is_retryable(ValueError("something else")) is False

    def test_status_via_status_attr(self) -> None:
        """Some SDKs use .status instead of .status_code."""
        err = Exception("rate limit")
        err.status = 429  # type: ignore[attr-defined]
        assert _is_retryable(err) is True


# -- Tests: llm_retry decorator behavior -----------------------------------

class TestLLMRetryDecorator:
    async def test_retries_then_succeeds(self) -> None:
        call_count = 0

        @llm_retry
        async def flaky_call() -> str:
            nonlocal call_count
            call_count += 1
            if call_count < 3:
                err = Exception("rate limit")
                err.status_code = 429  # type: ignore[attr-defined]
                raise err
            return "ok"

        result = await flaky_call()
        assert result == "ok"
        assert call_count == 3

    async def test_non_retryable_raises_immediately(self) -> None:
        call_count = 0

        @llm_retry
        async def bad_request() -> str:
            nonlocal call_count
            call_count += 1
            err = Exception("bad request")
            err.status_code = 400  # type: ignore[attr-defined]
            raise err

        with pytest.raises(Exception, match="bad request"):
            await bad_request()
        assert call_count == 1

    async def test_max_retries_then_raises(self) -> None:
        call_count = 0

        @llm_retry
        async def always_fails() -> str:
            nonlocal call_count
            call_count += 1
            err = Exception("overloaded")
            err.status_code = 429  # type: ignore[attr-defined]
            raise err

        with pytest.raises(Exception, match="overloaded"):
            await always_fails()
        assert call_count == 4  # 1 initial + 3 retries
