"""Unit tests for QuotaService (U5 — per-department LLM quotas).

Covers:
- Happy path: set → get → list → delete quotas
- Quota check: token_limit (allowed / denied)
- Quota check: cost_limit (allowed / denied)
- Quota check: no quota set → always allowed
- model_whitelist serialization (list → JSON → list)
- Upsert: set same quota twice updates the value
- Validation: invalid quota_type / period raises ValueError
- is_model_allowed: whitelist enforcement
- Singleton helpers (get/set_quota_service)
"""

from __future__ import annotations

import uuid
from pathlib import Path

import pytest

from agentkit.server.admin.quota_service import (
    QuotaService,
    get_quota_service,
    set_quota_service,
)
from agentkit.server.auth.models import init_auth_db


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
async def fresh_db(tmp_path: Path) -> Path:
    """A brand-new auth DB on a fresh path (no data)."""
    db_path = tmp_path / "auth.db"
    await init_auth_db(db_path)
    return db_path


@pytest.fixture
def service() -> QuotaService:
    return QuotaService()


def _random_dept_id() -> str:
    return str(uuid.uuid4())


# ---------------------------------------------------------------------------
# Quota CRUD happy path
# ---------------------------------------------------------------------------


class TestQuotaCrudHappyPath:
    async def test_set_token_limit_returns_quota_dict(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        quota = await service.set_quota(fresh_db, dept_id, "token_limit", 1000, period="daily")
        assert quota["department_id"] == dept_id
        assert quota["quota_type"] == "token_limit"
        assert quota["limit_value"] == 1000
        assert quota["period"] == "daily"
        assert "id" in quota
        assert "updated_at" in quota

    async def test_set_cost_limit_returns_quota_dict(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        quota = await service.set_quota(fresh_db, dept_id, "cost_limit", 5000, period="monthly")
        assert quota["quota_type"] == "cost_limit"
        assert quota["limit_value"] == 5000
        assert quota["period"] == "monthly"

    async def test_set_model_whitelist_serializes_list(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        whitelist = ["gpt-4o", "claude-sonnet-4-20250514", "gemini-pro"]
        quota = await service.set_quota(
            fresh_db, dept_id, "model_whitelist", whitelist, period="daily"
        )
        assert quota["quota_type"] == "model_whitelist"
        # The deserialized value should match the input list.
        assert quota["limit_value"] == whitelist

    async def test_get_returns_set_quota(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
        fetched = await service.get_quota(fresh_db, dept_id, "token_limit")
        assert fetched is not None
        assert fetched["limit_value"] == 1000
        assert fetched["quota_type"] == "token_limit"

    async def test_get_returns_none_for_unset_quota(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        fetched = await service.get_quota(fresh_db, dept_id, "token_limit")
        assert fetched is None

    async def test_list_returns_all_quotas_for_department(
        self, service: QuotaService, fresh_db: Path
    ):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "token_limit", 1000, period="daily")
        await service.set_quota(fresh_db, dept_id, "cost_limit", 5000, period="monthly")
        await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o"], period="daily")
        quotas = await service.list_department_quotas(fresh_db, dept_id)
        assert len(quotas) == 3
        types = {q["quota_type"] for q in quotas}
        assert types == {"token_limit", "cost_limit", "model_whitelist"}

    async def test_list_returns_empty_for_department_with_no_quotas(
        self, service: QuotaService, fresh_db: Path
    ):
        dept_id = _random_dept_id()
        quotas = await service.list_department_quotas(fresh_db, dept_id)
        assert quotas == []

    async def test_delete_removes_quota(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
        deleted = await service.delete_quota(fresh_db, dept_id, "token_limit")
        assert deleted is True
        # Confirm it's gone.
        assert await service.get_quota(fresh_db, dept_id, "token_limit") is None

    async def test_delete_returns_false_for_unset_quota(
        self, service: QuotaService, fresh_db: Path
    ):
        dept_id = _random_dept_id()
        deleted = await service.delete_quota(fresh_db, dept_id, "token_limit")
        assert deleted is False


# ---------------------------------------------------------------------------
# Upsert behavior
# ---------------------------------------------------------------------------


class TestQuotaUpsert:
    async def test_set_same_quota_twice_updates_value(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
        updated = await service.set_quota(fresh_db, dept_id, "token_limit", 2000)
        assert updated["limit_value"] == 2000
        # Only one row should exist.
        quotas = await service.list_department_quotas(fresh_db, dept_id)
        assert len(quotas) == 1
        assert quotas[0]["limit_value"] == 2000

    async def test_set_same_quota_with_different_period_creates_new_row(
        self, service: QuotaService, fresh_db: Path
    ):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "token_limit", 1000, period="daily")
        await service.set_quota(fresh_db, dept_id, "token_limit", 30000, period="monthly")
        quotas = await service.list_department_quotas(fresh_db, dept_id)
        assert len(quotas) == 2
        periods = {q["period"] for q in quotas}
        assert periods == {"daily", "monthly"}

    async def test_upsert_model_whitelist_replaces_list(
        self, service: QuotaService, fresh_db: Path
    ):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o", "claude"])
        await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o-mini"])
        fetched = await service.get_quota(fresh_db, dept_id, "model_whitelist")
        assert fetched is not None
        assert fetched["limit_value"] == ["gpt-4o-mini"]


# ---------------------------------------------------------------------------
# Quota enforcement
# ---------------------------------------------------------------------------


class TestQuotaCheck:
    async def test_token_limit_allowed_when_under_limit(
        self, service: QuotaService, fresh_db: Path
    ):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
        allowed, reason = await service.check_quota(
            fresh_db, dept_id, "token_limit", "daily", current_usage=500
        )
        assert allowed is True
        assert reason == "ok"

    async def test_token_limit_denied_at_limit(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
        allowed, reason = await service.check_quota(
            fresh_db, dept_id, "token_limit", "daily", current_usage=1000
        )
        assert allowed is False
        assert "exceeded" in reason

    async def test_token_limit_denied_over_limit(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
        allowed, reason = await service.check_quota(
            fresh_db, dept_id, "token_limit", "daily", current_usage=1500
        )
        assert allowed is False
        assert "exceeded" in reason

    async def test_cost_limit_allowed_when_under_limit(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "cost_limit", 5000, period="monthly")
        allowed, reason = await service.check_quota(
            fresh_db, dept_id, "cost_limit", "monthly", current_usage=2500
        )
        assert allowed is True
        assert reason == "ok"

    async def test_cost_limit_denied_at_limit(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "cost_limit", 5000, period="monthly")
        allowed, _ = await service.check_quota(
            fresh_db, dept_id, "cost_limit", "monthly", current_usage=5000
        )
        assert allowed is False

    async def test_no_quota_set_always_allowed(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        allowed, reason = await service.check_quota(
            fresh_db, dept_id, "token_limit", "daily", current_usage=999999
        )
        assert allowed is True
        assert reason == "ok"

    async def test_model_whitelist_not_a_quota_check(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o"], period="daily")
        allowed, reason = await service.check_quota(
            fresh_db, dept_id, "model_whitelist", "daily", current_usage=0
        )
        assert allowed is False
        assert "is_model_allowed" in reason


# ---------------------------------------------------------------------------
# Model whitelist enforcement
# ---------------------------------------------------------------------------


class TestModelWhitelist:
    async def test_is_model_allowed_when_in_whitelist(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o", "claude"])
        allowed, reason = await service.is_model_allowed(fresh_db, dept_id, "gpt-4o")
        assert allowed is True
        assert reason == "ok"

    async def test_is_model_denied_when_not_in_whitelist(
        self, service: QuotaService, fresh_db: Path
    ):
        dept_id = _random_dept_id()
        await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o", "claude"])
        allowed, reason = await service.is_model_allowed(fresh_db, dept_id, "gemini-pro")
        assert allowed is False
        assert "not in" in reason

    async def test_is_model_allowed_when_no_whitelist(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        # No whitelist set — all models allowed.
        allowed, reason = await service.is_model_allowed(fresh_db, dept_id, "any-model")
        assert allowed is True
        assert reason == "ok"


# ---------------------------------------------------------------------------
# Validation
# ---------------------------------------------------------------------------


class TestQuotaValidation:
    async def test_set_invalid_quota_type_raises(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        with pytest.raises(ValueError, match="Invalid quota_type"):
            await service.set_quota(fresh_db, dept_id, "invalid_type", 1000)

    async def test_set_invalid_period_raises(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        with pytest.raises(ValueError, match="Invalid period"):
            await service.set_quota(fresh_db, dept_id, "token_limit", 1000, period="weekly")

    async def test_set_token_limit_with_list_raises(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        with pytest.raises(ValueError, match="must be an int"):
            await service.set_quota(fresh_db, dept_id, "token_limit", ["gpt-4o"])

    async def test_set_model_whitelist_with_int_raises(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        with pytest.raises(ValueError, match="must be a list"):
            await service.set_quota(fresh_db, dept_id, "model_whitelist", 1000)

    async def test_get_invalid_quota_type_raises(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        with pytest.raises(ValueError, match="Invalid quota_type"):
            await service.get_quota(fresh_db, dept_id, "invalid_type")

    async def test_delete_invalid_period_raises(self, service: QuotaService, fresh_db: Path):
        dept_id = _random_dept_id()
        with pytest.raises(ValueError, match="Invalid period"):
            await service.delete_quota(fresh_db, dept_id, "token_limit", period="weekly")

    async def test_check_quota_invalid_quota_type_raises(
        self, service: QuotaService, fresh_db: Path
    ):
        dept_id = _random_dept_id()
        with pytest.raises(ValueError, match="Invalid quota_type"):
            await service.check_quota(fresh_db, dept_id, "invalid_type", "daily", current_usage=0)


# ---------------------------------------------------------------------------
# Singleton helpers
# ---------------------------------------------------------------------------


class TestSingletonHelpers:
    def test_get_quota_service_returns_singleton(self):
        # Save the original singleton so we don't disturb other tests.
        original = get_quota_service()
        try:
            custom = QuotaService()
            set_quota_service(custom)
            assert get_quota_service() is custom
            # Clearing falls back to a new lazy instance.
            set_quota_service(None)
            new_one = get_quota_service()
            assert new_one is not custom
        finally:
            set_quota_service(original)
