# core/rate_limit.py

from fastapi import Request, WebSocket, HTTPException, status
from starlette.requests import HTTPConnection
from db.redis import redis_client

class RateLimiter:
    def __init__(self, requests: int, window_seconds: int, scope: str = "global"):
        """
        :param requests: number of requests allowed
        :param window_seconds: window time in seconds
        :param scope: type of limit ("global" for all site, "endpoint" for a specific path)
        """
        self.requests = requests
        self.window_seconds = window_seconds
        self.scope = scope

    async def __call__(self, connection: HTTPConnection):
        client_ip = connection.client.host if connection.client else "127.0.0.1"
        
        real_ip = connection.headers.get("x-real-ip", connection.headers.get("x-forwarded-for", client_ip))
        real_ip = real_ip.split(",")[0].strip()
        if self.scope == "global":
            key = f"rate_limit:global:{real_ip}"
        else:
            path = connection.scope["path"]
            key = f"rate_limit:endpoint:{real_ip}:{path}"

        current_count = await redis_client.incr(key)

        if current_count == 1:
            await redis_client.expire(key, self.window_seconds)

        if current_count > self.requests:
            await redis_client.expire(key, self.window_seconds)
            
            # If it's a WebSocket connection, we might want to raise WebSocketException
            # But Starlette's HTTPException is also handled by FastAPI for WebSockets by closing the connection.
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail="تعداد درخواست‌های شما از حد مجاز گذشته است. لطفاً بعداً دوباره تلاش کنید."
            )