"""
Coalesced reseed of spend counters from the authoritative DB.

When a Redis spend counter expires (or is missing on a fresh pod), enforcement
must read the current spend from somewhere. The in-process management cache
(`user_api_key_cache.team_membership.spend`, etc.) is per-pod and lags DB
writes from other pods, so trusting it allows budget bypass in multi-pod
deployments. This module reseeds from the authoritative DB instead.

A per-counter singleflight lock collapses concurrent reseeds on the same pod
to one DB query per cold-cache window. The lock dict is bounded LRU to cap
memory in long-lived deployments.
"""

import asyncio
from collections import OrderedDict
from collections.abc import Mapping
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, ClassVar, Final, Optional

from litellm._logging import verbose_proxy_logger
from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy._types import Litellm_EntityType
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.table_repositories import (
    BudgetWindowSpendRepository,
    SpendLogsRepository,
    TeamMembershipRepository,
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import (
    VerificationTokenRepository,
)

if TYPE_CHECKING:
    from litellm.caching.dual_cache import DualCache
    from litellm.proxy.utils import PrismaClient


_WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType(
    {
        "Key": Litellm_EntityType.KEY.value,
        "Team": Litellm_EntityType.TEAM.value,
    }
)

_WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
    {
        "Key": "api_key",
        "Team": "team_id",
    }
)


def _as_utc(value: datetime) -> datetime:
    return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc)


class SpendCounterReseed:
    """
    Reseeds spend counters from the authoritative DB and warms the cache,
    coalesced via per-counter singleflight locks.

    Counter key prefixes map to DB tables:
        spend:key:{token}                 -> LiteLLM_VerificationToken.spend
        spend:team:{team_id}              -> LiteLLM_TeamTable.spend
        spend:team_member:{uid}:{tid}     -> LiteLLM_TeamMembership.spend
        spend:user:{user_id}              -> LiteLLM_UserTable.spend
        spend:org:{org_id}                -> LiteLLM_OrganizationTable.spend

    End-user and tag spend counters intentionally do not reseed here. Their
    auth paths already load the corresponding objects via get_end_user_object()
    and get_tag_objects_batch(); callers pass those values as fallback_spend.
    """

    _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict()
    _registry_lock: ClassVar[asyncio.Lock | None] = None

    @staticmethod
    async def _get_lock(counter_key: str) -> asyncio.Lock:
        if SpendCounterReseed._registry_lock is None:
            SpendCounterReseed._registry_lock = asyncio.Lock()
        async with SpendCounterReseed._registry_lock:
            lock = SpendCounterReseed._locks.get(counter_key)
            if lock is not None:
                SpendCounterReseed._locks.move_to_end(counter_key)
                return lock
            lock = asyncio.Lock()
            SpendCounterReseed._locks[counter_key] = lock
            if len(SpendCounterReseed._locks) > SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE:
                SpendCounterReseed._locks.popitem(last=False)
            return lock

    @staticmethod
    async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None:
        """
        Read the authoritative spend for a counter from the DB.

        Returns the spend value (including 0.0) when the DB is reachable
        and the row exists. Returns None when prisma is unavailable, the
        row is missing, the key format is unrecognized, or the query
        raises. Callers use None to fall back to a caller-supplied source.
        """
        if prisma_client is None:
            return None
        # Per-window key/team counters share prefixes with primary counters
        # but don't correspond to a DB row. Do not reject arbitrary entity IDs
        # or tag names that merely contain ":window:".
        if SpendCounterReseed._is_key_or_team_window_counter(counter_key):
            return None
        try:
            if counter_key.startswith("spend:key:"):
                token: Final = counter_key[len("spend:key:") :]
                row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token})
            elif counter_key.startswith("spend:team_member:"):
                suffix: Final = counter_key[len("spend:team_member:") :]
                if ":" not in suffix:
                    return None
                user_id, team_id = suffix.rsplit(":", 1)
                row = await TeamMembershipRepository(prisma_client).table.find_unique(
                    where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}
                )
            elif counter_key.startswith("spend:team:"):
                team_id = counter_key[len("spend:team:") :]
                row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id})
            elif counter_key.startswith("spend:user:"):
                user_id = counter_key[len("spend:user:") :]
                row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id})
            elif counter_key.startswith("spend:end_user:") or counter_key.startswith("spend:tag:"):
                return None
            elif counter_key.startswith("spend:org:"):
                org_id: Final = counter_key[len("spend:org:") :]
                row = await OrganizationRepository(prisma_client).table.find_unique(where={"organization_id": org_id})
            else:
                return None
        except Exception:
            verbose_proxy_logger.exception("SpendCounterReseed.from_db: failed for %s", counter_key)
            return None
        if row is None:
            return None
        return float(getattr(row, "spend", 0.0) or 0.0)

    @staticmethod
    def _is_key_or_team_window_counter(counter_key: str) -> bool:
        for prefix in ("spend:key:", "spend:team:"):
            if not counter_key.startswith(prefix):
                continue
            _, separator, duration = counter_key.rpartition(":window:")
            if not separator or not duration:
                return False
            try:
                duration_in_seconds(duration)
            except Exception:
                return False
            return True
        return False

    @staticmethod
    async def coalesced(
        prisma_client: Optional["PrismaClient"],
        spend_counter_cache: "DualCache",
        counter_key: str,
        require_cache_warm: bool = False,
    ) -> float | None:
        """
        Reseed a cold spend counter from the DB and warm the cache,
        coalesced via a per-counter lock so concurrent callers (read path
        + write path) collapse to one DB query per cold-cache window.

        Returns the spend value (including 0.0 from a fresh budget reset)
        when the DB read succeeds, or None when the DB is unavailable.
        """
        lock: Final = await SpendCounterReseed._get_lock(counter_key)
        async with lock:
            # Re-check after acquiring the lock. Skip in-memory on a clean
            # Redis miss - in-memory is per-pod-stale.
            redis_clean_miss = False
            if spend_counter_cache.redis_cache is not None:
                try:
                    val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
                    if val is not None:
                        return float(val)
                    redis_clean_miss = True
                except Exception:
                    pass
            if not redis_clean_miss:
                val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
                if val is not None:
                    return float(val)

            db_spend: Final = await SpendCounterReseed.from_db(prisma_client, counter_key)
            if db_spend is None:
                return None
            # Warm even when 0 so subsequent reads hit cache, not DB.
            #
            # Seed via SET NX (cross-pod safe): only one pod initializes the
            # Redis key with db_spend; concurrent seeders read the winner's
            # value. INCRBYFLOAT-of-db_spend from N pods would multiply the
            # counter (N x db_spend) and trigger spurious budget alerts.
            current_value: float = float(db_spend)
            try:
                if spend_counter_cache.redis_cache is not None:
                    seeded: Final = await spend_counter_cache.redis_cache.async_set_cache(
                        key=counter_key,
                        value=db_spend,
                        nx=True,
                    )
                    if seeded:
                        current_value = float(db_spend)
                    else:
                        cached: Final = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
                        current_value = float(cached) if cached is not None else float(db_spend)
                    spend_counter_cache.in_memory_cache.set_cache(
                        key=counter_key,
                        value=current_value,
                    )
                else:
                    await spend_counter_cache.async_increment_cache(key=counter_key, value=db_spend, refresh_ttl=True)
            except Exception:
                verbose_proxy_logger.exception(
                    "SpendCounterReseed.coalesced: failed to warm counter %s",
                    counter_key,
                )
                if require_cache_warm:
                    raise
            return current_value

    @staticmethod
    async def window_from_table(
        prisma_client: Optional["PrismaClient"],
        entity_type: str,
        entity_id: str,
        window_duration: str,
        expected_window_start: datetime,
    ) -> float | None:
        """
        Read the maintained per-window spend row by primary key.

        Returns the row's spend only when the row belongs to the window the
        caller is enforcing, i.e. ``row.window_start >= expected_window_start``.
        A row at or past the expected start was rolled by a pod whose reset_at
        was at least as fresh as this caller's, so it is trusted; an older row
        means the window boundary was crossed and nothing has rolled the row
        yet, so its spend belongs to a previous window.

        Returns None for a missing, stale or unreadable row so the caller falls
        back to the spend-logs aggregate. ``entity_type`` is the counter-facing
        label ("Key"/"Team"); anything else has no row and returns None.
        """
        if prisma_client is None:
            return None
        row_entity_type: Final = _WINDOW_SPEND_ENTITY_TYPES.get(entity_type)
        if row_entity_type is None:
            return None

        try:
            row: Final = await BudgetWindowSpendRepository(prisma_client).table.find_unique(
                where={
                    "entity_type_entity_id_window_duration": {
                        "entity_type": row_entity_type,
                        "entity_id": entity_id,
                        "window_duration": window_duration,
                    }
                }
            )
        except Exception:  # noqa: BLE001  # any read failure (DB, stale prisma client) must degrade to the aggregate path
            verbose_proxy_logger.exception(
                "SpendCounterReseed.window_from_table: failed for %s=%s window=%s",
                entity_type,
                entity_id,
                window_duration,
            )
            return None

        if row is None:
            return None
        if _as_utc(row.window_start) < _as_utc(expected_window_start):
            return None
        return float(row.spend or 0.0)

    @staticmethod
    async def window_from_db(
        prisma_client: Optional["PrismaClient"],
        entity_type: str,
        entity_id: str,
        window_duration: str | None,
        window_start: datetime,
    ) -> float | None:
        """
        Authoritative window spend: the maintained row first, falling back to
        the spend-logs aggregate only when no current row exists.

        The aggregate range-scans an unindexed table, so it must stay a
        transitional path (window configured before the row existed) rather
        than a steady-state read.
        """
        if window_duration is not None:
            from_table: Final = await SpendCounterReseed.window_from_table(
                prisma_client=prisma_client,
                entity_type=entity_type,
                entity_id=entity_id,
                window_duration=window_duration,
                expected_window_start=window_start,
            )
            if from_table is not None:
                return from_table
        return await SpendCounterReseed.window_from_spend_logs(
            prisma_client=prisma_client,
            entity_type=entity_type,
            entity_id=entity_id,
            window_start=window_start,
        )

    @staticmethod
    async def window_from_spend_logs(
        prisma_client: Optional["PrismaClient"],
        entity_type: str,
        entity_id: str,
        window_start: datetime,
    ) -> float | None:
        if prisma_client is None:
            return None

        group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type)
        if group_field is None:
            return None
        where: Final = {
            group_field: entity_id,
            "startTime": {"gte": window_start},
        }

        try:
            response: Final = await SpendLogsRepository(prisma_client).table.group_by(
                by=[group_field],
                where=where,
                sum={"spend": True},
            )
        except Exception:
            verbose_proxy_logger.exception(
                "SpendCounterReseed.window_from_spend_logs: failed for %s=%s",
                entity_type,
                entity_id,
            )
            return None

        if not response:
            return 0.0
        first_row: Final = response[0]
        sum_row: Final = first_row.get("_sum") if isinstance(first_row, dict) else getattr(first_row, "_sum", None)
        spend: Final = sum_row.get("spend") if isinstance(sum_row, dict) else getattr(sum_row, "spend", None)
        return float(spend or 0.0)

    @staticmethod
    async def coalesced_window(
        prisma_client: Optional["PrismaClient"],
        spend_counter_cache: "DualCache",
        counter_key: str,
        entity_type: str,
        entity_id: str,
        window_duration: str | None,
        window_start: datetime,
    ) -> float | None:
        lock: Final = await SpendCounterReseed._get_lock(counter_key)
        async with lock:
            redis_clean_miss = False
            if spend_counter_cache.redis_cache is not None:
                try:
                    val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
                    if val is not None:
                        return float(val)
                    redis_clean_miss = True
                except Exception:
                    pass
            if not redis_clean_miss:
                val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
                if val is not None:
                    return float(val)

            window_spend: Final = await SpendCounterReseed.window_from_db(
                prisma_client=prisma_client,
                entity_type=entity_type,
                entity_id=entity_id,
                window_duration=window_duration,
                window_start=window_start,
            )
            if window_spend is None:
                return None
            try:
                if spend_counter_cache.redis_cache is not None:
                    seeded: Final = await spend_counter_cache.redis_cache.async_set_cache(
                        key=counter_key,
                        value=window_spend,
                        nx=True,
                    )
                    if seeded:
                        current_value = window_spend
                    else:
                        current_cached_value = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
                        if current_cached_value is None:
                            current_value = await spend_counter_cache.redis_cache.async_increment(
                                key=counter_key,
                                value=window_spend,
                            )
                        else:
                            current_value = float(current_cached_value)
                    spend_counter_cache.in_memory_cache.set_cache(
                        key=counter_key,
                        value=current_value,
                    )
                else:
                    await spend_counter_cache.async_increment_cache(key=counter_key, value=window_spend)
            except Exception:
                verbose_proxy_logger.exception(
                    "SpendCounterReseed.coalesced_window: failed to warm counter %s",
                    counter_key,
                )
                raise
            return window_spend