"""Tests for adaptive sliding window rate limiter with per-model weights.

Covers three layers:
  1. InMemoryBackend -- pure sliding window logic
  2. Model weight resolution
  3. AdaptiveRateLimitMiddleware -- integration via ASGI transport
"""

from __future__ import annotations

import time
from unittest.mock import patch

import pytest
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from httpx import ASGITransport, AsyncClient

from llmhosts.security.adaptive_rate_limit import (
    AdaptiveRateLimitMiddleware,
    InMemoryBackend,
    RateLimitConfig,
    RateLimitResult,
)

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


def _make_test_app(config: RateLimitConfig | None = None) -> FastAPI:
    """Build a minimal FastAPI app with the adaptive rate limiter."""
    app = FastAPI()
    app.add_middleware(AdaptiveRateLimitMiddleware, config=config)

    @app.post("/v1/chat/completions")
    async def chat(request: Request) -> JSONResponse:
        return JSONResponse({"choices": []})

    @app.get("/v1/models")
    async def models() -> JSONResponse:
        return JSONResponse({"data": []})

    @app.get("/health")
    async def health() -> JSONResponse:
        return JSONResponse({"status": "ok"})

    @app.get("/metrics")
    async def metrics() -> JSONResponse:
        return JSONResponse({"metrics": []})

    return app


async def _client_for(app: FastAPI) -> AsyncClient:
    """Create an httpx AsyncClient wired to *app* via ASGI transport."""
    transport = ASGITransport(app=app, raise_app_exceptions=False)
    return AsyncClient(transport=transport, base_url="http://test")


# ---------------------------------------------------------------------------
# InMemoryBackend unit tests
# ---------------------------------------------------------------------------


class TestInMemoryBackend:
    """Pure logic tests for the sliding window backend."""

    def test_first_request_allowed(self) -> None:
        backend = InMemoryBackend()
        result = backend.consume("user:1", weight=1, limit=10, window=60)
        assert result.allowed is True
        assert result.remaining == 9
        assert result.weight == 1

    def test_requests_within_limit(self) -> None:
        backend = InMemoryBackend()
        for _ in range(5):
            result = backend.consume("user:1", weight=1, limit=10, window=60)
        assert result.allowed is True
        assert result.remaining == 5

    def test_requests_exceed_limit(self) -> None:
        backend = InMemoryBackend()
        for _ in range(10):
            backend.consume("user:1", weight=1, limit=10, window=60)
        result = backend.consume("user:1", weight=1, limit=10, window=60)
        assert result.allowed is False
        assert result.remaining == 0

    def test_window_expiry_frees_capacity(self) -> None:
        backend = InMemoryBackend()
        # Fill the bucket.
        for _ in range(10):
            backend.consume("user:1", weight=1, limit=10, window=60)

        # Verify it is full.
        result = backend.consume("user:1", weight=1, limit=10, window=60)
        assert result.allowed is False

        # Advance the clock past the window so all requests expire.
        with patch("llmhosts.security.adaptive_rate_limit.time") as mock_time:
            mock_time.monotonic.return_value = time.monotonic() + 61
            result = backend.consume("user:1", weight=1, limit=10, window=60)
        assert result.allowed is True
        assert result.remaining == 9

    def test_weighted_consumption(self) -> None:
        backend = InMemoryBackend()
        # Weight=4 should consume 4 of 10 units.
        result = backend.consume("user:1", weight=4, limit=10, window=60)
        assert result.allowed is True
        assert result.remaining == 6

        # Another weight=4 -> 8 total, 2 remaining.
        result = backend.consume("user:1", weight=4, limit=10, window=60)
        assert result.allowed is True
        assert result.remaining == 2

        # weight=3 would exceed (8+3=11 > 10).
        result = backend.consume("user:1", weight=3, limit=10, window=60)
        assert result.allowed is False

    def test_remaining_decreases(self) -> None:
        backend = InMemoryBackend()
        r1 = backend.consume("user:1", weight=1, limit=5, window=60)
        r2 = backend.consume("user:1", weight=1, limit=5, window=60)
        r3 = backend.consume("user:1", weight=1, limit=5, window=60)
        assert r1.remaining == 4
        assert r2.remaining == 3
        assert r3.remaining == 2

    def test_retry_after_calculated(self) -> None:
        backend = InMemoryBackend()
        # Fill completely.
        for _ in range(10):
            backend.consume("user:1", weight=1, limit=10, window=60)
        result = backend.consume("user:1", weight=1, limit=10, window=60)
        assert result.allowed is False
        assert result.retry_after is not None
        assert result.retry_after > 0.0
        # retry_after should be at most the full window.
        assert result.retry_after <= 60.0

    def test_get_remaining_empty(self) -> None:
        backend = InMemoryBackend()
        remaining = backend.get_remaining("nonexistent", limit=100, window=60)
        assert remaining == 100

    def test_get_remaining_after_requests(self) -> None:
        backend = InMemoryBackend()
        backend.consume("user:1", weight=3, limit=10, window=60)
        backend.consume("user:1", weight=2, limit=10, window=60)
        remaining = backend.get_remaining("user:1", limit=10, window=60)
        assert remaining == 5

    def test_cleanup_idle_entries(self) -> None:
        backend = InMemoryBackend()
        backend.consume("user:1", weight=1, limit=10, window=60)
        backend.consume("user:2", weight=1, limit=10, window=60)
        assert len(backend._buckets) == 2

        # Advance clock past the idle TTL (3600s).
        with patch("llmhosts.security.adaptive_rate_limit.time") as mock_time:
            mock_time.monotonic.return_value = time.monotonic() + 3601
            evicted = backend.cleanup()
        assert evicted == 2
        assert len(backend._buckets) == 0

    def test_concurrent_keys_independent(self) -> None:
        backend = InMemoryBackend()
        # Fill user:1 completely.
        for _ in range(5):
            backend.consume("user:1", weight=1, limit=5, window=60)
        blocked = backend.consume("user:1", weight=1, limit=5, window=60)
        assert blocked.allowed is False

        # user:2 should still have full budget.
        result = backend.consume("user:2", weight=1, limit=5, window=60)
        assert result.allowed is True
        assert result.remaining == 4


# ---------------------------------------------------------------------------
# Model weight tests
# ---------------------------------------------------------------------------


class TestModelWeights:
    """Verify per-model weight lookup."""

    def test_gpt4_costs_4_units(self) -> None:
        config = RateLimitConfig()
        assert config.model_weights["gpt-4"] == 4

    def test_gpt35_costs_1_unit(self) -> None:
        config = RateLimitConfig()
        assert config.model_weights["gpt-3.5-turbo"] == 1

    def test_unknown_model_default_weight(self) -> None:
        config = RateLimitConfig()
        assert config.model_weights.get("some-unknown-model", config.default_weight) == 1

    def test_custom_weights(self) -> None:
        config = RateLimitConfig(model_weights={"my-model": 7}, default_weight=2)
        assert config.model_weights["my-model"] == 7
        assert config.default_weight == 2

    def test_claude_opus_costs_4_units(self) -> None:
        config = RateLimitConfig()
        assert config.model_weights["claude-3-opus"] == 4

    def test_claude_haiku_costs_1_unit(self) -> None:
        config = RateLimitConfig()
        assert config.model_weights["claude-3-haiku"] == 1

    def test_llama_70b_costs_2_units(self) -> None:
        config = RateLimitConfig()
        assert config.model_weights["llama-3-70b"] == 2


# ---------------------------------------------------------------------------
# Middleware integration tests (ASGI transport, no real server)
# ---------------------------------------------------------------------------


class TestMiddlewareIntegration:
    """Test the middleware in a FastAPI context using ASGI transport."""

    @pytest.mark.asyncio
    async def test_health_endpoint_exempt(self) -> None:
        app = _make_test_app(RateLimitConfig(enabled=True))
        async with await _client_for(app) as client:
            resp = await client.get("/health")
        assert resp.status_code == 200
        assert resp.json() == {"status": "ok"}
        # Health responses should NOT have rate-limit headers.
        assert "x-ratelimit-limit" not in resp.headers

    @pytest.mark.asyncio
    async def test_metrics_endpoint_exempt(self) -> None:
        app = _make_test_app(RateLimitConfig(enabled=True))
        async with await _client_for(app) as client:
            resp = await client.get("/metrics")
        assert resp.status_code == 200
        assert "x-ratelimit-limit" not in resp.headers

    @pytest.mark.asyncio
    async def test_rate_limit_headers_on_response(self) -> None:
        app = _make_test_app(RateLimitConfig(enabled=True))
        async with await _client_for(app) as client:
            resp = await client.get("/v1/models")
        assert resp.status_code == 200
        assert "x-ratelimit-limit" in resp.headers
        assert "x-ratelimit-remaining" in resp.headers
        assert "x-ratelimit-reset" in resp.headers
        assert "x-ratelimit-model-weight" in resp.headers
        # Default weight for GET requests is 1.
        assert resp.headers["x-ratelimit-model-weight"] == "1"

    @pytest.mark.asyncio
    async def test_429_when_limit_exceeded(self) -> None:
        # Tiny limit: 3 units per window.
        config = RateLimitConfig(enabled=True, window_seconds=60)
        app = _make_test_app(config)
        async with await _client_for(app) as client:
            # With free plan (30 units) and weight=1, send 31 requests.
            responses = []
            for _ in range(31):
                responses.append(await client.get("/v1/models"))
            # First 30 should succeed, 31st should be 429.
            assert all(r.status_code == 200 for r in responses[:30])
            assert responses[30].status_code == 429
            body = responses[30].json()
            assert body["error"] == "rate_limit_exceeded"
            assert body["weight"] == 1

    @pytest.mark.asyncio
    async def test_429_includes_retry_after(self) -> None:
        config = RateLimitConfig(enabled=True, window_seconds=60)
        app = _make_test_app(config)
        async with await _client_for(app) as client:
            # Exhaust free plan limit (30 units).
            for _ in range(30):
                await client.get("/v1/models")
            resp = await client.get("/v1/models")
        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
    async def test_model_weight_from_request_body(self) -> None:
        # Limit of 5 units.  GPT-4 costs 4.  Two requests = 8 > 5.
        config = RateLimitConfig(enabled=True, window_seconds=60)
        app = _make_test_app(config)
        async with await _client_for(app) as client:
            # First GPT-4 request (weight 4) should succeed against 30-unit free budget.
            r1 = await client.post(
                "/v1/chat/completions",
                json={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
            )
            assert r1.status_code == 200
            assert r1.headers["x-ratelimit-model-weight"] == "4"
            # Remaining should be 30 - 4 = 26.
            assert r1.headers["x-ratelimit-remaining"] == "26"

    @pytest.mark.asyncio
    async def test_heavy_model_exhausts_budget_faster(self) -> None:
        # Use a small custom limit to make the test fast.
        config = RateLimitConfig(
            enabled=True,
            window_seconds=60,
            model_weights={"heavy-model": 5},
            default_weight=1,
        )
        app = _make_test_app(config)

        # Monkey-patch plan limit to 10 for simpler test.
        with patch(
            "llmhosts.security.adaptive_rate_limit._PLAN_LIMITS",
            {"free": 10, "pro": 300, "team": 1000, "enterprise": 5000},
        ):
            async with await _client_for(app) as client:
                # weight=5: first request OK (5/10 used).
                r1 = await client.post(
                    "/v1/chat/completions",
                    json={"model": "heavy-model", "messages": []},
                )
                assert r1.status_code == 200
                assert r1.headers["x-ratelimit-remaining"] == "5"

                # Second heavy request (5+5=10): still OK.
                r2 = await client.post(
                    "/v1/chat/completions",
                    json={"model": "heavy-model", "messages": []},
                )
                assert r2.status_code == 200
                assert r2.headers["x-ratelimit-remaining"] == "0"

                # Third heavy request (10+5=15 > 10): blocked.
                r3 = await client.post(
                    "/v1/chat/completions",
                    json={"model": "heavy-model", "messages": []},
                )
                assert r3.status_code == 429

    @pytest.mark.asyncio
    async def test_disabled_config_passthrough(self) -> None:
        config = RateLimitConfig(enabled=False)
        app = _make_test_app(config)
        async with await _client_for(app) as client:
            resp = await client.get("/v1/models")
        assert resp.status_code == 200
        # When disabled, no rate-limit headers should be added.
        assert "x-ratelimit-limit" not in resp.headers

    @pytest.mark.asyncio
    async def test_malformed_body_uses_default_weight(self) -> None:
        config = RateLimitConfig(enabled=True)
        app = _make_test_app(config)
        async with await _client_for(app) as client:
            resp = await client.post(
                "/v1/chat/completions",
                content=b"not json at all",
                headers={"content-type": "application/json"},
            )
        # FastAPI will 422 the bad body, but the middleware should not crash.
        # The middleware runs before the route validation, so it should still
        # add headers (it processes and then forwards to the handler).
        # We just care that it does not return 500.
        assert resp.status_code != 500


# ---------------------------------------------------------------------------
# Config tests
# ---------------------------------------------------------------------------


class TestRateLimitConfig:
    """Pydantic config model validation."""

    def test_default_config(self) -> None:
        config = RateLimitConfig()
        assert config.enabled is True
        assert config.algorithm == "sliding_window"
        assert config.window_seconds == 60
        assert config.default_weight == 1
        assert len(config.model_weights) > 0

    def test_custom_model_weights(self) -> None:
        config = RateLimitConfig(model_weights={"my-model": 10})
        assert config.model_weights == {"my-model": 10}
        # Default weight unaffected.
        assert config.default_weight == 1

    def test_disabled_config(self) -> None:
        config = RateLimitConfig(enabled=False)
        assert config.enabled is False

    def test_custom_window(self) -> None:
        config = RateLimitConfig(window_seconds=120)
        assert config.window_seconds == 120

    def test_cleanup_interval_default(self) -> None:
        config = RateLimitConfig()
        assert config.cleanup_interval_seconds == 3600.0


# ---------------------------------------------------------------------------
# RateLimitResult dataclass tests
# ---------------------------------------------------------------------------


class TestRateLimitResult:
    """Basic dataclass sanity checks."""

    def test_allowed_result(self) -> None:
        result = RateLimitResult(allowed=True, limit=100, remaining=99, reset_at=1000.0, weight=1)
        assert result.allowed is True
        assert result.retry_after is None

    def test_blocked_result(self) -> None:
        result = RateLimitResult(allowed=False, limit=100, remaining=0, reset_at=1000.0, weight=4, retry_after=30.0)
        assert result.allowed is False
        assert result.retry_after == 30.0
        assert result.weight == 4
