"""Tests for Rate Limiting.

测试速率限制功能（令牌桶算法）。
"""

from __future__ import annotations

import time

import pytest
from fastapi import HTTPException

from sagellm_gateway.config import RateLimitConfig, SecurityConfig, set_security_config
from sagellm_gateway.rate_limit import RateLimiter, TokenBucket


@pytest.fixture
def mock_request():
    """创建模拟的 FastAPI Request."""

    class MockRequest:
        def __init__(self):
            self.url = type("URL", (), {"path": "/v1/chat/completions"})()
            self.method = "POST"

    return MockRequest()


class TestTokenBucket:
    """令牌桶算法测试."""

    def test_initial_capacity(self):
        """测试初始容量."""
        bucket = TokenBucket(capacity=100, refill_rate=10.0)
        assert bucket.get_available_tokens() == 100

    def test_consume_tokens(self):
        """测试消费令牌."""
        bucket = TokenBucket(capacity=100, refill_rate=10.0)

        # 消费 10 个令牌
        assert bucket.consume(10) is True
        assert 89.9 <= bucket.get_available_tokens() <= 90.1  # 允许浮点误差

        # 再消费 50 个令牌
        assert bucket.consume(50) is True
        assert 39.9 <= bucket.get_available_tokens() <= 40.1  # 允许浮点误差

    def test_consume_exceed_capacity(self):
        """测试消费超过容量."""
        bucket = TokenBucket(capacity=100, refill_rate=10.0)

        # 尝试消费 150 个令牌（超过容量）
        assert bucket.consume(150) is False
        assert bucket.get_available_tokens() == 100  # 未消费

    def test_refill_tokens(self):
        """测试令牌补充."""
        bucket = TokenBucket(capacity=100, refill_rate=10.0)  # 每秒 10 个令牌

        # 消费 50 个令牌
        bucket.consume(50)
        assert 49.9 <= bucket.get_available_tokens() <= 50.1  # 允许浮点误差

        # 等待 1 秒，应该补充 10 个令牌
        time.sleep(1.1)
        tokens = bucket.get_available_tokens()
        assert 59 <= tokens <= 62  # 允许小误差（考虑时间精度）

    def test_refill_cap_at_capacity(self):
        """测试令牌补充不超过容量."""
        bucket = TokenBucket(capacity=100, refill_rate=10.0)

        # 等待 20 秒，理论上补充 200 个令牌，但最多 100
        time.sleep(2.0)
        assert bucket.get_available_tokens() <= 100


class TestRateLimiter:
    """速率限制器测试."""

    def test_disabled_rate_limiter(self, mock_request):
        """测试禁用速率限制."""
        config = SecurityConfig(
            rate_limit=RateLimitConfig(enabled=False),
        )
        set_security_config(config)

        limiter = RateLimiter()
        # 不应该抛出异常
        import asyncio

        asyncio.run(limiter.check_rate_limit(mock_request, "user1", estimated_tokens=1000))

    @pytest.mark.asyncio
    async def test_global_rpm_limit(self, mock_request):
        """测试全局 RPM 限制."""
        config = SecurityConfig(
            rate_limit=RateLimitConfig(
                enabled=True,
                rpm=10,  # 每分钟 10 个请求
            ),
        )
        set_security_config(config)

        limiter = RateLimiter()

        # 前 10 个请求应该成功
        for _ in range(10):
            await limiter.check_rate_limit(mock_request, "user1", estimated_tokens=100)

        # 第 11 个请求应该失败
        with pytest.raises(HTTPException) as exc_info:
            await limiter.check_rate_limit(mock_request, "user1", estimated_tokens=100)

        assert exc_info.value.status_code == 429
        assert "rate_limit_exceeded" in str(exc_info.value.detail)

    @pytest.mark.asyncio
    async def test_global_tpm_limit(self, mock_request):
        """测试全局 TPM 限制."""
        config = SecurityConfig(
            rate_limit=RateLimitConfig(
                enabled=True,
                tpm=1000,  # 每分钟 1000 个令牌
            ),
        )
        set_security_config(config)

        limiter = RateLimiter()

        # 消费 500 个令牌（成功）
        await limiter.check_rate_limit(mock_request, "user1", estimated_tokens=500)

        # 消费 400 个令牌（成功）
        await limiter.check_rate_limit(mock_request, "user2", estimated_tokens=400)

        # 消费 200 个令牌（失败，总计 1100 > 1000）
        with pytest.raises(HTTPException) as exc_info:
            await limiter.check_rate_limit(mock_request, "user3", estimated_tokens=200)

        assert exc_info.value.status_code == 429
        assert "token_rate_limit_exceeded" in str(exc_info.value.detail)

    @pytest.mark.asyncio
    async def test_user_rpm_limit(self, mock_request):
        """测试用户级别 RPM 限制."""
        config = SecurityConfig(
            rate_limit=RateLimitConfig(
                enabled=True,
                rpm_per_user={
                    "user_alice": 5,  # Alice 每分钟 5 个请求
                    "user_bob": 10,  # Bob 每分钟 10 个请求
                },
            ),
        )
        set_security_config(config)

        limiter = RateLimiter()

        # Alice 的前 5 个请求应该成功
        for _ in range(5):
            await limiter.check_rate_limit(mock_request, "user_alice", estimated_tokens=100)

        # Alice 的第 6 个请求应该失败
        with pytest.raises(HTTPException) as exc_info:
            await limiter.check_rate_limit(mock_request, "user_alice", estimated_tokens=100)

        assert exc_info.value.status_code == 429
        assert "user_rate_limit_exceeded" in str(exc_info.value.detail)

        # Bob 应该仍然可以发送请求
        await limiter.check_rate_limit(mock_request, "user_bob", estimated_tokens=100)

    @pytest.mark.asyncio
    async def test_user_tpm_limit(self, mock_request):
        """测试用户级别 TPM 限制."""
        config = SecurityConfig(
            rate_limit=RateLimitConfig(
                enabled=True,
                tpm_per_user={
                    "user_alice": 500,  # Alice 每分钟 500 个令牌
                },
            ),
        )
        set_security_config(config)

        limiter = RateLimiter()

        # 消费 300 个令牌（成功）
        await limiter.check_rate_limit(mock_request, "user_alice", estimated_tokens=300)

        # 消费 250 个令牌（失败，总计 550 > 500）
        with pytest.raises(HTTPException) as exc_info:
            await limiter.check_rate_limit(mock_request, "user_alice", estimated_tokens=250)

        assert exc_info.value.status_code == 429

    @pytest.mark.asyncio
    async def test_multiple_users_isolated(self, mock_request):
        """测试多用户隔离."""
        config = SecurityConfig(
            rate_limit=RateLimitConfig(
                enabled=True,
                rpm_per_user={
                    "user1": 5,
                    "user2": 5,
                },
            ),
        )
        set_security_config(config)

        limiter = RateLimiter()

        # user1 消费所有配额
        for _ in range(5):
            await limiter.check_rate_limit(mock_request, "user1", estimated_tokens=100)

        # user1 超限
        with pytest.raises(HTTPException):
            await limiter.check_rate_limit(mock_request, "user1", estimated_tokens=100)

        # user2 仍然可以正常使用
        await limiter.check_rate_limit(mock_request, "user2", estimated_tokens=100)
