import time

from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
from collections import defaultdict


class RateLimitMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, requests_per_minute: int = 10):
        super().__init__(app)
        self.requests_per_minute = requests_per_minute
        self.requests: dict = defaultdict(list)

    async def dispatch(self, request: Request, call_next):
        if request.url.path.startswith("/api/health"):
            return await call_next(request)

        user_key = request.headers.get("X-Forwarded-For", request.client.host if request.client else "unknown")
        now = time.time()
        window = now - 60

        self.requests[user_key] = [t for t in self.requests[user_key] if t > window]
        self.requests[user_key].append(now)

        if len(self.requests[user_key]) > self.requests_per_minute:
            return JSONResponse(status_code=429, content={"detail": "Rate limit exceeded"})

        return await call_next(request)