"""Tests for plan-based rate limiting."""

from __future__ import annotations

import pytest
from httpx import ASGITransport, AsyncClient

from llmhosts.config import LLMHostsConfig, SecurityConfig
from llmhosts.proxy.app import create_app
from llmhosts.proxy.dispatcher import BackendDispatcher


@pytest.fixture
def auth_app():
    config = LLMHostsConfig()
    config.security = SecurityConfig(require_auth=True, rate_limit_rpm=60)
    return create_app(config)


@pytest.fixture
async def auth_client(auth_app):
    dispatcher = BackendDispatcher()
    await dispatcher.startup()
    auth_app.state.dispatcher = dispatcher

    transport = ASGITransport(app=auth_app, raise_app_exceptions=False)
    async with AsyncClient(transport=transport, base_url="http://test") as c:
        yield c

    await dispatcher.shutdown()


@pytest.mark.asyncio
@pytest.mark.unit
async def test_free_plan_gets_60_rpm(auth_client: AsyncClient) -> None:
    # Create a free key
    create_resp = await auth_client.post("/api/keys", json={"label": "free-key", "plan": "free"})
    raw_key = create_resp.json()["raw_key"]
    resp = await auth_client.get(
        "/v1/models",
        headers={"Authorization": f"Bearer {raw_key}"},
    )
    assert resp.status_code == 200
    assert resp.headers.get("x-ratelimit-limit") == "30"


@pytest.mark.asyncio
@pytest.mark.unit
async def test_pro_plan_gets_300_rpm(auth_client: AsyncClient) -> None:
    create_resp = await auth_client.post("/api/keys", json={"label": "pro-key", "plan": "pro"})
    raw_key = create_resp.json()["raw_key"]
    resp = await auth_client.get(
        "/v1/models",
        headers={"Authorization": f"Bearer {raw_key}"},
    )
    assert resp.status_code == 200
    assert resp.headers.get("x-ratelimit-limit") == "300"


@pytest.mark.asyncio
@pytest.mark.unit
async def test_team_plan_gets_600_rpm(auth_client: AsyncClient) -> None:
    create_resp = await auth_client.post("/api/keys", json={"label": "team-key", "plan": "team"})
    raw_key = create_resp.json()["raw_key"]
    resp = await auth_client.get(
        "/v1/models",
        headers={"Authorization": f"Bearer {raw_key}"},
    )
    assert resp.status_code == 200
    assert resp.headers.get("x-ratelimit-limit") == "1000"


@pytest.mark.asyncio
@pytest.mark.unit
async def test_enterprise_plan_gets_unlimited_rpm(auth_client: AsyncClient) -> None:
    create_resp = await auth_client.post("/api/keys", json={"label": "ent-key", "plan": "enterprise"})
    raw_key = create_resp.json()["raw_key"]
    resp = await auth_client.get(
        "/v1/models",
        headers={"Authorization": f"Bearer {raw_key}"},
    )
    assert resp.status_code == 200
    # Enterprise gets high budget (5000 weighted units per window)
    assert int(resp.headers.get("x-ratelimit-limit", "0")) == 5000
