"""Tests for ``mgf.fastapi._rate_limit`` — F-1 rate-limit middleware
(v0.6.0)."""

from __future__ import annotations

import time

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient

from mgf.fastapi._rate_limit import RateLimitMiddleware, TokenBucket

# ---------------------------------------------------------------------------
# TokenBucket — algorithmic correctness
# ---------------------------------------------------------------------------


class TestTokenBucketConstruction:
    def test_zero_capacity_rejected(self) -> None:
        with pytest.raises(ValueError, match="capacity must be > 0"):
            TokenBucket(capacity=0, refill_per_second=1.0)

    def test_negative_capacity_rejected(self) -> None:
        with pytest.raises(ValueError, match="capacity must be > 0"):
            TokenBucket(capacity=-5, refill_per_second=1.0)

    def test_zero_refill_rejected(self) -> None:
        with pytest.raises(ValueError, match="refill_per_second must be > 0"):
            TokenBucket(capacity=10, refill_per_second=0)


class TestTokenBucketConsumption:
    def test_starts_full_allowing_burst(self) -> None:
        bucket = TokenBucket(capacity=5, refill_per_second=1.0)
        for _ in range(5):
            allowed, _ = bucket.try_consume("k")
            assert allowed is True

    def test_denies_when_empty(self) -> None:
        bucket = TokenBucket(capacity=3, refill_per_second=1.0)
        for _ in range(3):
            bucket.try_consume("k")
        allowed, retry_after = bucket.try_consume("k")
        assert allowed is False
        # 1 token needs ~1.0s to refill (at rate 1/s).
        assert 0.9 <= retry_after <= 1.1

    def test_per_key_isolation(self) -> None:
        bucket = TokenBucket(capacity=2, refill_per_second=1.0)
        # Exhaust key "a"; key "b" should be unaffected.
        bucket.try_consume("a")
        bucket.try_consume("a")
        a_third, _ = bucket.try_consume("a")
        b_first, _ = bucket.try_consume("b")
        assert a_third is False
        assert b_first is True

    def test_refill_over_time(self) -> None:
        bucket = TokenBucket(capacity=2, refill_per_second=100.0)  # fast
        bucket.try_consume("k")
        bucket.try_consume("k")
        allowed_before, _ = bucket.try_consume("k")
        assert allowed_before is False
        time.sleep(0.02)  # 20ms → ~2 tokens at rate 100/s
        allowed_after, _ = bucket.try_consume("k")
        assert allowed_after is True

    def test_refill_caps_at_capacity(self) -> None:
        bucket = TokenBucket(capacity=3, refill_per_second=100.0)
        time.sleep(0.1)  # 10 tokens worth — should cap at 3
        # Consume capacity quickly; 4th should fail (cap = 3).
        for _ in range(3):
            allowed, _ = bucket.try_consume("k")
            assert allowed is True
        allowed_4th, _ = bucket.try_consume("k")
        assert allowed_4th is False


class TestTokenBucketReset:
    def test_reset_all(self) -> None:
        bucket = TokenBucket(capacity=2, refill_per_second=1.0)
        for _ in range(2):
            bucket.try_consume("k")
        bucket.reset()
        # Bucket starts fresh — key "k" has full capacity again.
        allowed, _ = bucket.try_consume("k")
        assert allowed is True

    def test_reset_specific_key(self) -> None:
        bucket = TokenBucket(capacity=2, refill_per_second=1.0)
        for _ in range(2):
            bucket.try_consume("a")
        for _ in range(2):
            bucket.try_consume("b")
        bucket.reset("a")
        a_fresh, _ = bucket.try_consume("a")
        b_still_empty, _ = bucket.try_consume("b")
        assert a_fresh is True
        assert b_still_empty is False


# ---------------------------------------------------------------------------
# RateLimitMiddleware — wired into a real FastAPI app
# ---------------------------------------------------------------------------


def _build_app(*, bucket, key) -> FastAPI:
    app = FastAPI()
    app.add_middleware(RateLimitMiddleware, bucket=bucket, key=key)

    @app.get("/ping")
    def ping():
        return {"status": "ok"}

    return app


class TestRateLimitMiddleware:
    def test_passes_through_under_limit(self) -> None:
        bucket = TokenBucket(capacity=3, refill_per_second=1.0)
        app = _build_app(bucket=bucket, key=lambda req: "fixed-key")
        client = TestClient(app)

        for _ in range(3):
            r = client.get("/ping")
            assert r.status_code == 200

    def test_429_when_exhausted(self) -> None:
        bucket = TokenBucket(capacity=2, refill_per_second=1.0)
        app = _build_app(bucket=bucket, key=lambda req: "fixed-key")
        client = TestClient(app)

        client.get("/ping")
        client.get("/ping")
        r = client.get("/ping")
        assert r.status_code == 429
        assert "Retry-After" in r.headers
        body = r.json()
        assert body["error"] == "rate-limit-exceeded"
        assert body["retry_after"] > 0

    def test_retry_after_header_is_integer_string(self) -> None:
        bucket = TokenBucket(capacity=1, refill_per_second=0.5)
        app = _build_app(bucket=bucket, key=lambda req: "k")
        client = TestClient(app)
        client.get("/ping")
        r = client.get("/ping")
        assert r.status_code == 429
        # Retry-After per RFC 7231: integer delta-seconds.
        assert r.headers["Retry-After"].isdigit()
        assert int(r.headers["Retry-After"]) >= 1

    def test_none_key_disables_limiting(self) -> None:
        # When `key(request)` returns None, requests pass freely.
        bucket = TokenBucket(capacity=1, refill_per_second=0.001)
        app = _build_app(bucket=bucket, key=lambda req: None)
        client = TestClient(app)
        # Many requests — none rate-limited.
        for _ in range(20):
            r = client.get("/ping")
            assert r.status_code == 200

    def test_per_key_isolation_via_middleware(self) -> None:
        # Use header to distinguish keys.
        bucket = TokenBucket(capacity=1, refill_per_second=0.001)
        app = _build_app(
            bucket=bucket,
            key=lambda req: req.headers.get("X-API-Key"),
        )
        client = TestClient(app)

        # Both clients get their first request.
        r1 = client.get("/ping", headers={"X-API-Key": "alice"})
        r2 = client.get("/ping", headers={"X-API-Key": "bob"})
        assert r1.status_code == 200
        assert r2.status_code == 200

        # Each gets denied on their second.
        r3 = client.get("/ping", headers={"X-API-Key": "alice"})
        r4 = client.get("/ping", headers={"X-API-Key": "bob"})
        assert r3.status_code == 429
        assert r4.status_code == 429
