"""Redis-based distributed rate limiting with in-memory fallback."""

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

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


class RedisRateLimitMiddleware(BaseHTTPMiddleware):
    """Distributed rate limiting using Redis token bucket.

    Falls back to in-memory rate limiting if Redis is unavailable.
    Supports per-tenant rate limits via request.state.tenant.

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

    The middleware will automatically use Redis from app.state.redis
    if available, otherwise falls back to in-memory.
    """

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

        Args:
            app: The FastAPI application.
            requests_per_minute: Default sustained request rate.
            burst_size: Maximum burst of requests allowed.
            exclude_paths: Paths to exclude from rate limiting.
            key_prefix: Redis key prefix for rate limit buckets.
        """
        super().__init__(app)
        self.requests_per_minute = requests_per_minute
        self.burst_size = burst_size
        self.exclude_paths = exclude_paths or ["/health", "/ready"]
        self.key_prefix = key_prefix

        # Fallback in-memory bucket for when Redis unavailable
        self._fallback_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
        tenant = getattr(request.state, "tenant", None)
        rate_limit = tenant.rate_limit_per_minute if tenant else self.requests_per_minute

        # Try Redis, fall back to in-memory
        redis = getattr(request.app.state, "redis", None)
        allowed = await self._check_redis(redis, client_id, rate_limit) if redis else None
        if allowed is None:
            allowed = self._check_memory(client_id, rate_limit)

        if not allowed:
            return JSONResponse(
                status_code=429,
                content={
                    "error": "Rate limit exceeded",
                    "retry_after_seconds": 60 / rate_limit,
                },
                headers={"Retry-After": str(int(60 / rate_limit))},
            )

        # Add rate limit headers to response
        response = await call_next(request)
        response.headers["X-RateLimit-Limit"] = str(rate_limit)

        return response

    async def _check_redis(
        self, redis: Any, client_id: str, rate_limit: int
    ) -> bool | None:
        """Check rate limit using Redis token bucket.

        Args:
            redis: Redis client.
            client_id: Client identifier.
            rate_limit: Rate limit (requests per minute).

        Returns:
            True if allowed, False if rate limited, None if Redis unavailable.
        """
        try:
            key = f"{self.key_prefix}{client_id}"
            now = time.time()

            # Lua script for atomic token bucket operation
            lua_script = """
            local key = KEYS[1]
            local rate = tonumber(ARGV[1])
            local burst = tonumber(ARGV[2])
            local now = tonumber(ARGV[3])

            local bucket = redis.call('HMGET', key, 'tokens', 'last_update')
            local tokens = tonumber(bucket[1]) or burst
            local last_update = tonumber(bucket[2]) or now

            -- Refill tokens based on time passed
            local time_passed = now - last_update
            tokens = math.min(burst, tokens + (time_passed * rate / 60))

            if tokens >= 1 then
                tokens = tokens - 1
                redis.call('HMSET', key, 'tokens', tokens, 'last_update', now)
                redis.call('EXPIRE', key, 120)  -- 2 minute TTL
                return 1
            else
                redis.call('HMSET', key, 'tokens', tokens, 'last_update', now)
                redis.call('EXPIRE', key, 120)
                return 0
            end
            """

            result = await redis.eval(
                lua_script, 1, key, rate_limit, self.burst_size, now
            )
            return bool(result)

        except Exception:
            # Redis unavailable, fall back to in-memory
            return None

    def _check_memory(self, client_id: str, rate_limit: int) -> bool:
        """Fallback in-memory rate limiting.

        Args:
            client_id: Client identifier.
            rate_limit: Rate limit (requests per minute).

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

        # Refill tokens based on time passed
        time_passed = now - bucket["last_update"]
        tokens_to_add = time_passed * (rate_limit / 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

    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}"
