"""Tests for quota manager: token/request quotas, TTL calculation, reset schedules."""

import time

import pytest
from freezegun import freeze_time

from llm_rotator._types import Candidate, Usage
from llm_rotator.backends import InMemoryBackend
from llm_rotator.config import (
    ClientType,
    KeyConfig,
    ModelGroupConfig,
    ProviderConfig,
    RequestQuotaConfig,
    TokenQuotaConfig,
)
from llm_rotator.exceptions import QuotaExceededError
from llm_rotator.quota import QuotaManager


def _make_candidate(
    key_alias: str = "main",
    model: str = "gpt-4o",
    provider_name: str = "openai",
    model_group: str | None = "flagship",
) -> Candidate:
    return Candidate(
        provider_name=provider_name,
        model=model,
        key_alias=key_alias,
        key_token="sk-test",
        model_group=model_group,
        base_url="https://api.openai.com/v1",
        client_type=ClientType.OPENAI,
    )


def _make_usage(total: int = 100) -> Usage:
    return Usage(prompt_tokens=total // 2, completion_tokens=total - total // 2, total_tokens=total)


def _make_provider_with_group_quota(
    limit: int = 1000, reset: str = "daily_utc"
) -> ProviderConfig:
    return ProviderConfig(
        name="openai",
        client_type="openai",
        priority=1,
        model_groups=[
            ModelGroupConfig(
                name="flagship",
                tier=1,
                models=["gpt-4o", "gpt-5"],
                token_quota=TokenQuotaConfig(limit=limit, reset=reset),
            ),
        ],
        keys=[KeyConfig(token="sk-test", alias="main")],
    )


def _make_provider_with_key_quota(
    limit: int = 40, reset: str = "daily_utc"
) -> ProviderConfig:
    return ProviderConfig(
        name="openrouter",
        client_type="openai",
        priority=2,
        models=["llama-3"],
        keys=[
            KeyConfig(
                token="sk-or-test",
                alias="main",
                request_quota=RequestQuotaConfig(limit=limit, reset=reset),
            ),
        ],
    )


class TestQuotaManagerPreCheck:
    @freeze_time("2026-02-21 12:00:00", tz_offset=0)
    async def test_no_quota_always_allows(self):
        backend = InMemoryBackend()
        provider = ProviderConfig(
            name="openai",
            client_type="openai",
            priority=1,
            models=["gpt-4o"],
            keys=[KeyConfig(token="sk-test", alias="main")],
        )
        qm = QuotaManager(backend, [provider])
        candidate = _make_candidate(model_group=None)
        assert await qm.pre_check(candidate) is True

    @freeze_time("2026-02-21 12:00:00", tz_offset=0)
    async def test_under_limit_allows(self):
        backend = InMemoryBackend()
        provider = _make_provider_with_group_quota(limit=1000)
        qm = QuotaManager(backend, [provider])
        candidate = _make_candidate()
        assert await qm.pre_check(candidate) is True

    @freeze_time("2026-02-21 12:00:00", tz_offset=0)
    async def test_at_limit_blocks(self):
        backend = InMemoryBackend()
        provider = _make_provider_with_group_quota(limit=100)
        qm = QuotaManager(backend, [provider])
        candidate = _make_candidate()

        # Record usage to fill quota
        await qm.record_usage(candidate, _make_usage(total=100))

        assert await qm.pre_check(candidate) is False

    @freeze_time("2026-02-21 12:00:00", tz_offset=0)
    async def test_over_limit_blocks(self):
        backend = InMemoryBackend()
        provider = _make_provider_with_group_quota(limit=100)
        qm = QuotaManager(backend, [provider])
        candidate = _make_candidate()

        await qm.record_usage(candidate, _make_usage(total=150))

        assert await qm.pre_check(candidate) is False


class TestQuotaManagerRecordUsage:
    @freeze_time("2026-02-21 12:00:00", tz_offset=0)
    async def test_token_increment(self):
        backend = InMemoryBackend()
        provider = _make_provider_with_group_quota(limit=1000)
        qm = QuotaManager(backend, [provider])
        candidate = _make_candidate()

        await qm.record_usage(candidate, _make_usage(total=100))
        await qm.record_usage(candidate, _make_usage(total=50))

        # Check that quota accumulated
        assert await qm.pre_check(candidate) is True
        # 150 used of 1000

    @freeze_time("2026-02-21 12:00:00", tz_offset=0)
    async def test_request_quota_increment(self):
        backend = InMemoryBackend()
        provider = _make_provider_with_key_quota(limit=2)
        qm = QuotaManager(backend, [provider])
        candidate = _make_candidate(
            key_alias="main", provider_name="openrouter", model="llama-3", model_group=None
        )

        await qm.record_usage(candidate, _make_usage(total=100))
        assert await qm.pre_check(candidate) is True

        await qm.record_usage(candidate, _make_usage(total=100))
        assert await qm.pre_check(candidate) is False  # 2/2 requests


class TestQuotaManagerGroupScope:
    @freeze_time("2026-02-21 12:00:00", tz_offset=0)
    async def test_group_scope_shared_across_models(self):
        """Two models in same group share the token quota."""
        backend = InMemoryBackend()
        provider = _make_provider_with_group_quota(limit=200)
        qm = QuotaManager(backend, [provider])

        c1 = _make_candidate(model="gpt-4o")
        c2 = _make_candidate(model="gpt-5")

        await qm.record_usage(c1, _make_usage(total=120))
        await qm.record_usage(c2, _make_usage(total=90))

        # Both should be blocked — 210 > 200
        assert await qm.pre_check(c1) is False
        assert await qm.pre_check(c2) is False


class TestQuotaManagerKeyScope:
    @freeze_time("2026-02-21 12:00:00", tz_offset=0)
    async def test_key_scope_independent(self):
        """Different keys have independent request quotas."""
        backend = InMemoryBackend()
        provider = ProviderConfig(
            name="openrouter",
            client_type="openai",
            priority=1,
            models=["llama-3"],
            keys=[
                KeyConfig(
                    token="sk-a",
                    alias="key_a",
                    request_quota=RequestQuotaConfig(limit=2, reset="daily_utc"),
                ),
                KeyConfig(
                    token="sk-b",
                    alias="key_b",
                    request_quota=RequestQuotaConfig(limit=2, reset="daily_utc"),
                ),
            ],
        )
        qm = QuotaManager(backend, [provider])
        ca = _make_candidate(
            key_alias="key_a", provider_name="openrouter", model="llama-3", model_group=None
        )
        cb = _make_candidate(
            key_alias="key_b", provider_name="openrouter", model="llama-3", model_group=None
        )

        await qm.record_usage(ca, _make_usage(total=100))
        await qm.record_usage(ca, _make_usage(total=100))

        assert await qm.pre_check(ca) is False  # key_a exhausted
        assert await qm.pre_check(cb) is True  # key_b untouched


class TestQuotaManagerDailyReset:
    async def test_daily_utc_resets_at_midnight(self):
        backend = InMemoryBackend(clock=lambda: time.time())
        provider = _make_provider_with_group_quota(limit=100)

        with freeze_time("2026-02-21 23:30:00", tz_offset=0):
            qm = QuotaManager(backend, [provider])
            candidate = _make_candidate()
            await qm.record_usage(candidate, _make_usage(total=100))
            assert await qm.pre_check(candidate) is False

        with freeze_time("2026-02-22 00:01:00", tz_offset=0):
            # After midnight UTC — quota should have reset via TTL
            assert await qm.pre_check(candidate) is True


class TestQuotaManagerRollingWindow:
    async def test_rolling_window_expires(self):
        backend = InMemoryBackend(clock=lambda: time.time())
        provider = ProviderConfig(
            name="openai",
            client_type="openai",
            priority=1,
            model_groups=[
                ModelGroupConfig(
                    name="flagship",
                    tier=1,
                    models=["gpt-4o"],
                    token_quota=TokenQuotaConfig(
                        limit=100, reset="rolling_window", window_seconds=3600
                    ),
                ),
            ],
            keys=[KeyConfig(token="sk-test", alias="main")],
        )

        with freeze_time("2026-02-21 12:00:00", tz_offset=0):
            qm = QuotaManager(backend, [provider])
            candidate = _make_candidate()
            await qm.record_usage(candidate, _make_usage(total=100))
            assert await qm.pre_check(candidate) is False

        with freeze_time("2026-02-21 13:01:00", tz_offset=0):
            # After 1 hour window — quota should have reset
            assert await qm.pre_check(candidate) is True


class TestQuotaManagerExceededError:
    @freeze_time("2026-02-21 12:00:00", tz_offset=0)
    async def test_check_and_raise_throws(self):
        backend = InMemoryBackend()
        provider = _make_provider_with_group_quota(limit=100)
        qm = QuotaManager(backend, [provider])
        candidate = _make_candidate()

        await qm.record_usage(candidate, _make_usage(total=100))

        with pytest.raises(QuotaExceededError) as exc_info:
            await qm.check_and_raise(candidate)

        err = exc_info.value
        assert err.current == 100
        assert err.limit == 100
        assert err.resets_at is not None

    @freeze_time("2026-02-21 12:00:00", tz_offset=0)
    async def test_check_and_raise_passes_when_ok(self):
        backend = InMemoryBackend()
        provider = _make_provider_with_group_quota(limit=1000)
        qm = QuotaManager(backend, [provider])
        candidate = _make_candidate()

        # Should not raise
        await qm.check_and_raise(candidate)
