"""Tests for rate limiting middleware (TokenBucket + RateLimitMiddleware)."""

from __future__ import annotations

from unittest.mock import patch

import pytest
from fastapi import FastAPI
from httpx import ASGITransport, AsyncClient

from llmhosts.proxy.ratelimit import RateLimitMiddleware, TokenBucket

# ---------------------------------------------------------------------------
# Minimal test app
# ---------------------------------------------------------------------------

_app = FastAPI()
_app.add_middleware(RateLimitMiddleware, rpm=60, burst=2)


@_app.get("/health")
async def health() -> dict[str, str]:
    return {"status": "ok"}


@_app.get("/v1/test")
async def test_endpoint() -> dict[str, bool]:
    return {"ok": True}


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


@pytest.fixture()
async def client() -> AsyncClient:
    """Async HTTP client wired to the minimal test app."""
    transport = ASGITransport(app=_app, raise_app_exceptions=False)
    async with AsyncClient(transport=transport, base_url="http://test") as ac:
        # Reset middleware state between tests.
        # Walk the middleware stack to find our RateLimitMiddleware.
        inner = _app.middleware_stack
        while inner is not None:
            if isinstance(inner, RateLimitMiddleware):
                inner._buckets.clear()
                inner._request_count = 0
                break
            inner = getattr(inner, "app", None)
        yield ac


# ---------------------------------------------------------------------------
# TokenBucket unit tests
# ---------------------------------------------------------------------------


@pytest.mark.unit
def test_token_bucket_allows_under_limit() -> None:
    """Bucket with burst=5 should allow 3 consecutive consumes."""
    bucket = TokenBucket(rate=1.0, burst=5)
    results = [bucket.consume() for _ in range(3)]
    for allowed, retry_after in results:
        assert allowed is True
        assert retry_after == 0.0


@pytest.mark.unit
def test_token_bucket_blocks_over_limit() -> None:
    """Bucket with burst=2 should block the third consume."""
    bucket = TokenBucket(rate=1.0, burst=2)
    r1 = bucket.consume()
    r2 = bucket.consume()
    r3 = bucket.consume()

    assert r1[0] is True
    assert r2[0] is True
    assert r3[0] is False
    assert r3[1] > 0.0  # retry_after should be positive


@pytest.mark.unit
def test_token_bucket_refills() -> None:
    """After exhaustion, advancing time should refill tokens."""
    bucket = TokenBucket(rate=1.0, burst=2)

    # Exhaust the bucket.
    bucket.consume()
    bucket.consume()
    allowed, _ = bucket.consume()
    assert allowed is False

    # Advance monotonic clock by 2 seconds -- should refill 2 tokens.
    initial_time = bucket.last_refill
    with patch("llmhosts.proxy.ratelimit.time") as mock_time:
        mock_time.monotonic.return_value = initial_time + 2.0
        allowed, retry_after = bucket.consume()

    assert allowed is True
    assert retry_after == 0.0


# ---------------------------------------------------------------------------
# RateLimitMiddleware integration tests
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
@pytest.mark.unit
async def test_ratelimit_allows_normal_request(client: AsyncClient) -> None:
    """A single request under the limit should pass with 200."""
    resp = await client.get("/v1/test")
    assert resp.status_code == 200
    assert resp.json() == {"ok": True}


@pytest.mark.asyncio
@pytest.mark.unit
async def test_ratelimit_includes_headers(client: AsyncClient) -> None:
    """Successful responses should include X-RateLimit-* headers."""
    resp = await client.get("/v1/test")
    assert resp.status_code == 200
    assert "x-ratelimit-limit" in resp.headers
    assert "x-ratelimit-remaining" in resp.headers
    assert resp.headers["x-ratelimit-limit"] == "60"
    # After 1 consume from burst=2, remaining should be 1.
    assert resp.headers["x-ratelimit-remaining"] == "1"


@pytest.mark.asyncio
@pytest.mark.unit
async def test_ratelimit_blocks_over_limit(client: AsyncClient) -> None:
    """With burst=2, the third rapid request should get 429."""
    r1 = await client.get("/v1/test")
    r2 = await client.get("/v1/test")
    r3 = await client.get("/v1/test")

    assert r1.status_code == 200
    assert r2.status_code == 200
    assert r3.status_code == 429

    body = r3.json()
    assert body["error"]["type"] == "rate_limit_error"
    assert body["error"]["code"] == "rate_limit_exceeded"


@pytest.mark.asyncio
@pytest.mark.unit
async def test_ratelimit_retry_after(client: AsyncClient) -> None:
    """429 response should include a Retry-After header."""
    # Exhaust burst.
    await client.get("/v1/test")
    await client.get("/v1/test")
    resp = await client.get("/v1/test")

    assert resp.status_code == 429
    assert "retry-after" in resp.headers
    retry_val = int(resp.headers["retry-after"])
    assert retry_val >= 1


@pytest.mark.asyncio
@pytest.mark.unit
async def test_ratelimit_exempt_health(client: AsyncClient) -> None:
    """/health should never be rate limited, even when over limit."""
    # Exhaust the burst on a normal path first.
    await client.get("/v1/test")
    await client.get("/v1/test")
    r_blocked = await client.get("/v1/test")
    assert r_blocked.status_code == 429

    # /health should still work.
    resp = await client.get("/health")
    assert resp.status_code == 200
    assert resp.json() == {"status": "ok"}
