"""Tests for Quota Management.

测试配额管理功能。
"""

from __future__ import annotations

import pytest
from fastapi import HTTPException

from sagellm_gateway.config import QuotaConfig, SecurityConfig, set_security_config
from sagellm_gateway.quota import QuotaManager, QuotaTracker


@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 TestQuotaTracker:
    """配额追踪器测试."""

    def test_initial_usage(self):
        """测试初始使用量为零."""
        tracker = QuotaTracker()
        requests, tokens = tracker.get_usage("user1")
        assert requests == 0
        assert tokens == 0

    def test_increment_usage(self):
        """测试增加使用量."""
        tracker = QuotaTracker()

        # 增加使用量
        tracker.increment_usage("user1", requests=1, tokens=500)
        requests, tokens = tracker.get_usage("user1")
        assert requests == 1
        assert tokens == 500

        # 再次增加
        tracker.increment_usage("user1", requests=2, tokens=1000)
        requests, tokens = tracker.get_usage("user1")
        assert requests == 3
        assert tokens == 1500

    def test_multiple_users(self):
        """测试多用户隔离."""
        tracker = QuotaTracker()

        # user1 使用
        tracker.increment_usage("user1", requests=5, tokens=1000)

        # user2 使用
        tracker.increment_usage("user2", requests=3, tokens=500)

        # 验证隔离
        requests1, tokens1 = tracker.get_usage("user1")
        requests2, tokens2 = tracker.get_usage("user2")

        assert requests1 == 5 and tokens1 == 1000
        assert requests2 == 3 and tokens2 == 500

    def test_model_specific_usage(self):
        """测试模型级别使用量追踪."""
        tracker = QuotaTracker()

        # user1 使用 model1
        tracker.increment_usage("user1", requests=5, tokens=1000, model_name="model1")

        # user1 使用 model2
        tracker.increment_usage("user1", requests=2, tokens=300, model_name="model2")

        # 获取总体使用量
        total_requests, total_tokens = tracker.get_usage("user1")
        assert total_requests == 7
        assert total_tokens == 1300

        # 获取 model1 使用量
        model1_requests, model1_tokens = tracker.get_usage("user1", model_name="model1")
        assert model1_requests == 5
        assert model1_tokens == 1000

        # 获取 model2 使用量
        model2_requests, model2_tokens = tracker.get_usage("user1", model_name="model2")
        assert model2_requests == 2
        assert model2_tokens == 300


class TestQuotaManager:
    """配额管理器测试."""

    @pytest.mark.asyncio
    async def test_disabled_quota(self, mock_request):
        """测试禁用配额管理."""
        config = SecurityConfig(
            quota=QuotaConfig(enabled=False),
        )
        set_security_config(config)

        manager = QuotaManager()

        # 不应该抛出异常
        await manager.check_quota(mock_request, "user1", estimated_tokens=1000000)

    @pytest.mark.asyncio
    async def test_daily_request_quota(self, mock_request):
        """测试每日请求配额."""
        config = SecurityConfig(
            quota=QuotaConfig(
                enabled=True,
                daily_requests_per_user={
                    "user1": 10,  # user1 每日 10 个请求
                },
            ),
        )
        set_security_config(config)

        manager = QuotaManager()

        # 前 10 个请求应该成功
        for _ in range(10):
            await manager.check_quota(mock_request, "user1", estimated_tokens=100)
            await manager.record_usage("user1", requests=1, tokens=100)

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

        assert exc_info.value.status_code == 403
        assert "daily_request_quota_exceeded" in str(exc_info.value.detail)

    @pytest.mark.asyncio
    async def test_daily_token_quota(self, mock_request):
        """测试每日令牌配额."""
        config = SecurityConfig(
            quota=QuotaConfig(
                enabled=True,
                daily_tokens_per_user={
                    "user1": 1000,  # user1 每日 1000 个令牌
                },
            ),
        )
        set_security_config(config)

        manager = QuotaManager()

        # 消费 600 个令牌（成功）
        await manager.check_quota(mock_request, "user1", estimated_tokens=600)
        await manager.record_usage("user1", requests=1, tokens=600)

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

        assert exc_info.value.status_code == 403
        assert "daily_token_quota_exceeded" in str(exc_info.value.detail)

    @pytest.mark.asyncio
    async def test_model_specific_quota(self, mock_request):
        """测试模型级别配额."""
        config = SecurityConfig(
            quota=QuotaConfig(
                enabled=True,
                model_specific_quotas={
                    "gpt-4": {
                        "user1": 5,  # user1 对 gpt-4 每日 5 个请求
                    },
                },
            ),
        )
        set_security_config(config)

        manager = QuotaManager()

        # 前 5 个 gpt-4 请求应该成功
        for _ in range(5):
            await manager.check_quota(
                mock_request, "user1", estimated_tokens=100, model_name="gpt-4"
            )
            await manager.record_usage("user1", requests=1, tokens=100, model_name="gpt-4")

        # 第 6 个 gpt-4 请求应该失败
        with pytest.raises(HTTPException) as exc_info:
            await manager.check_quota(
                mock_request, "user1", estimated_tokens=100, model_name="gpt-4"
            )

        assert exc_info.value.status_code == 403
        assert "model_quota_exceeded" in str(exc_info.value.detail)

        # 其他模型应该仍然可以使用
        await manager.check_quota(mock_request, "user1", estimated_tokens=100, model_name="gpt-3.5")

    @pytest.mark.asyncio
    async def test_record_usage(self, mock_request):
        """测试记录使用量."""
        config = SecurityConfig(
            quota=QuotaConfig(enabled=True),
        )
        set_security_config(config)

        manager = QuotaManager()

        # 记录使用量
        await manager.record_usage("user1", requests=5, tokens=1000)

        # 验证追踪器中的数据
        requests, tokens = manager.tracker.get_usage("user1")
        assert requests == 5
        assert tokens == 1000

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

        manager = QuotaManager()

        # user1 消费所有配额
        for _ in range(5):
            await manager.check_quota(mock_request, "user1", estimated_tokens=100)
            await manager.record_usage("user1", requests=1, tokens=100)

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

        # user2 仍然可以正常使用
        await manager.check_quota(mock_request, "user2", estimated_tokens=100)
        await manager.record_usage("user2", requests=1, tokens=100)

    @pytest.mark.asyncio
    async def test_combined_quotas(self, mock_request):
        """测试组合配额（请求 + 令牌 + 模型）."""
        config = SecurityConfig(
            quota=QuotaConfig(
                enabled=True,
                daily_requests_per_user={"user1": 100},
                daily_tokens_per_user={"user1": 10000},
                model_specific_quotas={
                    "gpt-4": {"user1": 10},
                },
            ),
        )
        set_security_config(config)

        manager = QuotaManager()

        # 10 个 gpt-4 请求（不超过请求和令牌配额）
        for _ in range(10):
            await manager.check_quota(
                mock_request, "user1", estimated_tokens=500, model_name="gpt-4"
            )
            await manager.record_usage("user1", requests=1, tokens=500, model_name="gpt-4")

        # 模型配额用尽
        with pytest.raises(HTTPException) as exc_info:
            await manager.check_quota(
                mock_request, "user1", estimated_tokens=500, model_name="gpt-4"
            )

        assert "model_quota_exceeded" in str(exc_info.value.detail)
