src/middleware/rate_limiting.py

"""
API Rate Limiting Middleware for MediGuard AI.
Implements token bucket and sliding window rate limiting algorithms.
"""

import asyncio
import logging
import time
from collections import deque

import redis.asyncio as redis
from fastapi import HTTPException, Request, status
from starlette.middleware.base import BaseHTTPMiddleware

from src.settings import get_settings

logger = logging.getLogger(__name__)


class RateLimitStrategy:
    """Base class for rate limiting strategies."""

    def is_allowed(self, key: str, limit: int, window: int) -> tuple[bool, dict]:
        """Check if request is allowed.
        
        Returns:
            Tuple of (is_allowed, info_dict)
        """
        raise NotImplementedError


class TokenBucketStrategy(RateLimitStrategy):
    """Token bucket rate limiting algorithm."""

    def __init__(self, redis_client: redis.Redis | None = None):
        self.redis = redis_client
        self.memory_buckets: dict[str, dict] = {}

    async def is_allowed(self, key: str, limit: int, window: int) -> tuple[bool, dict]:
        """Check if request is allowed using token bucket."""
        now = time.time()

        if self.redis:
            return await self._redis_token_bucket(key, limit, window, now)
        else:
            return self._memory_token_bucket(key, limit, window, now)

    async def _redis_token_bucket(self, key: str, limit: int, window: int, now: float) -> tuple[bool, dict]:
        """Token bucket implementation using Redis."""
        bucket_key = f"rate_limit:bucket:{key}"

        # Get current bucket state
        bucket_data = await self.redis.hgetall(bucket_key)

        if bucket_data:
            tokens = float(bucket_data.get('tokens', limit))
            last_refill = float(bucket_data.get('last_refill', now))
        else:
            tokens = limit
            last_refill = now

        # Calculate tokens to add based on time elapsed
        time_elapsed = now - last_refill
        tokens_to_add = time_elapsed * (limit / window)
        tokens = min(limit, tokens + tokens_to_add)

        # Check if request can be processed
        if tokens >= 1:
            tokens -= 1
            await self.redis.hset(bucket_key, mapping={
                'tokens': tokens,
                'last_refill': now
            })
            await self.redis.expire(bucket_key, window * 2)

            return True, {
                'tokens': tokens,
                'limit': limit,
                'window': window,
                'retry_after': 0
            }
        else:
            # Calculate retry after
            retry_after = (1 - tokens) / (limit / window)

            return False, {
                'tokens': tokens,
                'limit': limit,
                'window': window,
                'retry_after': retry_after
            }

    def _memory_token_bucket(self, key: str, limit: int, window: int, now: float) -> tuple[bool, dict]:
        """Token bucket implementation in memory."""
        if key not in self.memory_buckets:
            self.memory_buckets[key] = {
                'tokens': limit,
                'last_refill': now
            }

        bucket = self.memory_buckets[key]

        # Calculate tokens to add
        time_elapsed = now - bucket['last_refill']
        tokens_to_add = time_elapsed * (limit / window)
        bucket['tokens'] = min(limit, bucket['tokens'] + tokens_to_add)
        bucket['last_refill'] = now

        # Check if request can be processed
        if bucket['tokens'] >= 1:
            bucket['tokens'] -= 1
            return True, {
                'tokens': bucket['tokens'],
                'limit': limit,
                'window': window,
                'retry_after': 0
            }
        else:
            retry_after = (1 - bucket['tokens']) / (limit / window)
            return False, {
                'tokens': bucket['tokens'],
                'limit': limit,
                'window': window,
                'retry_after': retry_after
            }


class SlidingWindowStrategy(RateLimitStrategy):
    """Sliding window rate limiting algorithm."""

    def __init__(self, redis_client: redis.Redis | None = None):
        self.redis = redis_client
        self.memory_windows: dict[str, deque] = {}

    async def is_allowed(self, key: str, limit: int, window: int) -> tuple[bool, dict]:
        """Check if request is allowed using sliding window."""
        now = time.time()
        window_start = now - window

        if self.redis:
            return await self._redis_sliding_window(key, limit, window, now, window_start)
        else:
            return self._memory_sliding_window(key, limit, window, now, window_start)

    async def _redis_sliding_window(self, key: str, limit: int, window: int, now: float, window_start: float) -> tuple[bool, dict]:
        """Sliding window implementation using Redis."""
        window_key = f"rate_limit:window:{key}"

        # Remove old entries
        await self.redis.zremrangebyscore(window_key, 0, window_start)

        # Count current requests
        current_count = await self.redis.zcard(window_key)

        if current_count < limit:
            # Add current request
            await self.redis.zadd(window_key, {str(now): now})
            await self.redis.expire(window_key, window)

            return True, {
                'count': current_count + 1,
                'limit': limit,
                'window': window,
                'remaining': limit - current_count - 1,
                'retry_after': 0
            }
        else:
            # Get oldest request time
            oldest = await self.redis.zrange(window_key, 0, 0, withscores=True)
            if oldest:
                retry_after = window - (now - oldest[0][1]) + 1
            else:
                retry_after = window

            return False, {
                'count': current_count,
                'limit': limit,
                'window': window,
                'remaining': 0,
                'retry_after': retry_after
            }

    def _memory_sliding_window(self, key: str, limit: int, window: int, now: float, window_start: float) -> tuple[bool, dict]:
        """Sliding window implementation in memory."""
        if key not in self.memory_windows:
            self.memory_windows[key] = deque()

        request_times = self.memory_windows[key]

        # Remove old requests
        while request_times and request_times[0] < window_start:
            request_times.popleft()

        if len(request_times) < limit:
            request_times.append(now)
            return True, {
                'count': len(request_times),
                'limit': limit,
                'window': window,
                'remaining': limit - len(request_times),
                'retry_after': 0
            }
        else:
            oldest_time = request_times[0]
            retry_after = window - (now - oldest_time) + 1

            return False, {
                'count': len(request_times),
                'limit': limit,
                'window': window,
                'remaining': 0,
                'retry_after': retry_after
            }


class RateLimiter:
    """Main rate limiter class."""

    def __init__(self, strategy: RateLimitStrategy):
        self.strategy = strategy
        self.rules: dict[str, dict] = {}

    def add_rule(self, path_pattern: str, limit: int, window: int, scope: str = "ip"):
        """Add a rate limiting rule."""
        self.rules[path_pattern] = {
            'limit': limit,
            'window': window,
            'scope': scope
        }

    def get_rule(self, path: str) -> dict | None:
        """Get rate limiting rule for a path."""
        # Exact match first
        if path in self.rules:
            return self.rules[path]

        # Pattern matching
        for pattern, rule in self.rules.items():
            if pattern.endswith('*') and path.startswith(pattern[:-1]):
                return rule

        return None

    async def check_rate_limit(self, request: Request) -> tuple[bool, dict]:
        """Check if request is allowed."""
        path = request.url.path
        rule = self.get_rule(path)

        if not rule:
            return True, {}

        # Generate key based on scope
        if rule['scope'] == 'ip':
            key = self._get_client_ip(request)
        elif rule['scope'] == 'user':
            key = self._get_user_id(request)
        elif rule['scope'] == 'api_key':
            key = self._get_api_key(request)
        else:
            key = self._get_client_ip(request)

        # Add path to key for per-path limiting
        key = f"{key}:{path}"

        return await self.strategy.is_allowed(key, rule['limit'], rule['window'])

    def _get_client_ip(self, request: Request) -> str:
        """Get client IP address."""
        # Check for forwarded headers
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            return forwarded_for.split(",")[0].strip()

        real_ip = request.headers.get("X-Real-IP")
        if real_ip:
            return real_ip

        # Fall back to client IP
        return request.client.host if request.client else "unknown"

    def _get_user_id(self, request: Request) -> str:
        """Get user ID from request."""
        # This would typically come from JWT token or session
        return request.headers.get("X-User-ID", "anonymous")

    def _get_api_key(self, request: Request) -> str:
        """Get API key from request."""
        return request.headers.get("X-API-Key", "none")


class RateLimitMiddleware(BaseHTTPMiddleware):
    """FastAPI middleware for rate limiting."""

    def __init__(self, app, redis_url: str | None = None):
        super().__init__(app)
        self.redis_client = None

        # Initialize Redis if available
        if redis_url:
            try:
                self.redis_client = redis.from_url(redis_url)
                asyncio.create_task(self._test_redis())
            except Exception as e:
                logger.warning(f"Redis not available for rate limiting: {e}")

        # Initialize strategy and limiter
        strategy = TokenBucketStrategy(self.redis_client)
        self.limiter = RateLimiter(strategy)

        # Add default rules
        self._setup_default_rules()

    async def _test_redis(self):
        """Test Redis connection."""
        try:
            await self.redis_client.ping()
            logger.info("Rate limiting: Redis connected")
        except Exception as e:
            logger.warning(f"Rate limiting: Redis connection failed: {e}")
            self.redis_client = None

    def _setup_default_rules(self):
        """Setup default rate limiting rules."""
        settings = get_settings()

        # API endpoints
        self.limiter.add_rule("/analyze/*", limit=100, window=60, scope="ip")
        self.limiter.add_rule("/ask", limit=50, window=60, scope="ip")
        self.limiter.add_rule("/search", limit=200, window=60, scope="ip")

        # Health endpoints (no limit)
        self.limiter.add_rule("/health*", limit=1000, window=60, scope="ip")

        # Admin endpoints (stricter)
        self.limiter.add_rule("/admin/*", limit=10, window=60, scope="user")

        # Global fallback
        self.limiter.add_rule("*", limit=1000, window=60, scope="ip")

    async def dispatch(self, request: Request, call_next):
        """Process request with rate limiting."""
        # Skip rate limiting for certain paths
        if self._should_skip(request):
            return await call_next(request)

        # Check rate limit
        allowed, info = await self.limiter.check_rate_limit(request)

        if not allowed:
            # Log rate limit violation
            logger.warning(
                f"Rate limit exceeded for {self.limiter._get_client_ip(request)} "
                f"on {request.url.path}: {info}"
            )

            # Return rate limit error
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail={
                    "error": "Rate limit exceeded",
                    "limit": info.get('limit'),
                    "window": info.get('window'),
                    "retry_after": info.get('retry_after')
                },
                headers={
                    "Retry-After": str(int(info.get('retry_after', 1)))
                }
            )

        # Add rate limit headers
        response = await call_next(request)
        response.headers["X-RateLimit-Limit"] = str(info.get('limit', ''))
        response.headers["X-RateLimit-Remaining"] = str(info.get('remaining', info.get('tokens', '')))
        response.headers["X-RateLimit-Window"] = str(info.get('window', ''))

        return response

    def _should_skip(self, request: Request) -> bool:
        """Check if rate limiting should be skipped for this request."""
        skip_paths = ["/docs", "/redoc", "/openapi.json", "/metrics", "/favicon.ico"]
        return any(request.url.path.startswith(path) for path in skip_paths)


# Factory function for easy initialization
def create_rate_limiter(app, redis_url: str | None = None) -> RateLimitMiddleware:
    """Create and configure rate limiter middleware."""
    return RateLimitMiddleware(app, redis_url)