# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""TTL + LRU cache with single-flight loading (plan §5.2).

Single-flight is not optional here. A crawl ``Enumerate``s and then ``Stat``s
every item, and a thumbnail probe fires up to four ``Stat``s per asset for
candidate ``.thumbs`` paths — all against the same folder, concurrently. Without
collapsing concurrent loads of one key onto a single upstream request, laziness
turns one preload into a burst of duplicates.
"""

from __future__ import annotations

import asyncio
import time
from collections import OrderedDict
from collections.abc import Awaitable, Callable, Hashable
from dataclasses import dataclass, field
from typing import Generic, TypeVar

K = TypeVar("K", bound=Hashable)
V = TypeVar("V")


class Missing:
    """A cached negative result.

    Distinct from ``None`` so a miss can be cached and served without re-hitting
    HuggingFace: the thumbnail probe generates a steady stream of legitimate
    misses that would otherwise be upstream traffic every time.
    """

    __slots__ = ()

    def __bool__(self) -> bool:
        return False


MISSING = Missing()


@dataclass
class CacheStats:
    hits: int = 0
    misses: int = 0
    loads: int = 0
    negative_hits: int = 0
    evictions: int = 0

    @property
    def hit_rate(self) -> float:
        total = self.hits + self.misses
        return self.hits / total if total else 0.0


@dataclass
class _Entry(Generic[V]):
    value: V | Missing
    expires_at: float


class SingleFlightCache(Generic[K, V]):
    """Async cache: at most one in-flight load per key, TTL + LRU eviction."""

    def __init__(
        self,
        *,
        ttl: float,
        negative_ttl: float,
        max_size: int,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        self._ttl = ttl
        self._negative_ttl = negative_ttl
        self._max_size = max_size
        self._clock = clock
        self._entries: OrderedDict[K, _Entry[V]] = OrderedDict()
        self._inflight: dict[K, asyncio.Future[V | Missing]] = {}
        self.stats = CacheStats()

    def _get_live(self, key: K) -> _Entry[V] | None:
        entry = self._entries.get(key)
        if entry is None:
            return None
        if entry.expires_at <= self._clock():
            del self._entries[key]
            return None
        self._entries.move_to_end(key)
        return entry

    def peek(self, key: K) -> V | Missing | None:
        """Return a cached value without loading. ``None`` means "not cached"."""
        entry = self._get_live(key)
        return entry.value if entry else None

    def put(self, key: K, value: V | Missing) -> None:
        ttl = self._negative_ttl if isinstance(value, Missing) else self._ttl
        self._entries[key] = _Entry(value=value, expires_at=self._clock() + ttl)
        self._entries.move_to_end(key)
        while len(self._entries) > self._max_size:
            self._entries.popitem(last=False)
            self.stats.evictions += 1

    async def get(self, key: K, loader: Callable[[], Awaitable[V | Missing]]) -> V | Missing:
        """Return the cached value, or load it — once, however many callers race."""
        entry = self._get_live(key)
        if entry is not None:
            self.stats.hits += 1
            if isinstance(entry.value, Missing):
                self.stats.negative_hits += 1
            return entry.value

        self.stats.misses += 1

        inflight = self._inflight.get(key)
        if inflight is not None:
            # Someone else is already loading this key: wait for their result
            # rather than issuing a duplicate upstream request.
            return await asyncio.shield(inflight)

        future: asyncio.Future[V | Missing] = asyncio.get_running_loop().create_future()
        self._inflight[key] = future
        try:
            self.stats.loads += 1
            value = await loader()
        except BaseException as exc:
            # Failures are not cached — a transient 503 must not pin a miss for
            # the whole TTL.
            if not future.done():
                future.set_exception(exc)
            raise
        else:
            self.put(key, value)
            if not future.done():
                future.set_result(value)
            return value
        finally:
            self._inflight.pop(key, None)
            # Nobody may be awaiting a failed future; keep the loop quiet about it.
            if future.done() and future.exception() is not None:
                future.exception()

    def clear(self) -> None:
        self._entries.clear()

    def __len__(self) -> int:
        return len(self._entries)


@dataclass
class TokenBucket:
    """Simple rate limiter for upstream calls (plan §8.5)."""

    rate: float
    capacity: float
    _tokens: float = field(default=0.0, init=False)
    _updated: float = field(default=0.0, init=False)

    def __post_init__(self) -> None:
        self._tokens = self.capacity
        self._updated = time.monotonic()

    async def acquire(self) -> None:
        while True:
            now = time.monotonic()
            self._tokens = min(self.capacity, self._tokens + (now - self._updated) * self.rate)
            self._updated = now
            if self._tokens >= 1:
                self._tokens -= 1
                return
            await asyncio.sleep((1 - self._tokens) / self.rate)