"""Token and request quota management."""

from __future__ import annotations

import math
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta

from llm_rotator._types import Candidate, Usage
from llm_rotator.backends import AbstractStateBackend
from llm_rotator.config import (
    ProviderConfig,
    RequestQuotaConfig,
    ResetSchedule,
    TokenQuotaConfig,
)
from llm_rotator.exceptions import QuotaExceededError

# Index entry: (scope_key, limit, reset_schedule, quota_config)
_QuotaEntry = tuple[str, int, ResetSchedule, TokenQuotaConfig | RequestQuotaConfig]


@dataclass(frozen=True)
class QuotaWarning:
    """Emitted when quota usage crosses the warning threshold."""

    scope: str
    current: int
    limit: int
    percentage: float


class QuotaManager:
    """Manages token and request quotas with support for daily_utc and rolling_window resets."""

    def __init__(
        self,
        backend: AbstractStateBackend,
        providers: list[ProviderConfig],
        warning_threshold: float = 0.8,
    ) -> None:
        self._backend = backend
        self._index = _build_quota_index(providers)
        self._warning_threshold = warning_threshold
        self._warned_scopes: set[str] = set()

    async def pre_check(self, candidate: Candidate) -> bool:
        """Check if candidate has quota remaining. Does not record usage."""
        entries = self._index.get(_candidate_key(candidate))
        if not entries:
            return True

        for scope, limit, _, _ in entries:
            usage = await self._backend.get_quota_usage(scope)
            if usage >= limit:
                return False
        return True

    async def record_usage(
        self, candidate: Candidate, usage: Usage
    ) -> list[QuotaWarning]:
        """Record usage against all applicable quotas for this candidate.

        Returns a list of :class:`QuotaWarning` for each quota that crossed
        the warning threshold during this call.
        """
        entries = self._index.get(_candidate_key(candidate))
        if not entries:
            return []

        warnings: list[QuotaWarning] = []
        now = datetime.now(UTC)
        for scope, limit, reset, quota_cfg in entries:
            ttl = _calculate_ttl(reset, now, quota_cfg)
            if isinstance(quota_cfg, TokenQuotaConfig):
                new_value = await self._backend.increment_quota(
                    scope, usage.total_tokens, ttl
                )
            elif isinstance(quota_cfg, RequestQuotaConfig):
                new_value = await self._backend.increment_quota(scope, 1, ttl)
            else:
                continue

            # Check warning threshold
            if limit > 0 and scope not in self._warned_scopes:
                pct = new_value / limit
                if pct >= self._warning_threshold:
                    self._warned_scopes.add(scope)
                    warnings.append(
                        QuotaWarning(
                            scope=scope,
                            current=new_value,
                            limit=limit,
                            percentage=pct,
                        )
                    )

        return warnings

    async def check_and_raise(self, candidate: Candidate) -> None:
        """Raise QuotaExceededError if any quota is exhausted."""
        entries = self._index.get(_candidate_key(candidate))
        if not entries:
            return

        now = datetime.now(UTC)
        for scope, limit, reset, quota_cfg in entries:
            current = await self._backend.get_quota_usage(scope)
            if current >= limit:
                ttl = _calculate_ttl(reset, now, quota_cfg)
                resets_at = now + timedelta(seconds=ttl)
                raise QuotaExceededError(
                    scope=scope, current=current, limit=limit, resets_at=resets_at
                )


# --- Internal helpers ---


def _candidate_key(candidate: Candidate) -> str:
    """Build lookup key for a candidate."""
    return f"{candidate.provider_name}:{candidate.key_alias}:{candidate.model_group or '_'}"


def _build_quota_index(
    providers: list[ProviderConfig],
) -> dict[str, list[_QuotaEntry]]:
    """Build a lookup index: candidate_key -> list of quota entries."""
    index: dict[str, list[_QuotaEntry]] = {}

    for provider in providers:
        # Group-level token quotas
        if provider.model_groups:
            for group in provider.model_groups:
                if group.token_quota:
                    scope = f"{provider.name}/{group.name}"
                    entry: _QuotaEntry = (
                        scope,
                        group.token_quota.limit,
                        group.token_quota.reset,
                        group.token_quota,
                    )
                    for key in provider.keys:
                        ck = f"{provider.name}:{key.alias}:{group.name}"
                        index.setdefault(ck, []).append(entry)

        # Key-level request quotas
        for key in provider.keys:
            if key.request_quota:
                scope = f"{provider.name}/{key.alias}/requests"
                entry = (
                    scope,
                    key.request_quota.limit,
                    key.request_quota.reset,
                    key.request_quota,
                )
                if provider.model_groups:
                    for group in provider.model_groups:
                        ck = f"{provider.name}:{key.alias}:{group.name}"
                        index.setdefault(ck, []).append(entry)
                else:
                    ck = f"{provider.name}:{key.alias}:_"
                    index.setdefault(ck, []).append(entry)

    return index


def _calculate_ttl(
    reset: ResetSchedule,
    now: datetime,
    quota_cfg: TokenQuotaConfig | RequestQuotaConfig,
) -> int:
    """Calculate TTL in seconds until next reset."""
    if reset == ResetSchedule.DAILY_UTC:
        tomorrow_midnight = now.replace(hour=0, minute=0, second=0, microsecond=0)
        tomorrow_midnight += timedelta(days=1)
        return max(1, math.ceil((tomorrow_midnight - now).total_seconds()))
    if reset == ResetSchedule.ROLLING_WINDOW:
        window = getattr(quota_cfg, "window_seconds", None)
        if window is None:
            return 3600
        return window
    return 3600
