"""Rate limiting middleware for LLMHosts proxy.

Provides a token-bucket rate limiter keyed by API key or client IP.
Middleware is added to the FastAPI app in :func:`llmhost.proxy.app.create_app`.
"""

from __future__ import annotations

import logging
import time
from typing import TYPE_CHECKING

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

if TYPE_CHECKING:
    from starlette.requests import Request
    from starlette.responses import Response

from llmhosts.plans import get_plan_limits

logger = logging.getLogger(__name__)

# Paths that are never rate-limited.
_EXEMPT_PATHS: frozenset[str] = frozenset({"/health"})

# Path prefixes that are never rate-limited (management endpoints).
_EXEMPT_PREFIXES: tuple[str, ...] = ("/api/keys",)

# Remove idle buckets after this many seconds.
_BUCKET_IDLE_TTL: float = 3600.0  # 1 hour

# Run cleanup every N requests.
_CLEANUP_INTERVAL: int = 100


class TokenBucket:
    """Simple token-bucket rate limiter.

    Parameters
    ----------
    rate:
        Tokens added per second.
    burst:
        Maximum number of tokens the bucket can hold.
    """

    def __init__(self, rate: float, burst: int) -> None:
        self.rate = rate
        self.burst = burst
        self.tokens: float = float(burst)
        self.last_refill: float = time.monotonic()
        self.last_used: float = time.monotonic()

    def consume(self) -> tuple[bool, float]:
        """Try to consume one token.

        Returns
        -------
        tuple[bool, float]
            ``(allowed, retry_after_seconds)``.  When *allowed* is ``True``,
            *retry_after_seconds* is ``0.0``.
        """
        now = time.monotonic()

        # Refill tokens based on elapsed time.
        elapsed = now - self.last_refill
        self.tokens = min(float(self.burst), self.tokens + elapsed * self.rate)
        self.last_refill = now

        if self.tokens >= 1.0:
            self.tokens -= 1.0
            self.last_used = now
            return True, 0.0

        # Not enough tokens -- calculate how long before one token is available.
        retry_after = (1.0 - self.tokens) / self.rate if self.rate > 0 else 0.0
        self.last_used = now
        return False, retry_after


class RateLimitMiddleware(BaseHTTPMiddleware):
    """Per-key / per-IP token-bucket rate limiter.

    Parameters
    ----------
    app:
        The ASGI application.
    rpm:
        Maximum requests per minute (converted to tokens/sec internally).
    burst:
        Maximum burst size (number of requests allowed in a rapid burst).
    """

    def __init__(self, app: object, rpm: int = 60, burst: int = 10) -> None:
        super().__init__(app)  # type: ignore[arg-type]
        self.rpm = rpm
        self.rate: float = rpm / 60.0  # tokens per second
        self.burst = burst
        self._buckets: dict[str, TokenBucket] = {}
        self._request_count: int = 0

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _get_bucket(self, key: str, plan: str = "free") -> tuple[TokenBucket, int]:
        """Return ``(bucket, effective_rpm)`` for *key*, creating one if needed.

        When the plan's RPM differs from the middleware's default, a
        plan-specific rate and burst are used.  Otherwise the middleware's
        configured ``rpm`` / ``burst`` are preserved (important for tests
        and custom configurations).
        """
        plan_limits = get_plan_limits(plan)
        plan_rpm = plan_limits.rpm if plan_limits.rpm > 0 else 999_999

        # Use configured defaults when plan matches, plan-specific otherwise.
        if plan_rpm == self.rpm:
            effective_rate = self.rate
            effective_burst = self.burst
        else:
            effective_rate = plan_rpm / 60.0
            effective_burst = max(10, plan_rpm // 6)  # ~10s burst window

        bucket = self._buckets.get(key)
        # Re-create if the plan (and thus rate) has changed.
        if bucket is None or bucket.rate != effective_rate:
            bucket = TokenBucket(rate=effective_rate, burst=effective_burst)
            self._buckets[key] = bucket
        return bucket, plan_rpm

    def _cleanup_idle_buckets(self) -> None:
        """Remove buckets that have not been used for ``_BUCKET_IDLE_TTL``."""
        now = time.monotonic()
        stale_keys = [k for k, b in self._buckets.items() if (now - b.last_used) > _BUCKET_IDLE_TTL]
        for key in stale_keys:
            del self._buckets[key]
        if stale_keys:
            logger.debug("Cleaned up %d idle rate-limit buckets", len(stale_keys))

    @staticmethod
    def _client_key(request: Request) -> str:
        """Derive the rate-limit key from the request.

        Uses ``request.state.api_key_id`` (set by auth middleware) when
        available, otherwise falls back to the client IP address.
        """
        api_key_id: str | None = getattr(request.state, "api_key_id", None)
        if api_key_id:
            return f"key:{api_key_id}"
        host = request.client.host if request.client else "unknown"
        return f"ip:{host}"

    # ------------------------------------------------------------------
    # Dispatch
    # ------------------------------------------------------------------

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        path = request.url.path

        # Exempt paths bypass rate limiting entirely.
        if path in _EXEMPT_PATHS or path.startswith(_EXEMPT_PREFIXES):
            return await call_next(request)

        # Periodic cleanup of idle buckets.
        self._request_count += 1
        if self._request_count % _CLEANUP_INTERVAL == 0:
            self._cleanup_idle_buckets()

        key = self._client_key(request)
        plan = getattr(request.state, "plan", "free") if hasattr(request, "state") else "free"
        bucket, effective_rpm = self._get_bucket(key, plan)
        allowed, retry_after = bucket.consume()

        if not allowed:
            retry_after_int = max(1, int(retry_after + 0.5))
            reset_timestamp = int(time.time() + retry_after)
            logger.warning("Rate limit exceeded for %s plan=%s (retry_after=%.1fs)", key, plan, retry_after)
            return JSONResponse(
                status_code=429,
                content={
                    "error": {
                        "message": "Rate limit exceeded",
                        "type": "rate_limit_error",
                        "code": "rate_limit_exceeded",
                    }
                },
                headers={
                    "Retry-After": str(retry_after_int),
                    "X-RateLimit-Limit": str(effective_rpm),
                    "X-RateLimit-Remaining": "0",
                    "X-RateLimit-Reset": str(reset_timestamp),
                },
            )

        response = await call_next(request)

        # Attach rate-limit headers to every successful response.
        remaining = max(0, int(bucket.tokens))
        effective_rate = bucket.rate
        reset_timestamp = (
            int(time.time() + (bucket.burst - bucket.tokens) / effective_rate) if effective_rate > 0 else 0
        )
        response.headers["X-RateLimit-Limit"] = str(effective_rpm)
        response.headers["X-RateLimit-Remaining"] = str(remaining)
        response.headers["X-RateLimit-Reset"] = str(reset_timestamp)

        return response
