""" 比率 limiting middleware for the Pixel API.

Provides:
- Redis-based rate limiting per user/IP
- Configurable rate limits per endpoint
- Sliding window algorithm
- Rate limit headers in responses
""" 
import logging
import time
from collections import defaultdict, deque
from typing import Callable, Optional, Dict, Deque
from fastapi import Request, status
from fastapi.responses import JSONResponse
from redis import asyncio as aioredis
from redis.exceptions import RedisError

from src.utils.errors import ErrorCode

logger = logging.getLogger(__name__)


class RateLimiter:
    """
    Redis-based rate limiter using sliding window algorithm.
    
    Tracks request counts per user/IP within a time window and
    enforces configurable rate limits.
    """
    
    def __init__(self, redis_url: str = "redis://localhost:6379"):
        """ 初始化 the rate limiter.
        
        Args:
            redis_url: Redis connection URL
        """
        self.redis_url = redis_url
        self._redis: Optional[aioredis.Redis] = None
        self._connected = False
        
        # 默认 rate limit configurations (requests per window)
        # 格式化: {endpoint_pattern: (limit, window_seconds)}
        self.rate_limits: Dict[str, tuple[int, int]] = {
            "default": (100, 60),  # 100 requests per minute
            "/generations/image": (10, 60),  # 10 image generations per minute
            "/generations/video": (5, 60),  # 5 video generations per minute
            "/chat": (30, 60),  # 30 chat messages per minute
            "/projects": (50, 60),  # 50 project operations per minute
        }
        # Local fallback only applies to costly endpoints when Redis is unavailable.
        self.local_fallback_prefixes = (
            "/generations/image",
            "/generations/video",
        )
        self._local_windows: Dict[str, Deque[float]] = defaultdict(deque)
    
    async def connect(self):
        """Establish connection to Redis."""
        if not self._connected:
            try:
                self._redis = await aioredis.from_url(
                    self.redis_url,
                    decode_responses=False  # We'll use binary keys
                )
                # Test connection
                await self._redis.ping()
                self._connected = True
                logger.info(f"Rate limiter connected to Redis at {self.redis_url}")
            except RedisError as e:
                logger.error(f"Failed to connect to Redis for rate limiting: {e}")
                self._connected = False
    
    async def disconnect(self):
        """Close connection to Redis."""
        if self._redis:
            await self._redis.close()
            self._connected = False
            logger.info("Rate limiter disconnected from Redis")
    
    def get_rate_limit(self, path: str) -> tuple[int, int]:
        """ 获取 rate limit configuration for a path.
        
        Args:
            path: Request path
            
        Returns:
            Tuple of (limit, window_seconds)
        """
        normalized_path = self._normalize_path(path)

        # Check for exact match first
        if normalized_path in self.rate_limits:
            return self.rate_limits[normalized_path]
        
        # Check for prefix match
        for pattern, config in self.rate_limits.items():
            if pattern != "default" and normalized_path.startswith(pattern):
                return config
        
        # 返回 default
        return self.rate_limits["default"]

    def _canonicalize_path(self, path: str) -> str:
        """Strip version prefix so limits are stable across API mount paths."""
        if path.startswith("/api/v1"):
            trimmed = path[len("/api/v1"):]
            return trimmed or "/"
        return path

    def _normalize_path(self, path: str) -> str:
        """Normalize dynamic path segments to reduce key cardinality."""
        path = self._canonicalize_path(path)
        parts = path.split("/")
        normalized_parts = []

        for part in parts:
            if not part:
                continue
            if self._is_dynamic_segment(part):
                normalized_parts.append("{id}")
            else:
                normalized_parts.append(part)

        return "/" + "/".join(normalized_parts) if normalized_parts else "/"

    def _is_dynamic_segment(self, segment: str) -> bool:
        if not segment:
            return False
        if segment.isdigit():
            return True
        if len(segment) == 36 and segment.count("-") == 4:
            return True
        return False

    def _should_use_local_fallback(self, path: str) -> bool:
        normalized_path = self._normalize_path(path)
        return any(normalized_path.startswith(prefix) for prefix in self.local_fallback_prefixes)

    def _check_local_fallback(
        self,
        identifier: str,
        path: str,
        current_time: float
    ) -> tuple[bool, int, int, int]:
        """
        In-memory fallback limiter used only when Redis is unavailable.
        Returns same tuple format as Redis limiter.
        """
        normalized_path = self._normalize_path(path)
        limit, window = self.get_rate_limit(normalized_path)

        if limit <= 0:
            return False, 0, 0, 0

        key = f"local_rate_limit:{identifier}:{normalized_path}"
        bucket = self._local_windows[key]
        window_start = current_time - window

        while bucket and bucket[0] <= window_start:
            bucket.popleft()

        current_count = len(bucket)
        is_limited = current_count >= limit

        if not is_limited:
            bucket.append(current_time)

        reset_time = int(current_time + window)
        return is_limited, current_count, limit, reset_time
    
    def get_identifier(self, request: Request) -> str:
        """ 获取 unique identifier for rate limiting.
        
        Uses user_id if available, otherwise falls back to IP address.
        
        Args:
            request: FastAPI request object
            
        Returns:
            Unique identifier string
        """
        # to get user_id from request state (set by auth middleware)
        user_id = getattr(request.state, "user_id", None)
        if user_id:
            return f"user:{user_id}"
        
        # Fall back to IP address
        # Check for X-Forwarded-For header (proxy/load balancer)
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            # Take the first IP in the chain
            ip = forwarded_for.split(",")[0].strip()
        else:
            ip = request.client.host if request.client else "unknown"
        
        return f"ip:{ip}"
    
    async def is_rate_limited(
        self,
        identifier: str,
        path: str
    ) -> tuple[bool, int, int, int]:
        """
        Check if request should be rate limited.
        
        Uses sliding window algorithm with Redis.
        
        Args:
            identifier: Unique identifier (user_id or IP)
            path: Request path
            
        Returns:
            Tuple of (is_limited, current_count, limit, reset_time)
        """
        normalized_path = self._normalize_path(path)

        if not self._connected:
            if self._should_use_local_fallback(normalized_path):
                return self._check_local_fallback(identifier, normalized_path, time.time())
            # Redis is not connected, non-critical paths fail open.
            logger.warning("Redis not connected, rate limiting disabled for non-critical endpoints")
            return False, 0, 0, 0
        
        limit, window = self.get_rate_limit(normalized_path)
        now = time.time()
        current_time = int(now)
        window_start = now - window
        
        # Redis key for this identifier and path
        key = f"rate_limit:{identifier}:{normalized_path}"
        
        try:
            # Use Redis pipeline for atomic operations
            pipe = self._redis.pipeline()
            
            # Remove old entries outside the window
            pipe.zremrangebyscore(key, 0, window_start)
            
            # 计数 requests in current window
            pipe.zcard(key)
            
            # Add current request
            pipe.zadd(key, {f"{now}:{time.time_ns()}": now})
            
            # 集合 expiry on the key
            pipe.expire(key, window)
            
            # Execute pipeline
            results = await pipe.execute()
            
            # 获取 count after removing old entries
            current_count = results[1]
            
            # Check if limit exceeded
            is_limited = current_count >= limit
            
            # Calculate reset time (start of next window)
            reset_time = current_time + window
            
            if is_limited:
                logger.warning(
                    f"Rate limit exceeded for {identifier} on {normalized_path}: "
                    f"{current_count}/{limit} requests in {window}s window"
                )
            
            return is_limited, current_count, limit, reset_time
            
        except RedisError as e:
            logger.error(f"Redis error in rate limiting: {e}")
            if self._should_use_local_fallback(normalized_path):
                return self._check_local_fallback(identifier, normalized_path, now)
            # On error, non-critical paths fail open.
            return False, 0, 0, 0
    
    def configure_rate_limit(self, path: str, limit: int, window: int):
        """ 配置ure rate limit for a specific path.
        
        Args:
            path: Request path or pattern
            limit: Maximum number of requests
            window: Time window in seconds
        """
        self.rate_limits[path] = (limit, window)
        logger.info(f"Configured rate limit for {path}: {limit} requests per {window}s")


# 全局的 rate limiter instance
_rate_limiter: Optional[RateLimiter] = None


def get_rate_limiter() -> RateLimiter:
    """ 获取 the global rate limiter instance.
    
    Returns:
        RateLimiter instance
    """ 
    global _rate_limiter
    if _rate_limiter is None:
        from src.config.settings import REDIS_URL
        _rate_limiter = RateLimiter(redis_url=REDIS_URL)
    return _rate_limiter


async def rate_limit_middleware(request: Request, call_next: Callable):
    """ 比率 limiting middleware.
    
    Checks rate limits before processing requests and adds rate limit
    headers to responses.
    """ 
    rate_limiter = get_rate_limiter()
    
    # 获取 identifier and check rate limit
    identifier = rate_limiter.get_identifier(request)
    path = request.url.path
    
    is_limited, current_count, limit, reset_time = await rate_limiter.is_rate_limited(
        identifier, path
    )
    
    if is_limited:
        # 返回 429 Too Many Requests
        remaining = 0
        retry_after = reset_time - int(time.time())
        
        error_response = {
            "error_code": ErrorCode.RATE_LIMIT_EXCEEDED,
            "message": "Rate limit exceeded",
            "details": {
                "limit": limit,
                "window_seconds": rate_limiter.get_rate_limit(path)[1],
                "retry_after": retry_after
            },
            "request_id": getattr(request.state, "request_id", None)
        }
        
        return JSONResponse(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            content=error_response,
            headers={
                "X-RateLimit-Limit": str(limit),
                "X-RateLimit-Remaining": str(remaining),
                "X-RateLimit-Reset": str(reset_time),
                "Retry-After": str(retry_after)
            }
        )
    
    # 进程 request
    response = await call_next(request)
    
    # Add rate limit headers to response
    remaining = max(0, limit - current_count - 1)
    response.headers["X-RateLimit-Limit"] = str(limit)
    response.headers["X-RateLimit-Remaining"] = str(remaining)
    response.headers["X-RateLimit-Reset"] = str(reset_time)
    
    return response


def setup_rate_limiter(app):
    """ 集合up rate limiter middleware for the FastAPI application.
    
    Args:
        app: FastAPI application instance
    """
    # Register middleware
    app.middleware("http")(rate_limit_middleware)
    
    logger.info("Rate limiter middleware registered")


async def init_rate_limiter():
    """ 初始化 rate limiter connection on startup."""
    rate_limiter = get_rate_limiter()
    await rate_limiter.connect()


async def shutdown_rate_limiter():
    """ 清理up rate limiter connection on shutdown."""
    rate_limiter = get_rate_limiter()
    await rate_limiter.disconnect()
