import time

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

from app.config import settings
from app.utils.redis import get_redis


RATE_LIMITS = {
    "POST:/api/auth/login": (10, 60),
    "POST:/api/auth/change-password": (5, 60),
    "POST:/api/generation-records/optimize": (20, 60),
    "POST:/api/generation-records/*/generate": (10, 60),
    "POST:/api/generation-records/*/retry": (10, 60),
}


def _match_rate_limit(path: str, method: str) -> tuple[int, int] | None:
    key = f"{method}:{path}"
    if key in RATE_LIMITS:
        return RATE_LIMITS[key]
    # Check wildcard patterns
    for pattern, limit in RATE_LIMITS.items():
        pattern_method, pattern_path = pattern.split(":", 1)
        if method != pattern_method:
            continue
        pattern_parts = pattern_path.split("/")
        key_parts = path.split("/")
        if len(pattern_parts) != len(key_parts):
            continue
        match = True
        for pp, kp in zip(pattern_parts, key_parts):
            if pp != "*" and pp != kp:
                match = False
                break
        if match:
            return limit
    return None


class RateLimitMiddleware(BaseHTTPMiddleware):
    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        if not settings.RATE_LIMIT_ENABLED:
            return await call_next(request)

        limit_info = _match_rate_limit(request.url.path, request.method)
        if not limit_info:
            return await call_next(request)

        max_requests, window = limit_info
        # Use user_id from token or IP as identifier
        identifier = request.client.host if request.client else "unknown"

        redis = get_redis()
        if not redis:
            return await call_next(request)

        redis_key = f"ratelimit:{request.method}:{request.url.path}:{identifier}"

        try:
            current = await redis.incr(redis_key)
            if current == 1:
                await redis.expire(redis_key, window)
            if current > max_requests:
                ttl = await redis.ttl(redis_key)
                return JSONResponse(
                    status_code=429,
                    content={"detail": "请求过于频繁，请稍后重试"},
                    headers={"Retry-After": str(max(ttl, 1))},
                )
        except Exception:
            pass  # Redis unavailable, skip rate limiting

        return await call_next(request)