"""IP-based rate limiting middleware for sensitive routes.

This provides additional rate limiting on top of the plan-based rate limiting in
llmhost.proxy.ratelimit. This middleware is IP-based and protects against abuse
even from unauthenticated requests.
"""

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

logger = logging.getLogger(__name__)

# Route patterns and their IP-based rate limits (requests per minute)
_ROUTE_LIMITS: dict[str, int] = {
    "/api/auth/": 10,  # Auth endpoints: 10 req/min per IP
    "/api/keys/": 20,  # Key management: 20 req/min per IP (user-based limit also applies)
    "/v1/chat/completions": 100,  # Chat completions: 100 req/min per IP (backup to plan limits)
    "/v1/messages": 100,  # Anthropic messages: 100 req/min per IP
}

# 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 SlidingWindowBucket:
    """Sliding window rate limiter.

    Tracks timestamps of requests in a sliding window. More memory-intensive
    than token bucket, but provides precise per-minute limits.

    Parameters
    ----------
    limit:
        Maximum number of requests allowed per minute.
    """

    def __init__(self, limit: int) -> None:
        self.limit = limit
        self.window_seconds: float = 60.0
        self.requests: list[float] = []
        self.last_used: float = time.monotonic()

    def is_allowed(self) -> tuple[bool, int]:
        """Check if a request is allowed.

        Returns
        -------
        tuple[bool, int]
            ``(allowed, retry_after_seconds)``. When *allowed* is ``True``,
            *retry_after_seconds* is ``0``.
        """
        now = time.time()
        self.last_used = time.monotonic()

        # Remove requests outside the sliding window
        cutoff = now - self.window_seconds
        self.requests = [ts for ts in self.requests if ts > cutoff]

        if len(self.requests) < self.limit:
            self.requests.append(now)
            return True, 0

        # Calculate retry-after based on oldest request in window
        oldest = self.requests[0]
        retry_after = int(oldest + self.window_seconds - now) + 1
        return False, max(1, retry_after)


class IPRateLimitMiddleware(BaseHTTPMiddleware):
    """IP-based rate limiting for sensitive routes.

    This middleware applies stricter per-IP limits to specific routes to prevent
    abuse. It works alongside the plan-based rate limiting (which is per-key).

    Parameters
    ----------
    app:
        The ASGI application.
    """

    def __init__(self, app: object) -> None:
        super().__init__(app)  # type: ignore[arg-type]
        self._buckets: dict[str, SlidingWindowBucket] = {}
        self._request_count: int = 0

    def _get_bucket(self, ip: str, limit: int) -> SlidingWindowBucket:
        """Return bucket for *ip*, creating one if needed."""
        key = f"{ip}:{limit}"
        bucket = self._buckets.get(key)
        if bucket is None:
            bucket = SlidingWindowBucket(limit=limit)
            self._buckets[key] = bucket
        return bucket

    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("IPRateLimit: cleaned up %d idle buckets", len(stale_keys))

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

        # Check if this path requires IP-based rate limiting
        limit: int | None = None
        for pattern, rpm in _ROUTE_LIMITS.items():
            if path.startswith(pattern):
                limit = rpm
                break

        if limit is None:
            # No IP-based rate limiting for this route
            return await call_next(request)

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

        # Get client IP
        ip = request.client.host if request.client else "unknown"
        bucket = self._get_bucket(ip, limit)
        allowed, retry_after = bucket.is_allowed()

        if not allowed:
            logger.warning("IP rate limit exceeded for %s on %s (limit=%d/min)", ip, path, limit)
            return JSONResponse(
                status_code=429,
                content={
                    "error": {
                        "message": "Too many requests from your IP address. Please try again later.",
                        "type": "rate_limit_error",
                        "code": "ip_rate_limit_exceeded",
                    }
                },
                headers={
                    "Retry-After": str(retry_after),
                    "X-RateLimit-Limit-IP": str(limit),
                },
            )

        return await call_next(request)
