"""Rate-limit middleware — token-bucket-per-key (F-1, v0.6.0).

Stability tier: **experimental** (AP-11) — added in 0.6.0 to
close the F-1 federation row
(`mgf-common/docs/inprogress/pre_release_additions.md` §6.1).
PlasmaMapper is the immediate adopter; inthewords runs the
Django equivalent (D-1, in mgf-django 0.4.0).

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

- :class:`TokenBucket` — pure-Python implementation of the
  Hierarchical Token Bucket algorithm. Thread-safe; suitable
  for in-process rate limiting. Production deployments wanting
  multi-replica enforcement wire a Redis-backed bucket
  externally — the **API** stays the same, the **backend**
  changes.
- :class:`RateLimitMiddleware` — Starlette / FastAPI middleware
  that consults a TokenBucket per-request. Returns ``429 Too
  Many Requests`` with a ``Retry-After`` header when the bucket
  is exhausted.
- :func:`rate_limit_dependency` — alternative shape: a FastAPI
  dependency that enforces per-route. Useful when you want
  finer-grained policies (e.g., "POST /transfers is 5/min;
  GET /things is 1000/min").

Key extraction
--------------

The middleware accepts a **key callable** that maps an incoming
``Request`` to a string identifying the bucket. Common shapes:

- ``lambda req: req.client.host``                # per-IP
- ``lambda req: req.state.user_id`` (or None)    # per-user
- ``lambda req: req.headers.get("X-API-Key")``   # per-API-key

Returning ``None`` from the key function disables rate limiting
for that request (useful for unauthenticated → no limit OR
authenticated → user-keyed shapes).

.. warning::
   The key callable is consumer-controlled and the middleware trusts
   whatever it returns. If you derive the key from a client-supplied
   header — ``X-Forwarded-For`` / ``X-Real-IP`` / ``X-API-Key`` — that
   value is **attacker-controlled** unless a trusted reverse proxy
   strips/overwrites the client's copy. Keying on a spoofable header
   lets an attacker rotate the value per request and bypass the limit
   entirely. Prefer ``request.client.host`` (only behind a proxy you
   trust to set the peer address) or an *authenticated* identity from
   ``request.state``. This mirrors the ``trust_proxy`` footgun called
   out on :class:`mgf.fastapi.security` IpAllowlist.

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

- **AP-11** — stability tiers.
- **mgf.common.flags** — wire a kill-switch to disable the
  middleware at runtime without redeploying.

Cross-consumer pattern: PlasmaMapper (FastAPI SaaS) + inthewords
(Django via D-1 sibling impl).
"""

from __future__ import annotations

import logging
import threading
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    from starlette.requests import Request
    from starlette.responses import Response
    from starlette.types import ASGIApp

_LOGGER = logging.getLogger(__name__)


__all__ = [
    "RateLimitMiddleware",
    "TokenBucket",
]


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

    Internal — not part of the public API. ``tokens`` is the
    current allowance (float for fractional refill); ``last_refill_at``
    is the monotonic timestamp of the last refill.
    """

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


class TokenBucket:
    """Token-bucket rate limiter, in-process, thread-safe.

    The bucket starts FULL (``capacity`` tokens). Each
    :meth:`try_consume` call attempts to remove 1 token; if
    available, returns ``True``. If empty, returns ``False``
    AND a ``retry_after_seconds`` hint (how long until 1 token
    refills).

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

    Per-key state: one bucket per ``key`` (any hashable). 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's state in
    an LRU cache externally OR migrate to a Redis-backed impl.

    Parameters
    ----------
    capacity
        Maximum tokens in the bucket. The "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" — useful intuition: for
        ``capacity=10, refill_per_second=1``, an exhausted
        bucket takes 10 seconds to refill fully.
    """

    def __init__(self, *, capacity: int, refill_per_second: float) -> 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._states: dict[str, _BucketState] = {}
        self._states_lock = threading.Lock()

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

        Returns ``(allowed, retry_after_seconds)``:

        - ``(True, 0.0)`` when a token was available + consumed.
        - ``(False, N)`` when empty; ``N`` is how many seconds
          until 1 token refills (caller surfaces via
          ``Retry-After: N``).
        """
        now = time.monotonic()
        state = self._get_or_create(key, now)

        with state.lock:
            # Refill: add (elapsed * rate) tokens, cap at capacity.
            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

            # Empty. How long until 1 token?
            deficit = 1.0 - state.tokens
            retry_after = deficit / self.refill_per_second
            return False, retry_after

    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.

        With ``key=None``: clear ALL state (every bucket goes
        back to full).
        With ``key="x"``: clear just that bucket.

        Test-fixture ergonomics; production code should avoid
        this (resetting an active client's bucket gives them an
        unintended free burst).
        """
        with self._states_lock:
            if key is None:
                self._states.clear()
            else:
                self._states.pop(key, None)


class RateLimitMiddleware(BaseHTTPMiddleware):
    """Starlette / FastAPI middleware enforcing a TokenBucket per key.

    Usage::

        from mgf.fastapi import RateLimitMiddleware, TokenBucket

        bucket = TokenBucket(capacity=60, refill_per_second=1)
        app.add_middleware(
            RateLimitMiddleware,
            bucket=bucket,
            key=lambda req: req.client.host,
        )

    Behaviour:

    - Calls ``key(request)`` to extract the rate-limit key.
    - If the key is ``None``: middleware passes through (no
      limit). Useful for "anonymous traffic unlimited; per-user
      traffic limited" shapes.
    - If allowed: passes through to the next middleware /
      route.
    - If denied: returns ``429 Too Many Requests`` with
      ``Retry-After: <seconds>`` header + a JSON body
      ``{"error": "rate-limit-exceeded", "retry_after": N}``.

    The 429 response is library-managed; consumers wanting a
    custom body shape can subclass + override
    :meth:`_make_denied_response`.
    """

    def __init__(
        self,
        app: ASGIApp,
        *,
        bucket: TokenBucket,
        key: Callable[[Request], str | None],
    ) -> None:
        super().__init__(app)
        self.bucket = bucket
        self.key = key

    async def dispatch(
        self,
        request: Request,
        call_next: Callable[[Request], Awaitable[Response]],
    ) -> Response:
        key = self.key(request)
        if key is None:
            return await call_next(request)

        allowed, retry_after = self.bucket.try_consume(key)
        if allowed:
            return await call_next(request)

        # Denied — log + 429.
        _LOGGER.info(
            "rate-limit exceeded",
            extra={
                "event": "mgf.fastapi.rate_limit_exceeded",
                "key": key,
                "retry_after_seconds": round(retry_after, 3),
            },
        )
        return self._make_denied_response(retry_after)

    def _make_denied_response(self, retry_after: float) -> Response:
        # Round up to whole seconds for the Retry-After header
        # (per RFC 7231 §7.1.3, the value is delta-seconds, an
        # integer). The JSON body carries the float for clients
        # that can use the precision.
        retry_after_secs = max(1, int(retry_after + 0.999))
        return JSONResponse(
            status_code=429,
            content={
                "error": "rate-limit-exceeded",
                "retry_after": round(retry_after, 3),
            },
            headers={"Retry-After": str(retry_after_secs)},
        )
