"""Rate limiting middleware."""

import time
from collections import defaultdict
from collections.abc import Callable

from fastapi import Request, Response
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware


class RateLimitMiddleware(BaseHTTPMiddleware):
    """Simple in-memory rate limiting middleware.

    For production, use Redis-based rate limiting for
    distributed deployments.

    Example usage:
        ```python
        app.add_middleware(
            RateLimitMiddleware,
            requests_per_minute=60,
            burst_size=10,
        )
        ```
    """

    def __init__(
        self,
        app: Callable,
        requests_per_minute: int = 60,
        burst_size: int = 10,
        exclude_paths: list[str] | None = None,
    ) -> None:
        """Initialize rate limiter.

        Args:
            app: The FastAPI application.
            requests_per_minute: Sustained request rate.
            burst_size: Maximum burst of requests allowed.
            exclude_paths: Paths to exclude from rate limiting.
        """
        super().__init__(app)
        self.requests_per_minute = requests_per_minute
        self.burst_size = burst_size
        self.exclude_paths = exclude_paths or ["/health", "/ready"]

        # Token bucket state per client
        self._buckets: dict[str, dict] = defaultdict(
            lambda: {"tokens": burst_size, "last_update": time.time()}
        )

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        """Process request through rate limiter.

        Args:
            request: The incoming request.
            call_next: The next middleware/handler.

        Returns:
            Response or 429 if rate limited.
        """
        # Skip rate limiting for excluded paths
        if any(request.url.path.startswith(path) for path in self.exclude_paths):
            return await call_next(request)

        # Get client identifier
        client_id = self._get_client_id(request)

        # Get tenant-specific rate limit if available (set by AuthMiddleware)
        tenant = getattr(request.state, "tenant", None)
        tenant_rate = tenant.rate_limit_per_minute if tenant else None
        effective_rate = tenant_rate or self.requests_per_minute

        # Check rate limit
        if not self._allow_request(client_id, tenant_rate):
            return JSONResponse(
                status_code=429,
                content={
                    "error": "Rate limit exceeded",
                    "retry_after_seconds": 60 / effective_rate,
                },
                headers={"Retry-After": str(int(60 / effective_rate))},
            )

        # Add rate limit headers to response
        response = await call_next(request)

        bucket = self._buckets[client_id]
        response.headers["X-RateLimit-Limit"] = str(effective_rate)
        response.headers["X-RateLimit-Remaining"] = str(int(bucket["tokens"]))

        return response

    def _get_client_id(self, request: Request) -> str:
        """Get unique client identifier.

        Uses user ID if authenticated, otherwise IP address.

        Args:
            request: The incoming request.

        Returns:
            Client identifier string.
        """
        # Try to get user ID from auth
        user_id = getattr(request.state, "user_id", None)
        if user_id:
            return f"user:{user_id}"

        # Fall back to IP address
        forwarded = request.headers.get("X-Forwarded-For")
        if forwarded:
            return f"ip:{forwarded.split(',')[0].strip()}"

        client_host = request.client.host if request.client else "unknown"
        return f"ip:{client_host}"

    def _allow_request(self, client_id: str, rate_limit: int | None = None) -> bool:
        """Check if request should be allowed using token bucket.

        Args:
            client_id: The client identifier.
            rate_limit: Custom rate limit (requests per minute). Uses default if None.

        Returns:
            True if request is allowed, False if rate limited.
        """
        bucket = self._buckets[client_id]
        now = time.time()

        # Use custom rate limit or default
        effective_rate = rate_limit or self.requests_per_minute

        # Refill tokens based on time passed
        time_passed = now - bucket["last_update"]
        tokens_to_add = time_passed * (effective_rate / 60)
        bucket["tokens"] = min(self.burst_size, bucket["tokens"] + tokens_to_add)
        bucket["last_update"] = now

        # Check if we have tokens available
        if bucket["tokens"] >= 1:
            bucket["tokens"] -= 1
            return True

        return False
