"""Rate limiting middleware for API protection."""

import time
import logging
from collections import defaultdict
from typing import Optional, Dict, Any
from pydantic import BaseModel
from fastapi import Request, HTTPException
from starlette.middleware.base import BaseHTTPMiddleware

try:
    import redis.asyncio as aioredis
    REDIS_AVAILABLE = True
except ImportError:
    REDIS_AVAILABLE = False

LOGGER = logging.getLogger(__name__)


class RateLimitConfig(BaseModel):
    """Configuration for rate limiting."""
    enabled: bool = True
    requests_per_minute: int = 60
    burst_size: int = 10
    redis_url: Optional[str] = None


class RateLimitMiddleware(BaseHTTPMiddleware):
    """Token bucket rate limiting middleware.
    
    Implements rate limiting with configurable backend (memory or Redis).
    Uses token bucket algorithm for smooth rate control with burst capacity.
    """
    
    def __init__(
        self,
        app,
        config: RateLimitConfig,
    ):
        """Initialize rate limit middleware.
        
        Args:
            app: FastAPI application
            config: Rate limit configuration
        """
        super().__init__(app)
        self.config = config
        
        if config.redis_url and REDIS_AVAILABLE:
            self.redis = aioredis.from_url(config.redis_url)
            self.backend = "redis"
            LOGGER.info("Rate limiting using Redis backend")
        else:
            if config.redis_url and not REDIS_AVAILABLE:
                LOGGER.warning(
                    "Redis URL provided but redis package not available. "
                    "Falling back to in-memory rate limiting."
                )
            self.buckets: Dict[str, Dict[str, Any]] = defaultdict(
                lambda: {"tokens": config.burst_size, "last_update": time.time()}
            )
            self.backend = "memory"
            LOGGER.info("Rate limiting using in-memory backend")
    
    async def dispatch(self, request: Request, call_next):
        """Process request with rate limiting.
        
        Args:
            request: Incoming HTTP request
            call_next: Next middleware in chain
            
        Returns:
            HTTP response with rate limit headers
            
        Raises:
            HTTPException: 429 if rate limit exceeded
        """
        if not self.config.enabled:
            return await call_next(request)
        
        client_id = self._get_client_id(request)
        
        # Check rate limit
        allowed = await self._check_rate_limit(client_id)
        if not allowed:
            LOGGER.warning(
                "Rate limit exceeded",
                extra={"context": {"client_id": client_id}},
            )
            raise HTTPException(
                status_code=429,
                detail="Rate limit exceeded. Try again later.",
                headers={"Retry-After": "60"},
            )
        
        # Process request
        response = await call_next(request)
        
        # Add rate limit headers
        response.headers["X-RateLimit-Limit"] = str(self.config.requests_per_minute)
        remaining = await self._get_remaining(client_id)
        response.headers["X-RateLimit-Remaining"] = str(remaining)
        
        return response
    
    def _get_client_id(self, request: Request) -> str:
        """Extract client identifier for rate limiting.
        
        Args:
            request: HTTP request
            
        Returns:
            Client identifier (user ID or IP address)
        """
        # Check for authenticated user context
        if hasattr(request.state, "user"):
            return f"user:{request.state.user.id}"
        
        # Fall back to IP address
        client_host = request.client.host if request.client else "unknown"
        return f"ip:{client_host}"
    
    async def _check_rate_limit(self, client_id: str) -> bool:
        """Check if request is allowed under rate limit.
        
        Args:
            client_id: Client identifier
            
        Returns:
            True if request is allowed, False otherwise
        """
        if self.backend == "redis":
            return await self._check_rate_limit_redis(client_id)
        return await self._check_rate_limit_memory(client_id)
    
    async def _check_rate_limit_memory(self, client_id: str) -> bool:
        """Memory-backed token bucket rate limiting.
        
        Args:
            client_id: Client identifier
            
        Returns:
            True if request is allowed, False otherwise
        """
        bucket = self.buckets[client_id]
        now = time.time()
        
        # Refill tokens based on time elapsed
        elapsed = now - bucket["last_update"]
        refill = (elapsed * self.config.requests_per_minute) / 60.0
        bucket["tokens"] = min(self.config.burst_size, bucket["tokens"] + refill)
        bucket["last_update"] = now
        
        # Check if request can proceed
        if bucket["tokens"] >= 1:
            bucket["tokens"] -= 1
            return True
        return False
    
    async def _check_rate_limit_redis(self, client_id: str) -> bool:
        """Redis-backed rate limiting using sliding window.
        
        Args:
            client_id: Client identifier
            
        Returns:
            True if request is allowed, False otherwise
        """
        key = f"ratelimit:{client_id}"
        now = int(time.time())
        window = 60  # seconds
        
        pipe = self.redis.pipeline()
        pipe.zadd(key, {str(now): now})
        pipe.zremrangebyscore(key, 0, now - window)
        pipe.zcard(key)
        pipe.expire(key, window + 1)
        results = await pipe.execute()
        
        request_count = results[2]
        return request_count <= self.config.requests_per_minute
    
    async def _get_remaining(self, client_id: str) -> int:
        """Get remaining requests for client.
        
        Args:
            client_id: Client identifier
            
        Returns:
            Number of remaining requests in current window
        """
        if self.backend == "memory":
            bucket = self.buckets.get(client_id)
            if bucket:
                return int(bucket["tokens"])
            return self.config.burst_size
        
        # For Redis, calculate from sliding window
        key = f"ratelimit:{client_id}"
        now = int(time.time())
        window = 60
        
        count = await self.redis.zcount(key, now - window, now)
        return max(0, self.config.requests_per_minute - count)
