"""Adaptive rate limiting with sliding window and per-model weight multipliers.

Replaces the token bucket rate limiter with a more accurate sliding window
algorithm that accounts for model cost differences.  Heavy models (GPT-4,
Claude Opus) consume more of a client's budget per request than lightweight
ones (GPT-3.5, Haiku), giving fairer resource allocation.

This module is part of the TrustShield security platform (Pillar 3: AdaptiveShield).
"""

from __future__ import annotations

import json
import logging
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable

from pydantic import BaseModel, Field
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

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


class RateLimitConfig(BaseModel):
    """Rate limiting configuration.

    Attributes
    ----------
    enabled:
        Master switch -- when ``False`` the middleware is a no-op pass-through.
    algorithm:
        Algorithm identifier (currently only ``sliding_window`` is implemented).
    window_seconds:
        Width of the sliding window in seconds.
    burst_multiplier:
        Reserved for future burst-detection logic.
    burst_window_seconds:
        Reserved for future burst-detection logic.
    model_weights:
        Per-model cost multiplier.  A request for a model with weight 4 counts
        as 4 units towards the client's budget; weight 1 counts as 1 unit.
    default_weight:
        Weight assigned to models not listed in *model_weights*.
    cleanup_interval_seconds:
        Entries idle longer than this are eligible for eviction.
    """

    enabled: bool = True
    algorithm: str = "sliding_window"
    window_seconds: int = 60
    burst_multiplier: float = 1.5
    burst_window_seconds: int = 10
    model_weights: dict[str, int] = Field(
        default_factory=lambda: {
            "gpt-4": 4,
            "gpt-4-turbo": 3,
            "gpt-4o": 2,
            "gpt-3.5-turbo": 1,
            "claude-3-opus": 4,
            "claude-3-sonnet": 2,
            "claude-3-haiku": 1,
            "llama-3-70b": 2,
            "llama-3-8b": 1,
        },
    )
    default_weight: int = 1
    cleanup_interval_seconds: float = 3600.0


# ---------------------------------------------------------------------------
# Result
# ---------------------------------------------------------------------------


@dataclass(slots=True)
class RateLimitResult:
    """Outcome of a single rate-limit evaluation."""

    allowed: bool
    limit: int
    remaining: int
    reset_at: float  # unix-epoch timestamp (monotonic clock based)
    weight: int
    retry_after: float | None = None  # seconds until next allowed request


# ---------------------------------------------------------------------------
# Backend protocol
# ---------------------------------------------------------------------------


@runtime_checkable
class RateLimitBackend(Protocol):
    """Protocol for rate-limit storage backends (in-memory, Redis, etc.)."""

    def consume(self, key: str, weight: int, limit: int, window: int) -> RateLimitResult: ...

    def get_remaining(self, key: str, limit: int, window: int) -> int: ...

    def cleanup(self) -> int: ...  # returns number of cleaned entries


# ---------------------------------------------------------------------------
# In-memory sliding window backend
# ---------------------------------------------------------------------------


@dataclass
class SlidingWindowEntry:
    """Tracks weighted requests within a sliding window for a single key."""

    requests: list[tuple[float, int]] = field(default_factory=list)  # (timestamp, weight)
    last_accessed: float = field(default_factory=time.monotonic)


class InMemoryBackend:
    """In-memory sliding window rate-limit backend.

    Uses a simple list of ``(timestamp, weight)`` tuples per key.  Expired
    entries (older than the window) are pruned on every ``consume`` call.
    """

    def __init__(self) -> None:
        self._buckets: dict[str, SlidingWindowEntry] = defaultdict(SlidingWindowEntry)
        self._last_cleanup: float = time.monotonic()

    def consume(self, key: str, weight: int, limit: int, window: int) -> RateLimitResult:
        """Attempt to consume *weight* units from *key*'s budget.

        Returns a :class:`RateLimitResult` indicating whether the request is
        allowed and, if not, how long the caller should wait.
        """
        now = time.monotonic()
        entry = self._buckets[key]
        entry.last_accessed = now

        # Prune expired requests (outside the current window).
        cutoff = now - window
        entry.requests = [(ts, w) for ts, w in entry.requests if ts > cutoff]

        # Calculate current usage within the window.
        current_usage = sum(w for _, w in entry.requests)

        if current_usage + weight > limit:
            retry_after = self._calculate_retry_after(entry.requests, weight, limit, window, now)
            return RateLimitResult(
                allowed=False,
                limit=limit,
                remaining=0,
                reset_at=now + window,
                weight=weight,
                retry_after=retry_after,
            )

        # Allow the request and record it.
        entry.requests.append((now, weight))
        return RateLimitResult(
            allowed=True,
            limit=limit,
            remaining=max(0, limit - current_usage - weight),
            reset_at=now + window,
            weight=weight,
        )

    def get_remaining(self, key: str, limit: int, window: int) -> int:
        """Return how many units *key* can still consume within the window."""
        now = time.monotonic()
        entry = self._buckets.get(key)
        if not entry:
            return limit
        cutoff = now - window
        current_usage = sum(w for ts, w in entry.requests if ts > cutoff)
        return max(0, limit - current_usage)

    def cleanup(self) -> int:
        """Remove entries idle for more than 1 hour.  Returns evicted count."""
        now = time.monotonic()
        idle_keys = [k for k, v in self._buckets.items() if now - v.last_accessed > 3600]
        for k in idle_keys:
            del self._buckets[k]
        if idle_keys:
            logger.debug("Cleaned up %d idle rate-limit entries", len(idle_keys))
        return len(idle_keys)

    # -- helpers --

    @staticmethod
    def _calculate_retry_after(
        requests: list[tuple[float, int]],
        weight: int,
        limit: int,
        window: int,
        now: float,
    ) -> float:
        """Calculate seconds until enough capacity frees up for *weight* units.

        Walks the request list chronologically: as each old request expires
        (leaves the window), its weight is freed.  We return the time at which
        the cumulative freed weight is enough to fit the new request.
        """
        if not requests:
            return 0.0

        current_usage = sum(w for _, w in requests)
        needed = current_usage + weight - limit

        if needed <= 0:
            return 0.0

        freed = 0
        for ts, w in sorted(requests):
            freed += w
            if freed >= needed:
                # This request expires at ``ts + window``.
                return max(0.0, (ts + window) - now)

        # All requests expiring still isn't enough -- full window wait.
        return float(window)


# ---------------------------------------------------------------------------
# Middleware
# ---------------------------------------------------------------------------

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

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

# Plan-based limits (weighted units per window).
_PLAN_LIMITS: dict[str, int] = {
    "free": 30,
    "pro": 300,
    "team": 1000,
    "enterprise": 5000,
}


class AdaptiveRateLimitMiddleware(BaseHTTPMiddleware):
    """Sliding window rate limiter with per-model weight multipliers.

    Parameters
    ----------
    app:
        The ASGI application.
    config:
        Optional :class:`RateLimitConfig`.  Uses sensible defaults when omitted.
    """

    def __init__(self, app: Any, config: RateLimitConfig | None = None) -> None:
        super().__init__(app)
        self.config = config or RateLimitConfig()
        self._backend = InMemoryBackend()
        self._request_count: int = 0

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        if not self.config.enabled:
            return await call_next(request)

        path = request.url.path

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

        # Determine rate-limit key (API key ID or client IP).
        key = self._get_rate_limit_key(request)

        # Determine model weight from request body (for POST /v1/* routes).
        weight = await self._get_model_weight(request)

        # Resolve plan-based budget.
        limit = self._get_plan_limit(request)

        # Periodic cleanup of idle entries.
        self._request_count += 1
        if self._request_count % 100 == 0:
            self._backend.cleanup()

        result = self._backend.consume(key, weight, limit, self.config.window_seconds)

        if not result.allowed:
            logger.warning(
                "Rate limit exceeded for %s weight=%d limit=%d retry_after=%.1fs",
                key,
                weight,
                limit,
                result.retry_after or 0.0,
            )
            return self._rate_limit_response(result)

        # Pass through to the actual handler and attach headers.
        response = await call_next(request)
        self._add_headers(response, result)
        return response

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

    @staticmethod
    def _get_rate_limit_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}"
        client_ip = request.client.host if request.client else "unknown"
        return f"ip:{client_ip}"

    async def _get_model_weight(self, request: Request) -> int:
        """Extract the model name from the request body and return its weight.

        Only inspects bodies on POST requests.  If the body cannot be read or
        parsed, returns the configured default weight (never raises).
        """
        if request.method != "POST":
            return self.config.default_weight

        try:
            body = await request.body()
            if body:
                data = json.loads(body)
                model = data.get("model", "")
                return self.config.model_weights.get(model, self.config.default_weight)
        except (json.JSONDecodeError, UnicodeDecodeError, AttributeError):
            logger.debug("Could not parse request body for model weight extraction")
        except Exception:
            logger.debug("Unexpected error reading request body for model weight", exc_info=True)
        return self.config.default_weight

    @staticmethod
    def _get_plan_limit(request: Request) -> int:
        """Resolve the rate-limit budget based on the user's plan."""
        plan: str = getattr(request.state, "plan", "free")
        return _PLAN_LIMITS.get(plan, _PLAN_LIMITS["free"])

    @staticmethod
    def _rate_limit_response(result: RateLimitResult) -> JSONResponse:
        """Build a 429 Too Many Requests response with standard headers."""
        retry_after_secs = int(result.retry_after or 60) + 1
        headers = {
            "X-RateLimit-Limit": str(result.limit),
            "X-RateLimit-Remaining": "0",
            "X-RateLimit-Reset": str(int(result.reset_at)),
            "X-RateLimit-Model-Weight": str(result.weight),
            "Retry-After": str(retry_after_secs),
        }
        return JSONResponse(
            status_code=429,
            content={
                "error": "rate_limit_exceeded",
                "message": (
                    f"Rate limit exceeded. This request costs {result.weight} unit(s). "
                    f"Try again in {retry_after_secs}s."
                ),
                "limit": result.limit,
                "weight": result.weight,
                "retry_after": retry_after_secs,
            },
            headers=headers,
        )

    @staticmethod
    def _add_headers(response: Response, result: RateLimitResult) -> None:
        """Attach rate-limit metadata headers to a successful response."""
        response.headers["X-RateLimit-Limit"] = str(result.limit)
        response.headers["X-RateLimit-Remaining"] = str(result.remaining)
        response.headers["X-RateLimit-Reset"] = str(int(result.reset_at))
        response.headers["X-RateLimit-Model-Weight"] = str(result.weight)
