"""Rate-limit primitive — TokenBucket (C-2, v0.37.0).

Stability tier: **experimental** (AP-11) — added in 0.37.0.
Wave 3 federation primitive: formalises the TokenBucket shape
that mgf-fastapi 0.6.0 (`mgf.fastapi.TokenBucket` /
`RateLimitMiddleware`) and mgf-django 0.4.0
(`mgf.django.TokenBucket` / `ThrottleMiddleware`) both ship
locally today.

Both siblings rebase on this module in their NEXT minor:

    # mgf-fastapi 0.7.0 + mgf-django 0.5.0
    from mgf.common.rate_limit import TokenBucket

    # Plus drop the local _rate_limit.py / _throttle.py
    # TokenBucket class (delete the duplicate).

The middleware shells stay in each sibling (the
Starlette / Django shapes differ); only the algorithmic core
deduplicates.

What ships here
---------------

- :class:`TokenBucket` — pure-Python token-bucket. Per-key
  state; thread-safe. Suitable for in-process rate limiting.
- :class:`Bucket` — Protocol abstraction so production
  consumers can wire a Redis-backed bucket externally without
  changing call sites. ``mgf.common.rate_limit.TokenBucket``
  satisfies this Protocol; future ``RedisBucket`` /
  ``MemcachedBucket`` siblings will too.

The Protocol surface (the contract) is:

.. code-block:: python

    @runtime_checkable
    class Bucket(Protocol):
        def try_consume(self, key: str) -> tuple[bool, float]: ...
        # Returns (allowed, retry_after_seconds).

Multi-replica deployments wanting consistent rate-limit
state across all replicas wire a Redis-backed implementation
on the same Protocol — middleware code doesn't change.

Cross-references
----------------

- **AP-11** — stability tiers.
- **C-2** federation row.
- :mod:`mgf.fastapi._rate_limit` — Starlette/FastAPI middleware
  shell that wires this primitive (rebases in 0.7+).
- :mod:`mgf.django._throttle` — Django middleware shell
  (rebases in 0.5+).
"""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Protocol, runtime_checkable

if TYPE_CHECKING:
    from collections.abc import Callable

__all__ = [
    "Bucket",
    "TokenBucket",
]


@runtime_checkable
class Bucket(Protocol):
    """Rate-limit bucket Protocol.

    The minimal shape every rate-limit backend must satisfy.
    Backends decide their state model (in-process, Redis,
    Memcached, ...); the read shape is universal:
    ``try_consume(key) -> (allowed, retry_after_seconds)``.
    """

    def try_consume(self, key: str) -> tuple[bool, float]:
        """Attempt to take 1 token from the bucket for ``key``.

        Returns:
            ``(True, 0.0)`` if a token was available + consumed.
            ``(False, retry_after_seconds)`` if empty.
        """
        ...


@dataclass
class _BucketState:
    """Per-key bucket state. Internal."""

    tokens: float
    last_refill_at: float
    lock: threading.Lock = field(default_factory=threading.Lock)


class TokenBucket:
    """In-process token-bucket rate limiter. Thread-safe.

    The bucket starts FULL (``capacity`` tokens). Each
    :meth:`try_consume` call attempts to remove 1 token; if
    available, returns ``(True, 0.0)``. If empty, returns
    ``(False, retry_after_seconds)`` — the caller surfaces the
    delay via ``Retry-After: N`` HTTP header or equivalent.

    Refill is continuous (``refill_per_second`` tokens / second);
    the bucket caps at ``capacity``.

    State: one bucket per ``key``. The state dict is unbounded
    by design — a long-running process accumulates one entry
    per unique key seen. For multi-tenant deployments wanting
    eviction, wrap the bucket externally OR use a Redis-backed
    implementation satisfying :class:`Bucket`.

    Parameters
    ----------
    capacity
        Maximum tokens (burst allowance). A client returning
        after idle can consume ``capacity`` in quick succession
        before being rate-limited.
    refill_per_second
        Refill rate. ``capacity / refill_per_second`` is the
        "fill-from-empty time" — for ``capacity=10,
        refill_per_second=1``, an exhausted bucket refills in
        10 seconds.
    time_source
        Monotonic clock used for refill accounting. Defaults to
        :func:`time.monotonic`. Injectable so tests can advance time
        deterministically (the refill-over-time behaviour is the whole
        point of the bucket and is otherwise untestable hermetically).
        MUST be non-decreasing — pass a monotonic source, not
        ``time.time``.
    """

    def __init__(
        self,
        *,
        capacity: int,
        refill_per_second: float,
        time_source: Callable[[], float] = time.monotonic,
    ) -> None:
        if capacity <= 0:
            raise ValueError(
                f"TokenBucket capacity must be > 0; got {capacity}",
            )
        if refill_per_second <= 0:
            raise ValueError(
                f"TokenBucket refill_per_second must be > 0; got {refill_per_second}",
            )
        self.capacity = capacity
        self.refill_per_second = float(refill_per_second)
        self._now = time_source
        self._states: dict[str, _BucketState] = {}
        self._states_lock = threading.Lock()

    def try_consume(self, key: str) -> tuple[bool, float]:
        """Take 1 token from ``key``'s bucket. See Protocol docstring."""
        now = self._now()
        state = self._get_or_create(key, now)
        with state.lock:
            elapsed = now - state.last_refill_at
            state.tokens = min(
                self.capacity,
                state.tokens + elapsed * self.refill_per_second,
            )
            state.last_refill_at = now
            if state.tokens >= 1.0:
                state.tokens -= 1.0
                return True, 0.0
            deficit = 1.0 - state.tokens
            return False, deficit / self.refill_per_second

    def _get_or_create(self, key: str, now: float) -> _BucketState:
        with self._states_lock:
            state = self._states.get(key)
            if state is None:
                state = _BucketState(
                    tokens=float(self.capacity),
                    last_refill_at=now,
                )
                self._states[key] = state
            return state

    def reset(self, key: str | None = None) -> None:
        """Reset bucket state.

        ``key=None`` clears ALL state. ``key="x"`` clears just
        that bucket. Test-fixture ergonomics; avoid in
        production (a reset gives the keyed client an
        unintended free burst).
        """
        with self._states_lock:
            if key is None:
                self._states.clear()
            else:
                self._states.pop(key, None)
