"""QuotaService — per-department LLM quota CRUD + enforcement (U5/U7).

This module is the single owner of the ``department_quotas`` table.
Web UI routes (``/api/v1/admin/departments/{id}/quotas``) and the CLI
``agentkit admin llm set-quota`` sub-app both call into
:class:`QuotaService` rather than touching the table directly, keeping
the validation rules (quota_type/period enums, JSON serialization of
``model_whitelist``) in one place.

Quota semantics
---------------
- ``token_limit`` (int): max tokens per ``period``. ``check_quota``
  compares ``current_usage`` (tokens used) against the limit.
- ``cost_limit`` (int, in cents): max spend per ``period``.
  ``check_quota`` compares ``current_usage`` (cost in cents) against
  the limit.
- ``model_whitelist`` (list[str]): allowed model names. Not a quota
  check — it's a model-access check (use :meth:`is_model_allowed`).

The service is a module-level singleton (see :func:`get_quota_service`)
so tests can inject a custom instance via :func:`set_quota_service`.
"""

from __future__ import annotations

import json
import logging
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import aiosqlite

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

#: Valid ``quota_type`` values.
QUOTA_TYPES: frozenset[str] = frozenset({"token_limit", "cost_limit", "model_whitelist"})

#: Valid ``period`` values.
PERIODS: frozenset[str] = frozenset({"daily", "monthly"})


def _now_iso() -> str:
    """Return current UTC time as ISO 8601 string."""
    return datetime.now(timezone.utc).isoformat()


def _new_id() -> str:
    """Return a new UUID4 string."""
    return str(uuid.uuid4())


def _serialize_limit_value(quota_type: str, limit_value: Any) -> str:
    """Serialize ``limit_value`` to a JSON string for storage.

    For ``token_limit`` / ``cost_limit``: ``limit_value`` is an int,
    stored as ``json.dumps(int)`` (e.g. ``"1000"``).
    For ``model_whitelist``: ``limit_value`` is a ``list[str]``, stored
    as ``json.dumps(list)`` (e.g. ``'["gpt-4o", "claude"]'``).
    """
    if quota_type == "model_whitelist":
        if not isinstance(limit_value, list):
            raise ValueError(
                f"model_whitelist limit_value must be a list, got {type(limit_value).__name__}"
            )
        return json.dumps([str(m) for m in limit_value])
    if not isinstance(limit_value, int):
        raise ValueError(
            f"{quota_type} limit_value must be an int, got {type(limit_value).__name__}"
        )
    return json.dumps(limit_value)


def _deserialize_limit_value(quota_type: str, raw: str) -> Any:
    """Deserialize ``limit_value`` from JSON string back to native type."""
    try:
        value = json.loads(raw)
    except json.JSONDecodeError as exc:
        logger.warning("Failed to deserialize limit_value %r: %s", raw, exc)
        return None
    if quota_type == "model_whitelist":
        return list(value) if isinstance(value, list) else []
    return int(value) if isinstance(value, (int, float)) else None


def _validate_quota_type(quota_type: str) -> None:
    if quota_type not in QUOTA_TYPES:
        raise ValueError(f"Invalid quota_type {quota_type!r}; must be one of {sorted(QUOTA_TYPES)}")


def _validate_period(period: str) -> None:
    if period not in PERIODS:
        raise ValueError(f"Invalid period {period!r}; must be one of {sorted(PERIODS)}")


class QuotaService:
    """CRUD + enforcement for per-department LLM quotas.

    All methods are async and take ``db_path: Path`` as the first
    argument (after ``self``). Each method opens its own short-lived
    :class:`aiosqlite.Connection` — there is no shared connection state,
    which keeps the service safe to call from any async context.
    """

    # ------------------------------------------------------------------
    # Quota CRUD
    # ------------------------------------------------------------------

    async def set_quota(
        self,
        db_path: Path,
        department_id: str,
        quota_type: str,
        limit_value: Any,
        period: str = "daily",
    ) -> dict[str, Any]:
        """Upsert a quota for a department.

        Args:
            db_path: Path to the auth SQLite DB.
            department_id: Department id.
            quota_type: One of ``token_limit`` / ``cost_limit`` /
                ``model_whitelist``.
            limit_value: ``int`` for token/cost limits, ``list[str]``
                for model_whitelist.
            period: ``daily`` or ``monthly`` (default ``daily``).

        Returns:
            The upserted quota as a dict (with deserialized
            ``limit_value``).
        """
        _validate_quota_type(quota_type)
        _validate_period(period)
        serialized = _serialize_limit_value(quota_type, limit_value)
        now = _now_iso()
        quota_id = _new_id()

        async with aiosqlite.connect(str(db_path)) as db:
            # Upsert: if (department_id, quota_type, period) exists,
            # update limit_value + updated_at; otherwise insert.
            await db.execute(
                "INSERT INTO department_quotas "
                "(id, department_id, quota_type, limit_value, period, updated_at) "
                "VALUES (?, ?, ?, ?, ?, ?) "
                "ON CONFLICT(department_id, quota_type, period) DO UPDATE SET "
                "limit_value = excluded.limit_value, "
                "updated_at = excluded.updated_at",
                (quota_id, department_id, quota_type, serialized, period, now),
            )
            await db.commit()

        return {
            "id": quota_id,
            "department_id": department_id,
            "quota_type": quota_type,
            "limit_value": _deserialize_limit_value(quota_type, serialized),
            "period": period,
            "updated_at": now,
        }

    async def get_quota(
        self,
        db_path: Path,
        department_id: str,
        quota_type: str,
        period: str = "daily",
    ) -> dict[str, Any] | None:
        """Return a single quota, or ``None`` if not set."""
        _validate_quota_type(quota_type)
        _validate_period(period)
        async with aiosqlite.connect(str(db_path)) as db:
            db.row_factory = aiosqlite.Row
            cursor = await db.execute(
                "SELECT * FROM department_quotas "
                "WHERE department_id = ? AND quota_type = ? AND period = ?",
                (department_id, quota_type, period),
            )
            row = await cursor.fetchone()
        if row is None:
            return None
        return self._row_to_dict(row)

    async def list_department_quotas(
        self,
        db_path: Path,
        department_id: str,
    ) -> list[dict[str, Any]]:
        """List all quotas for a department (all types/periods)."""
        async with aiosqlite.connect(str(db_path)) as db:
            db.row_factory = aiosqlite.Row
            cursor = await db.execute(
                "SELECT * FROM department_quotas WHERE department_id = ? "
                "ORDER BY quota_type ASC, period ASC",
                (department_id,),
            )
            rows = await cursor.fetchall()
        return [self._row_to_dict(row) for row in rows]

    async def delete_quota(
        self,
        db_path: Path,
        department_id: str,
        quota_type: str,
        period: str = "daily",
    ) -> bool:
        """Delete a quota. Returns ``True`` if a row was deleted."""
        _validate_quota_type(quota_type)
        _validate_period(period)
        async with aiosqlite.connect(str(db_path)) as db:
            cursor = await db.execute(
                "DELETE FROM department_quotas "
                "WHERE department_id = ? AND quota_type = ? AND period = ?",
                (department_id, quota_type, period),
            )
            await db.commit()
            return cursor.rowcount > 0

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

    async def check_quota(
        self,
        db_path: Path,
        department_id: str,
        quota_type: str,
        period: str,
        current_usage: int | float,
    ) -> tuple[bool, str]:
        """Check whether ``current_usage`` is within the configured quota.

        Args:
            db_path: Path to the auth SQLite DB.
            department_id: Department id.
            quota_type: ``token_limit`` or ``cost_limit``.
                ``model_whitelist`` is not a quota check — use
                :meth:`is_model_allowed` instead.
            period: ``daily`` or ``monthly``.
            current_usage: Current usage value (tokens for
                ``token_limit``, cents for ``cost_limit``).

        Returns:
            ``(allowed, reason)`` tuple. ``allowed`` is ``True`` if the
            usage is within the limit (or no quota is set). ``reason``
            is ``"ok"`` when allowed, or a human-readable denial reason.
        """
        _validate_quota_type(quota_type)
        _validate_period(period)
        if quota_type == "model_whitelist":
            return (
                False,
                "model_whitelist is not a quota check; use is_model_allowed()",
            )

        quota = await self.get_quota(db_path, department_id, quota_type, period)
        if quota is None:
            return (True, "ok")

        limit = quota["limit_value"]
        if not isinstance(limit, int):
            return (True, "ok")  # malformed limit; allow

        if current_usage >= limit:
            return (
                False,
                f"{quota_type} ({period}) exceeded: usage={current_usage}, limit={limit}",
            )
        return (True, "ok")

    async def is_model_allowed(
        self,
        db_path: Path,
        department_id: str,
        model: str,
    ) -> tuple[bool, str]:
        """Check whether ``model`` is in the department's whitelist.

        If no ``model_whitelist`` quota is set for the department (any
        period), all models are allowed. Otherwise the model must appear
        in the whitelist.

        Returns:
            ``(allowed, reason)`` tuple.
        """
        async with aiosqlite.connect(str(db_path)) as db:
            db.row_factory = aiosqlite.Row
            cursor = await db.execute(
                "SELECT * FROM department_quotas "
                "WHERE department_id = ? AND quota_type = 'model_whitelist' "
                "ORDER BY period ASC LIMIT 1",
                (department_id,),
            )
            row = await cursor.fetchone()
        if row is None:
            return (True, "ok")
        whitelist = _deserialize_limit_value("model_whitelist", row["limit_value"])
        if not isinstance(whitelist, list):
            return (True, "ok")
        if model in whitelist:
            return (True, "ok")
        return (
            False,
            f"model {model!r} not in department {department_id!r} whitelist",
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _row_to_dict(row: aiosqlite.Row | Any) -> dict[str, Any]:
        """Convert a ``department_quotas`` row to a JSON-safe dict."""
        quota_type = row["quota_type"]
        return {
            "id": row["id"],
            "department_id": row["department_id"],
            "quota_type": quota_type,
            "limit_value": _deserialize_limit_value(quota_type, row["limit_value"]),
            "period": row["period"],
            "updated_at": row["updated_at"],
        }


# ---------------------------------------------------------------------------
# Module-level singleton (overridable in tests via set_quota_service)
# ---------------------------------------------------------------------------


_quota_service: QuotaService | None = None


def get_quota_service() -> QuotaService:
    """Return the process-wide :class:`QuotaService` (lazy singleton)."""
    global _quota_service
    if _quota_service is None:
        _quota_service = QuotaService()
    return _quota_service


def set_quota_service(service: QuotaService | None) -> None:
    """Inject a custom :class:`QuotaService` (used by tests)."""
    global _quota_service
    _quota_service = service
