"""TokenBucket — distributed token-bucket rate limiter over NATS JetStream KV.

generalizes 14-eng-ai-survey's Redis Lua-script token bucket
(``RedisRateLimiter``) and its own separate, non-atomic burst-limit
reimplementation of the same idea
(``RateLimitingCallback._check_burst_limit``, a plain ``GET`` then
``SET`` -- not atomic, a real race window between the two calls) into
one reusable, correct-by-construction primitive: tokens refill
continuously at a fixed rate up to a capacity ceiling; each claim
consumes tokens atomically via a CAS read-modify-write loop, the same
shape :class:`KVLease`/:class:`IdempotencyKeyStore` (this module's
siblings) already use.

usage::

    bucket = TokenBucket(
        nats_client, bucket_name="llm_throughput", refill_rate=2.0, capacity=10.0
    )
    outcome = await bucket.claim("global")
    if not outcome.claimed:
        # outcome.retry_after_seconds tells the caller how long until a
        # token frees up -- useful for an HTTP 429 Retry-After header
        ...

    # or block until a token is available (or 5s elapses):
    outcome = await bucket.claim("global", max_wait_seconds=5.0)

all KV envelope payloads flow through
:func:`threetears.core.serialization.serialize_to_json` /
:func:`deserialize_from_json`, matching :class:`KVLease`'s convention.
"""

from __future__ import annotations

import asyncio
import random
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any, Final

from threetears.core.serialization import deserialize_from_json, serialize_to_json
from threetears.observe import get_logger

if TYPE_CHECKING:
    # From the submodule, not the package: these three are Protocols that
    # `threetears.nats` stopped re-exporting when its nats-py-backed surface went lazy.
    # Annotation-only, so the eager `kv` import here costs an L1 consumer nothing.
    from threetears.nats.kv import KvBucketLike, KvCapable

__all__ = [
    "TokenClaimResult",
    "TokenBucket",
    "TokenBucketConflict",
]

log = get_logger(__name__)

#: bucket-level KV TTL default -- an abandoned key (e.g. a per-user burst
#: bucket for a user who stopped sending requests) is forgotten after this
#: long, rather than growing the bucket's key set forever. the NEXT claim
#: against a forgotten key just starts fresh at full capacity (identical
#: observable behavior to a live key that happens to be at full capacity),
#: so this is a pure cleanup knob, not a correctness one.
_DEFAULT_KV_TTL: Final[timedelta] = timedelta(hours=1)

#: bounded CAS retry budget for claim()'s read-modify-write, when the CAS
#: update loses to a DIFFERENT concurrent claimer on the SAME key -- not
#: the same thing as "insufficient tokens" (see claim()'s docstring).
#: deliberately higher than IdempotencyKeyStore's 8 (this module's sibling):
#: this primitive's whole point is many pods hammering the SAME shared
#: bucket key -- a genuinely hotter access pattern than a mostly-distinct-
#: per-operation idempotency key. 30 matches 14-eng-ai-survey's own
#: empirically-tuned constant for this exact shape (IndexesData/
#: SplitAssignmentsData, tuned against a live 20-way concurrent
#: integration test after 8 proved insufficient under real multi-
#: connection contention -- confirmed here too, via this package's own
#: real-NATS integration test).
_CAS_MAX_RETRIES: Final[int] = 30

#: full-jitter backoff bound between CAS retries, seconds.
_CAS_RETRY_BACKOFF_SECONDS: Final[float] = 0.02


class TokenBucketConflict(RuntimeError):
    """raised when the CAS retry budget is exhausted claiming from a bucket.

    signals persistent write contention on ONE bucket key -- should be
    vanishingly rare given the bounded full-jitter backoff; exhausting
    the budget indicates unusually high concurrent claim pressure on a
    single key, rather than ordinary contention.
    """


@dataclass(frozen=True, slots=True)
class TokenClaimResult:
    """outcome of :meth:`TokenBucket.claim`.

    unlike :class:`KVLease`'s raise-on-denial convention, a rate-limit
    miss is an expected, common, per-request outcome rather than an
    exceptional one -- ``claim`` never raises for "not enough tokens
    right now," it always returns a result the caller inspects. this
    matches :class:`IdempotencyKeyStore`'s own claim()-returns-a-result
    convention more closely than :class:`KVLease`'s mutex semantics.

    :param claimed: whether tokens were successfully consumed
    :ptype claimed: bool
    :param tokens_remaining: bucket's token count immediately after this
        claim attempt (post-consume when claimed, current snapshot when not)
    :ptype tokens_remaining: float
    :param retry_after_seconds: estimated seconds until enough tokens will
        have refilled for a future claim of the same size to succeed;
        ``0.0`` when claimed
    :ptype retry_after_seconds: float
    """

    claimed: bool
    tokens_remaining: float
    retry_after_seconds: float


@dataclass
class _BucketState:
    """internal decoded KV value.

    :ivar tokens: token count as of ``last_refill``
    :ivar last_refill: timezone-aware datetime tokens were last computed at
    """

    tokens: float
    last_refill: datetime


def _encode_state(tokens: float, last_refill: datetime) -> bytes:
    """serialize bucket state to JSON bytes for KV storage.

    :param tokens: token count to store
    :ptype tokens: float
    :param last_refill: timezone-aware datetime this token count is as-of
    :ptype last_refill: datetime
    :return: JSON-encoded bytes suitable for a KV value
    :rtype: bytes
    """
    payload: dict[str, Any] = {"tokens": tokens, "last_refill": last_refill.isoformat()}
    return serialize_to_json(payload)


def _decode_state(data: bytes) -> _BucketState:
    """deserialize stored KV value bytes back to bucket state.

    :param data: bytes payload as returned from a KV bucket read
    :ptype data: bytes
    :return: decoded bucket state
    :rtype: _BucketState
    :raises ValueError: if payload is malformed or timestamp unparseable
    """
    raw = deserialize_from_json(data, field_types={})
    tokens = float(raw["tokens"])
    last_refill = datetime.fromisoformat(str(raw["last_refill"]))
    return _BucketState(tokens=tokens, last_refill=last_refill)


class TokenBucket:
    """distributed token-bucket rate limiter over one shared NATS JetStream KV bucket.

    one instance may serve many independent buckets, each keyed by a
    caller-supplied ``key`` (e.g. per-user, per-endpoint, or a single
    ``"global"`` key) -- ``refill_rate``/``capacity`` are fixed per
    :class:`TokenBucket` instance and apply uniformly to every key it
    serves; callers needing different rates for different keys construct
    separate instances (optionally sharing one underlying KV bucket_name
    only if their key namespaces cannot collide, though separate
    bucket_names is the safer default -- see IdempotencyKeyStore's own
    "pick a bucket dedicated to one domain" guidance).

    KV bucket binding is lazy (deferred to the first operation), matching
    :class:`KVLease`/:class:`IdempotencyKeyStore`'s construction style
    within this package.
    """

    def __init__(
        self,
        nats_client: "KvCapable",
        *,
        bucket_name: str,
        refill_rate: float,
        capacity: float,
        kv_ttl: timedelta | None = _DEFAULT_KV_TTL,
    ) -> None:
        """configure the bucket; defer KV bucket binding until first use.

        :param nats_client: connected canonical :class:`threetears.nats.kv.KvCapable`;
            the store opens its KV bucket through :meth:`KvCapable.kv_bucket`
        :ptype nats_client: KvCapable
        :param bucket_name: KV bucket suffix; the wrapper prefixes it with the
            namespace. pick a bucket dedicated to one rate-limiting domain so
            unrelated keys never collide across surfaces
        :ptype bucket_name: str
        :param refill_rate: tokens added per second, continuously (not in
            discrete ticks) -- a claim 0.5 seconds after the last sees
            ``0.5 * refill_rate`` additional tokens available, matching the
            Redis Lua script's own fractional-elapsed-time math
        :ptype refill_rate: float
        :param capacity: maximum tokens one key can accumulate (the burst
            ceiling); refill never pushes a key's token count above this
        :ptype capacity: float
        :param kv_ttl: bucket-level KV TTL for inactive keys; defaults to
            :data:`_DEFAULT_KV_TTL`. ``None`` means keys never expire --
            rarely correct when the key space grows with something
            unbounded (e.g. per-user keys), since an abandoned key then
            lingers in the KV bucket forever
        :ptype kv_ttl: timedelta | None
        :return: none
        :rtype: None
        :raises ValueError: if refill_rate or capacity is not positive
        """
        if refill_rate <= 0:
            raise ValueError(f"refill_rate must be positive, got {refill_rate}")
        if capacity <= 0:
            raise ValueError(f"capacity must be positive, got {capacity}")
        self._client = nats_client
        self._bucket_name = bucket_name
        self._refill_rate = refill_rate
        self._capacity = capacity
        self._kv_ttl = kv_ttl
        self._bucket: "KvBucketLike | None" = None
        self._bucket_lock = asyncio.Lock()

    @property
    def bucket_name(self) -> str:
        """the configured KV bucket suffix.

        :return: bucket name
        :rtype: str
        """
        return self._bucket_name

    async def claim(
        self, key: str = "default", *, tokens: float = 1.0, max_wait_seconds: float = 0.0
    ) -> TokenClaimResult:
        """attempt to claim ``tokens`` tokens from bucket ``key``.

        ``max_wait_seconds=0.0`` (the default): non-blocking, single
        attempt, returns immediately with ``claimed=False`` if
        insufficient tokens are available right now.

        ``max_wait_seconds>0.0``: blocks, retrying until claimed or the
        deadline elapses, matching :meth:`KVLease.acquire`'s own
        blocking-wait convention -- though here the wait sleeps toward
        the estimated refill time rather than a fixed poll interval,
        since the exact time a token becomes available is computable.

        never raises for "not enough tokens" (see :class:`TokenClaimResult`'s
        docstring for why) -- only raises on a genuine KV transport
        failure or exhausted CAS-contention retry budget.

        :param key: bucket key; independent token buckets never interact
        :ptype key: str
        :param tokens: number of tokens this claim consumes
        :ptype tokens: float
        :param max_wait_seconds: total seconds caller is willing to block;
            ``0.0`` disables blocking and forces a single fail-fast attempt
        :ptype max_wait_seconds: float
        :return: claim outcome; caller proceeds only when ``claimed`` is true
        :rtype: TokenClaimResult
        :raises ValueError: if ``tokens`` exceeds bucket capacity (can never
            be satisfied by any amount of waiting)
        :raises TokenBucketConflict: if the CAS retry budget is exhausted
        :raises threetears.nats.KvError: on a KV transport failure
        """
        if tokens > self._capacity:
            raise ValueError(f"cannot claim {tokens} tokens: exceeds bucket capacity {self._capacity}")
        bucket = await self._ensure_bucket()
        deadline = datetime.now(UTC) + timedelta(seconds=max_wait_seconds) if max_wait_seconds > 0 else None
        while True:
            result = await self._attempt(bucket, key, tokens)
            if result.claimed:
                return result
            if deadline is None:
                return result
            remaining = (deadline - datetime.now(UTC)).total_seconds()
            if remaining <= 0:
                return result
            await asyncio.sleep(max(0.0, min(result.retry_after_seconds, remaining)))

    async def refund(self, key: str = "default", *, tokens: float = 1.0) -> float:
        """Put back tokens claimed for work that never happened. Returns the resulting count.

        :meth:`claim` consumes, and until now the only recovery was refill over time -- so a
        caller cancelled between claiming its turn and doing the work held the bucket down for
        as long as the refill rate took to make it up. That is invisible in one instance and
        compounds under repeated cancellation: a pod restarting in a loop can hold a key's
        budget near zero while doing nothing at all.

        **Best effort, and never raises.** This exists to be called from an exception handler
        that is already unwinding, most often a cancellation. Raising there would replace a
        recoverable throughput dip with a lost error, so every failure -- KV outage, CAS
        exhaustion -- is logged and swallowed; the bucket then refills on its own, which is
        exactly the state this method exists to shorten rather than to guarantee.

        Capped at capacity, so a double refund or a refund of tokens that were never claimed
        cannot mint budget the bucket never had. That makes it safe to call unconditionally in
        a handler that may or may not have got as far as claiming.

        :param key: bucket key the tokens were claimed from
        :ptype key: str
        :param tokens: how many to return; must not exceed capacity
        :ptype tokens: float
        :return: tokens available after the refund. ``-1.0`` means it could not be applied and
            the bucket will refill instead -- a KV failure or CAS exhaustion. An untouched key
            returns capacity, since nothing was consumed from it to give back
        :rtype: float
        """
        if tokens <= 0:
            return -1.0
        try:
            bucket = await self._ensure_bucket()
            for attempt in range(_CAS_MAX_RETRIES):
                now = datetime.now(UTC)
                entry = await bucket.get_entry(key=key)
                if entry is None:
                    # No key means nothing was ever consumed from it; a refund would be
                    # inventing budget rather than returning it.
                    return self._capacity
                value, revision = entry
                state = _decode_state(value)
                elapsed = max(0.0, (now - state.last_refill).total_seconds())
                current = min(self._capacity, state.tokens + elapsed * self._refill_rate)
                restored = min(self._capacity, current + tokens)
                if await bucket.update(key=key, value=_encode_state(restored, now), revision=revision) is not None:
                    return restored
                if attempt < _CAS_MAX_RETRIES - 1:
                    backoff = random.uniform(0, _CAS_RETRY_BACKOFF_SECONDS)  # noqa: S311 - jitter, not security
                    await asyncio.sleep(backoff)
            log.warning("token bucket: exhausted CAS retries refunding %s tokens to %r", tokens, key)
        except Exception:  # noqa: BLE001 -- prawduct:allow prawduct/broad-except -- a refund runs while the caller is already unwinding; raising would turn a self-healing throughput dip into a lost error. Logged with its traceback below
            log.exception("token bucket: could not refund %s tokens to %r; it will refill instead", tokens, key)
        return -1.0

    async def _attempt(self, bucket: "KvBucketLike", key: str, tokens: float) -> TokenClaimResult:
        """one refill-then-maybe-consume pass.

        retries internally, bounded by :data:`_CAS_MAX_RETRIES`, only
        when the CAS update loses to a DIFFERENT concurrent claimer
        touching the identical key this instant -- NOT when tokens are
        merely insufficient, which is a definitive (non-retryable)
        outcome this method returns directly.

        :param bucket: backing wrapper KV bucket (from :meth:`_ensure_bucket`)
        :ptype bucket: KvBucketLike
        :param key: bucket key
        :ptype key: str
        :param tokens: number of tokens this claim consumes
        :ptype tokens: float
        :return: definitive claim outcome for this attempt
        :rtype: TokenClaimResult
        :raises TokenBucketConflict: if the CAS retry budget is exhausted
        """
        for attempt in range(_CAS_MAX_RETRIES):
            now = datetime.now(UTC)
            entry = await bucket.get_entry(key=key)
            if entry is None:
                current_tokens = self._capacity
                revision = None
            else:
                value, revision = entry
                state = _decode_state(value)
                elapsed = max(0.0, (now - state.last_refill).total_seconds())
                current_tokens = min(self._capacity, state.tokens + elapsed * self._refill_rate)

            if current_tokens < tokens:
                shortfall = tokens - current_tokens
                retry_after = shortfall / self._refill_rate
                return TokenClaimResult(claimed=False, tokens_remaining=current_tokens, retry_after_seconds=retry_after)

            new_tokens = current_tokens - tokens
            payload = _encode_state(new_tokens, now)
            if revision is None:
                # KvBucketLike.create returns the new revision on success or
                # None on CAS conflict (another claimer created the key
                # between our get_entry miss and this create) -- retry.
                new_revision = await bucket.create(key=key, value=payload)
            else:
                # KvBucketLike.update returns None on revision mismatch --
                # another claimer consumed from this key between our
                # get_entry and this update -- retry.
                new_revision = await bucket.update(key=key, value=payload, revision=revision)
            if new_revision is not None:
                return TokenClaimResult(claimed=True, tokens_remaining=new_tokens, retry_after_seconds=0.0)
            if attempt < _CAS_MAX_RETRIES - 1:
                backoff = random.uniform(0, _CAS_RETRY_BACKOFF_SECONDS)  # noqa: S311 - jitter, not security
                await asyncio.sleep(backoff)
        raise TokenBucketConflict(f"exhausted {_CAS_MAX_RETRIES} CAS retries claiming from bucket {key!r}")

    async def _ensure_bucket(self) -> "KvBucketLike":
        """open (or bind) the TTL'd KV bucket once; async-safe lazy init."""
        if self._bucket is not None:
            return self._bucket
        async with self._bucket_lock:
            if self._bucket is None:
                self._bucket = await self._client.kv_bucket(
                    name=self._bucket_name,
                    ttl=self._kv_ttl,
                    storage="memory",
                    create_if_missing=True,
                    history=1,
                )
                log.info("TokenBucket bound bucket %s", self._bucket_name)
        return self._bucket
