"""
cortex/api/rate_limiter.py — In-memory token bucket rate limiter (#71).

Provides per-owner or per-IP request rate limiting without any external
dependency (no Redis, no database).

Algorithm: Token Bucket
  - Each bucket starts with `burst` tokens.
  - Tokens refill at `requests_per_minute / 60` tokens/second.
  - Each request consumes 1 token.
  - When the bucket is empty, the request is rejected with HTTP 429.

Usage (as FastAPI middleware)::

    from cortex.api.rate_limiter import RateLimiterMiddleware
    app.add_middleware(
        RateLimiterMiddleware,
        requests_per_minute=60,
        burst=10,
    )

Per-request identity is resolved in this order:
  1. ``owner_id`` query parameter
  2. ``X-Owner-Id`` header
  3. Client IP address (``request.client.host``)

Returns HTTP 429 with a ``Retry-After`` header (integer seconds) when the
bucket for an identity is exhausted.
"""
from __future__ import annotations

import collections
import logging
import math
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Optional, Tuple

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Token bucket core (thread-safe, pure Python)
# ---------------------------------------------------------------------------


@dataclass
class _Bucket:
    """A single token-bucket entry."""
    tokens: float
    last_refill: float = field(default_factory=time.monotonic)


class TokenBucketStore:
    """
    Thread-safe in-memory store for per-identity token buckets.

    Uses an ``OrderedDict`` as an LRU cache with a configurable maximum
    size (default 10 000 buckets).  When the limit is reached the oldest
    (least-recently-used) bucket is evicted, preventing unbounded memory
    growth when many unique identities hit the server (#216).

    Args:
        requests_per_minute: Sustained request rate limit (tokens/min).
        burst: Maximum burst size (peak tokens; also initial fill).
        max_buckets: Maximum number of tracked identities (LRU eviction).
    """

    MAX_BUCKETS_DEFAULT = 10_000

    def __init__(
        self,
        requests_per_minute: int = 60,
        burst: int = 10,
        max_buckets: int = MAX_BUCKETS_DEFAULT,
    ) -> None:
        if requests_per_minute <= 0:
            raise ValueError("requests_per_minute must be > 0")
        if burst <= 0:
            raise ValueError("burst must be > 0")

        self.rpm = requests_per_minute
        self.burst = burst
        self.max_buckets = max_buckets
        self._refill_rate: float = requests_per_minute / 60.0  # tokens/sec
        self._buckets: collections.OrderedDict[str, _Bucket] = collections.OrderedDict()
        self._lock = threading.Lock()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def consume(self, identity: str) -> Tuple[bool, float]:
        """
        Attempt to consume one token for *identity*.

        Returns:
            (allowed, retry_after_seconds)
            - ``allowed=True``  → request is within limit; retry_after=0.
            - ``allowed=False`` → rate limit exceeded; retry_after > 0.
        """
        now = time.monotonic()
        with self._lock:
            bucket = self._buckets.get(identity)
            if bucket is None:
                # Evict oldest bucket(s) when at capacity (#216)
                while len(self._buckets) >= self.max_buckets:
                    evicted_key, _ = self._buckets.popitem(last=False)
                    logger.debug("Rate-limiter LRU eviction: %s", evicted_key)
                bucket = _Bucket(tokens=float(self.burst), last_refill=now)
                self._buckets[identity] = bucket
            else:
                # Move to end (most-recently-used)
                self._buckets.move_to_end(identity)

            # Refill tokens since last request
            elapsed = now - bucket.last_refill
            bucket.tokens = min(
                float(self.burst),
                bucket.tokens + elapsed * self._refill_rate,
            )
            bucket.last_refill = now

            if bucket.tokens >= 1.0:
                bucket.tokens -= 1.0
                return True, 0.0
            else:
                # How long until 1 token is available?
                tokens_needed = 1.0 - bucket.tokens
                retry_after = tokens_needed / self._refill_rate
                return False, retry_after

    def reset(self, identity: str) -> None:
        """Reset the bucket for *identity* (useful in tests)."""
        with self._lock:
            self._buckets.pop(identity, None)

    def clear(self) -> None:
        """Clear all buckets (useful in tests)."""
        with self._lock:
            self._buckets.clear()

    def bucket_count(self) -> int:
        """Return number of active buckets (for monitoring)."""
        with self._lock:
            return len(self._buckets)

    def update_limits(self, requests_per_minute: int, burst: int) -> None:
        """Update rate limits and clear all existing buckets."""
        if requests_per_minute <= 0:
            raise ValueError("requests_per_minute must be > 0")
        if burst <= 0:
            raise ValueError("burst must be > 0")
        with self._lock:
            self.rpm = requests_per_minute
            self.burst = burst
            self._refill_rate = requests_per_minute / 60.0
            self._buckets.clear()


# ---------------------------------------------------------------------------
# FastAPI / Starlette middleware
# ---------------------------------------------------------------------------

_SKIP_PATHS = frozenset({"/api/v1/health", "/docs", "/openapi.json", "/redoc"})

try:
    from starlette.middleware.base import BaseHTTPMiddleware
    from starlette.requests import Request
    from starlette.responses import JSONResponse, Response

    class RateLimiterMiddleware(BaseHTTPMiddleware):
        """
        Starlette/FastAPI middleware that enforces per-identity rate limits.

        Identity is resolved from (in priority order):
          1. ``owner_id`` query parameter
          2. ``X-Owner-Id`` request header
          3. Client IP

        Returns HTTP 429 with ``Retry-After`` and ``X-RateLimit-*`` headers
        when the bucket is exhausted.

        Args:
            app:                  ASGI app (injected by FastAPI).
            requests_per_minute:  Sustained request rate (default 60).
            burst:                Burst allowance on top of sustained rate (default 10).
            enabled:              Set False to disable (e.g., in test suites).
        """

        def __init__(
            self,
            app: Any,
            *,
            requests_per_minute: int = 60,
            burst: int = 10,
            enabled: bool = True,
        ) -> None:
            super().__init__(app)
            self._enabled = enabled
            self._store = TokenBucketStore(
                requests_per_minute=requests_per_minute,
                burst=burst,
            )

        # Expose store so http_server.py can update limits after config load
        @property
        def store(self) -> TokenBucketStore:
            return self._store

        def _get_identity(self, request: Request) -> str:
            """Resolve per-request identity string.

            Identity resolution order (most-specific → least-specific):
              1. owner_id query param    → "owner:<id>"  (API consumer identity)
              2. X-Owner-Id header       → "owner:<id>"  (API consumer identity)
              3. Client IP address       → "ip:<host>"   (network-level fallback)

            NOTE (#493): Rate limiting is keyed on owner_id/X-Owner-Id when
            present, ensuring different API consumers get separate buckets.
            Pure IP-based keying is only used when no owner identity is provided —
            e.g., unauthenticated or probe traffic.

            TODO (#493): If a dedicated API key header (X-Api-Key) is introduced
            in the future, add it as a resolution step between owner_id and IP
            to enable per-key rate limiting independent of owner_id.
            """
            # 1. owner_id query param
            owner_id = request.query_params.get("owner_id")
            if owner_id:
                return f"owner:{owner_id}"

            # 2. X-Owner-Id header
            header_owner = request.headers.get("x-owner-id")
            if header_owner:
                return f"owner:{header_owner}"

            # 3. Client IP
            client = request.client
            if client is not None:
                return f"ip:{client.host}"

            return "ip:unknown"

        async def dispatch(self, request: Request, call_next: Callable) -> Response:
            if not self._enabled:
                return await call_next(request)

            # Skip health/docs endpoints
            if request.url.path in _SKIP_PATHS:
                return await call_next(request)

            identity = self._get_identity(request)
            allowed, retry_after = self._store.consume(identity)

            if not allowed:
                retry_int = math.ceil(retry_after)
                return JSONResponse(
                    status_code=429,
                    content={
                        "detail": "Rate limit exceeded",
                        "retry_after": retry_int,
                    },
                    headers={
                        "Retry-After": str(retry_int),
                        "X-RateLimit-Limit": str(self._store.rpm),
                        "X-RateLimit-Remaining": "0",
                    },
                )

            response = await call_next(request)
            # Inject informational headers on successful responses
            response.headers["X-RateLimit-Limit"] = str(self._store.rpm)
            return response

except ImportError:
    # Starlette not available — provide a no-op stub
    class RateLimiterMiddleware:  # type: ignore[no-redef]
        def __init__(self, app: Any, **kwargs: Any) -> None:
            self.app = app
            self._store = TokenBucketStore(
                requests_per_minute=kwargs.get("requests_per_minute", 60),
                burst=kwargs.get("burst", 10),
            )

        @property
        def store(self) -> TokenBucketStore:
            return self._store

        async def __call__(self, scope: Any, receive: Any, send: Any) -> None:
            await self.app(scope, receive, send)
