middleware/rate_limit.py

from fastapi import Request, HTTPException
from datetime import datetime, timedelta
from collections import defaultdict
import asyncio

# In-memory store: {key: [timestamp, ...]}
_request_log: dict = defaultdict(list)
_lock = asyncio.Lock()

LIMITS = {
    "guest":  {"analyze": 3,  "run": 10, "challenge": 2},
    "free":   {"analyze": 15, "run": 50, "challenge": 10},
    "pro":    {"analyze": 999, "run": 999, "challenge": 999},
    "admin":  {"analyze": 999, "run": 999, "challenge": 999},
}


async def check_rate_limit(request: Request, action: str, user=None, role: str = "guest"):
    """Check if user/IP has exceeded their daily limit for a given action."""
    if role in ("admin", "pro"):
        return  # Unlimited

    identifier = str(user.id) if user else request.client.host
    key = f"{identifier}:{action}"
    now = datetime.utcnow()
    window_start = now - timedelta(hours=24)

    async with _lock:
        # Remove entries older than 24h
        _request_log[key] = [t for t in _request_log[key] if t > window_start]
        count = len(_request_log[key])
        limit = LIMITS.get(role, LIMITS["guest"]).get(action, 5)

        if count >= limit:
            raise HTTPException(
                status_code=429,
                detail=f"Daily limit of {limit} {action}s reached. Upgrade to Pro for unlimited access."
            )
        _request_log[key].append(now)