"""
基于内存的简易限流中间件（MVP不依赖Redis）
"""
import time
from collections import defaultdict
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse

class RateLimitMiddleware(BaseHTTPMiddleware):
    def __init__(self, app):
        super().__init__(app)
        # {key: [(timestamp, ...)]}
        self._requests = defaultdict(list)
        
        # 限流规则
        self.rules = {
            "auth": {  # /api/v1/auth/login, register, forgot-password
                "paths": ["/api/v1/auth/login", "/api/v1/auth/register", "/api/v1/auth/forgot-password"],
                "max_requests": 100,
                "window_seconds": 60,
            },
            "query_run": {  # run-now
                "paths": ["/run-now"],  # 用 endswith 匹配
                "max_requests": 10,
                "window_seconds": 3600,
            },
            "global": {
                "max_requests": 100,
                "window_seconds": 60,
            }
        }
    
    async def dispatch(self, request: Request, call_next):
        client_ip = request.client.host if request.client else "unknown"
        path = request.url.path
        now = time.time()
        
        # 健康检查不限流
        if path == "/health" or path.startswith("/docs") or path.startswith("/openapi"):
            return await call_next(request)
        
        # 检查认证接口限流
        if any(path == p for p in self.rules["auth"]["paths"]):
            key = f"auth:{client_ip}"
            if self._is_rate_limited(key, now, self.rules["auth"]):
                return JSONResponse(
                    status_code=429,
                    content={"detail": "请求过于频繁，请稍后再试"}
                )
        
        # 检查查询执行限流
        if path.endswith("/run-now") and request.method == "POST":
            key = f"query_run:{client_ip}"
            if self._is_rate_limited(key, now, self.rules["query_run"]):
                return JSONResponse(
                    status_code=429,
                    content={"detail": "查询执行过于频繁，请稍后再试"}
                )
        
        # 全局限流
        key = f"global:{client_ip}"
        if self._is_rate_limited(key, now, self.rules["global"]):
            return JSONResponse(
                status_code=429,
                content={"detail": "请求过于频繁，请稍后再试"}
            )
        
        return await call_next(request)
    
    def _is_rate_limited(self, key, now, rule):
        window = rule["window_seconds"]
        max_req = rule["max_requests"]
        
        # 清理过期记录
        self._requests[key] = [t for t in self._requests[key] if now - t < window]
        
        if len(self._requests[key]) >= max_req:
            return True
        
        self._requests[key].append(now)
        return False
