from __future__ import annotations

from collections.abc import Awaitable, Callable

import httpx
import pytest

from agent_platform.config import Settings
from agent_platform.runtime.provider import ModelProvider


@pytest.fixture(autouse=True)
def disable_retry_sleep(monkeypatch: pytest.MonkeyPatch) -> None:
    async def no_sleep(_: float) -> None:
        return None

    monkeypatch.setattr("agent_platform.runtime.provider.asyncio.sleep", no_sleep)


async def _complete_with_handler(
    settings: Settings,
    handler: Callable[[httpx.Request], httpx.Response | Awaitable[httpx.Response]],
) -> dict:
    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        provider = ModelProvider(settings, client)
        return await provider.complete(model="test-model", messages=[{"role": "user", "content": "hello"}])


async def test_complete_retries_transient_gateway_response(settings: Settings) -> None:
    calls = 0

    def handler(request: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        if calls == 1:
            return httpx.Response(524, request=request, headers={"Retry-After": "0"})
        return httpx.Response(
            200,
            request=request,
            json={"choices": [{"message": {"role": "assistant", "content": "ok"}}]},
        )

    result = await _complete_with_handler(settings, handler)

    assert calls == 2
    assert result["choices"][0]["message"]["content"] == "ok"


async def test_complete_retries_transport_error(settings: Settings) -> None:
    calls = 0

    def handler(request: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        if calls == 1:
            raise httpx.ConnectError("temporary disconnect", request=request)
        return httpx.Response(
            200,
            request=request,
            json={"choices": [{"message": {"role": "assistant", "content": "recovered"}}]},
        )

    result = await _complete_with_handler(settings, handler)

    assert calls == 2
    assert result["choices"][0]["message"]["content"] == "recovered"


async def test_complete_does_not_retry_client_error(settings: Settings) -> None:
    calls = 0

    def handler(request: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        return httpx.Response(400, request=request, json={"error": "bad request"})

    with pytest.raises(httpx.HTTPStatusError):
        await _complete_with_handler(settings, handler)

    assert calls == 1