Thunder/utils/rate_limiter.py

# Thunder/utils/rate_limiter.py

import time
import math
import asyncio
from collections import deque
from typing import Callable, Dict, Optional, Tuple
from pyrogram import Client
from pyrogram.types import Message
from pyrogram.errors import FloodWait, RPCError
from Thunder.utils.logger import logger
from Thunder.utils.database import db
from Thunder.utils.messages import (
    MSG_RATE_LIMIT_QUEUE_PRIORITY,
    MSG_RATE_LIMIT_QUEUE_REGULAR,
    MSG_RATE_LIMIT_QUEUE_FULL
)
from Thunder.vars import Var


class QueueFullError(Exception):
    pass


class RateLimiter:
    def __init__(self):
        self.request_queue: deque = deque()
        self.priority_queue: deque = deque()
        self.user_queue_counts: Dict[int, int] = {}

        self.request_event: asyncio.Event = asyncio.Event()
        self.request_lock: asyncio.Lock = asyncio.Lock()

        self.user_requests: Dict[int, deque] = {}
        self.global_requests: deque = deque()

        self.processing_times: deque = deque(maxlen=100)
        self.file_processing_times: Dict[str, deque] = {}
        self.average_processing_time: float = 1.0

        self.auth_cache: Dict[int, Tuple[bool, float]] = {}
        self.auth_cache_ttl_seconds: int = 300

        self._initialization_error = False
        self._load_configuration()

    def _load_configuration(self):
        try:
            self.max_requests_per_period = Var.MAX_FILES_PER_PERIOD
            self.rate_limit_period_seconds = Var.RATE_LIMIT_PERIOD_MINUTES * 60
            self.max_queue_size = Var.MAX_QUEUE_SIZE
            self.enabled = Var.RATE_LIMIT_ENABLED
            self.global_rate_limit_enabled = Var.GLOBAL_RATE_LIMIT
            self.max_global_requests_per_minute = Var.MAX_GLOBAL_REQUESTS_PER_MINUTE

            if not self._validate_configuration():
                logger.warning("Rate limiter disabled due to invalid configuration.")
                self.enabled = False
            else:
                logger.debug(f"Rate limiter initialized: enabled={self.enabled}, "
                             f"max_requests={self.max_requests_per_period}, "
                             f"period={self.rate_limit_period_seconds}s, "
                             f"queue_size={self.max_queue_size}, "
                             f"global_enabled={self.global_rate_limit_enabled}, "
                             f"max_global_requests={self.max_global_requests_per_minute}")
        except Exception as e:
            logger.critical(f"Critical error initializing rate limiter, using safe defaults: {e}", exc_info=True)
            self.max_requests_per_period = 5
            self.rate_limit_period_seconds = 60
            self.max_queue_size = 100
            self.enabled = False
            self.global_rate_limit_enabled = False
            self.max_global_requests_per_minute = 60
            self._initialization_error = True

    def _validate_configuration(self) -> bool:
        is_valid = True
        if self.max_requests_per_period <= 0:
            logger.error("Invalid MAX_FILES_PER_PERIOD: must be > 0.")
            is_valid = False
        if self.rate_limit_period_seconds <= 0:
            logger.error("Invalid RATE_LIMIT_PERIOD_MINUTES: must be > 0.")
            is_valid = False
        if self.max_queue_size <= 0:
            logger.error("Invalid MAX_QUEUE_SIZE: must be > 0.")
            is_valid = False
        if self.global_rate_limit_enabled and self.max_global_requests_per_minute <= 0:
            logger.error("Invalid MAX_GLOBAL_REQUESTS_PER_MINUTE: must be > 0 when global rate limit is enabled.")
            is_valid = False
        return is_valid

    def is_owner(self, user_id: int) -> bool:
        return user_id == Var.OWNER_ID

    async def is_authorized_user(self, user_id: int) -> bool:
        current_time = time.time()
        if user_id in self.auth_cache:
            is_auth, timestamp = self.auth_cache[user_id]
            if current_time - timestamp < self.auth_cache_ttl_seconds:
                return is_auth

        try:
            authorized_user = await db.authorized_users_col.find_one({"user_id": user_id})
            is_auth = bool(authorized_user)
            self.auth_cache[user_id] = (is_auth, current_time)
            return is_auth
        except Exception as e:
            logger.error(f"Database error checking authorized user {user_id}: {e}")
            return False

    async def get_user_priority(self, user_id: int) -> str:
        if self.is_owner(user_id):
            return 'owner'
        if await self.is_authorized_user(user_id):
            return 'authorized'
        return 'regular'

    async def check_limits(self, user_id: int, record: bool = True) -> bool:
        if not self.enabled or self._initialization_error or self.is_owner(user_id):
            return True

        current_time = time.time()

        if self.global_rate_limit_enabled:
            while self.global_requests and self.global_requests[0] <= current_time - 60:
                self.global_requests.popleft()
            if len(self.global_requests) >= self.max_global_requests_per_minute:
                return False

        user_timestamps = self.user_requests.setdefault(user_id, deque())
        while user_timestamps and user_timestamps[0] <= current_time - self.rate_limit_period_seconds:
            user_timestamps.popleft()
        if len(user_timestamps) >= self.max_requests_per_period:
            return False

        if record:
            if self.global_rate_limit_enabled:
                self.global_requests.append(current_time)
            user_timestamps.append(current_time)
        return True

    async def _requeue_request(self, request_data: dict, queue_type: str):
        async with self.request_lock:
            if queue_type == "priority":
                self.priority_queue.appendleft(request_data)
            else:
                self.request_queue.appendleft(request_data)
            self.request_event.set()
        logger.debug(f"Re-queued request for user {request_data['user_id']} to {queue_type} queue.")

    async def add_to_queue(self, func: Callable, user_id: int, file_identifier: Optional[str] = None, *args, **kwargs):
        if not self.enabled:
            await func(*args, **kwargs)
            return

        request_data = {
            'func': func, 'user_id': user_id, 'args': args, 'kwargs': kwargs,
            'timestamp': time.time(), 'user_priority': await self.get_user_priority(user_id),
            'file_identifier': file_identifier
        }

        async with self.request_lock:
            total_queued = len(self.request_queue) + len(self.priority_queue)
            if total_queued >= self.max_queue_size:
                raise QueueFullError("Queue is full")

            if request_data['user_priority'] == 'authorized':
                self.priority_queue.append(request_data)
                queue_name = "priority"
            else:
                self.request_queue.append(request_data)
                queue_name = "regular"

            self.user_queue_counts[user_id] = self.user_queue_counts.get(user_id, 0) + 1
            logger.debug(f"Added request for user {user_id} to {queue_name} queue. Total queued: {total_queued + 1}")
            self.request_event.set()

    async def request_executor(self):
        logger.debug("Request executor started.")
        while True:
            try:
                await self.request_event.wait()

                async with self.request_lock:
                    queue, queue_type = (self.priority_queue, "priority") if self.priority_queue else (self.request_queue, "regular")
                    if not queue:
                        self.request_event.clear()
                        continue
                    request_data = queue.popleft()

                user_id = request_data['user_id']
                processed = False
                if not self.is_owner(user_id):
                    if not await self.check_limits(user_id, record=True):
                        await self._requeue_request(request_data, queue_type)
                        await asyncio.sleep(0.5)
                        continue

                logger.debug(f"Processing request for user {user_id} from {queue_type} queue.")
                start_time = time.time()
                try:
                    await request_data['func'](*request_data['args'], **request_data['kwargs'])
                    processing_time = time.time() - start_time
                    self.processing_times.append(processing_time)
                    if self.processing_times:
                        self.average_processing_time = sum(self.processing_times) / len(self.processing_times)

                    file_identifier = request_data.get('file_identifier')
                    if file_identifier:
                        file_times = self.file_processing_times.setdefault(file_identifier, deque(maxlen=100))
                        file_times.append(processing_time)

                    processed = True

                except FloodWait as e:
                    logger.warning(f"FloodWait for user {user_id}, waiting {e.value}s before re-queuing.")
                    await asyncio.sleep(e.value)
                    await self._requeue_request(request_data, queue_type)
                except Exception as e:
                    logger.error(f"Error processing queued request for user {user_id}: {e}", exc_info=True)
                    processed = True
                finally:
                    async with self.request_lock:
                        if processed and user_id in self.user_queue_counts:
                            self.user_queue_counts[user_id] -= 1
                            if self.user_queue_counts[user_id] <= 0:
                                self.user_queue_counts.pop(user_id, None)

            except asyncio.CancelledError:
                logger.debug("Request executor cancelled, shutting down.")
                break
            except Exception as e:
                logger.critical(f"Critical error in request executor: {e}", exc_info=True)
                await asyncio.sleep(5)

    async def shutdown(self):
        logger.debug("Shutting down rate limiter and clearing queues...")
        async with self.request_lock:
            self.request_queue.clear()
            self.priority_queue.clear()
            self.user_queue_counts.clear()
            self.request_event.clear()
        logger.debug("Rate limiter queues cleared.")

    def get_queue_status(self) -> dict:
        return {
            'regular_queue_size': len(self.request_queue),
            'priority_queue_size': len(self.priority_queue),
            'total_queued': len(self.request_queue) + len(self.priority_queue),
            'max_queue_size': self.max_queue_size,
            'active_users_in_queue': len(self.user_queue_counts),
            'enabled': self.enabled,
        }

    async def get_user_queue_position(self, user_id: int) -> dict:
        user_priority = await self.get_user_priority(user_id)
        position = -1
        queue_to_search = self.priority_queue if user_priority == 'authorized' else self.request_queue
        
        for idx, req in enumerate(queue_to_search):
            if req.get('user_id') == user_id:
                position = idx + 1
                break

        effective_position = position
        if user_priority == 'regular' and position > -1:
            effective_position += len(self.priority_queue)

        return {
            'user_priority': user_priority,
            'position_in_own_queue': position if position > -1 else None,
            'effective_position': effective_position if effective_position > -1 else None,
            'priority_queue_size': len(self.priority_queue),
            'regular_queue_size': len(self.request_queue),
            'bypasses_rate_limit': user_priority == 'owner'
        }

    def _get_base_processing_time(self, file_identifier: Optional[str]) -> float:
        if file_identifier and file_identifier in self.file_processing_times:
            file_times = self.file_processing_times[file_identifier]
            if file_times:
                return sum(file_times) / len(file_times)
        return self.average_processing_time

    async def _calculate_queue_wait(self, user_id: int, effective_processing_time: float) -> float:
        pos_info = await self.get_user_queue_position(user_id)
        items_ahead = (pos_info['effective_position'] - 1) if pos_info['effective_position'] else 0
        return items_ahead * effective_processing_time

    def _calculate_user_rate_limit_wait(self, user_id: int, future_time: float) -> float:
        user_timestamps = self.user_requests.get(user_id, deque())
        future_user_timestamps = deque(ts for ts in user_timestamps if ts > future_time - self.rate_limit_period_seconds)

        if len(future_user_timestamps) >= self.max_requests_per_period:
            reset_time = future_user_timestamps[0] + self.rate_limit_period_seconds
            return max(0.0, reset_time - future_time)
        return 0.0

    def _calculate_global_rate_limit_wait(self, future_time: float) -> float:
        if not self.global_rate_limit_enabled:
            return 0.0

        future_global_requests = deque(ts for ts in self.global_requests if ts > future_time - 60)
        
        if len(future_global_requests) >= self.max_global_requests_per_minute:
            oldest_request_time = future_global_requests[0]
            reset_time = oldest_request_time + 60
            return max(0.0, reset_time - future_time)
        return 0.0

    async def estimate_wait_time(self, user_id: int, file_identifier: Optional[str] = None) -> float:
        if self.is_owner(user_id):
            return 0.0

        base_processing_time = self._get_base_processing_time(file_identifier)
        min_time_per_request = self.rate_limit_period_seconds / self.max_requests_per_period if self.max_requests_per_period > 0 else 0
        effective_processing_time = max(base_processing_time, min_time_per_request)
        
        if self.global_rate_limit_enabled and self.max_global_requests_per_minute > 0:
            min_time_per_global = 60 / self.max_global_requests_per_minute
            effective_processing_time = max(effective_processing_time, min_time_per_global)

        queue_wait = await self._calculate_queue_wait(user_id, effective_processing_time)
        future_time = time.time() + queue_wait

        rate_limit_wait = self._calculate_user_rate_limit_wait(user_id, future_time)
        global_wait = self._calculate_global_rate_limit_wait(future_time)

        return queue_wait + rate_limit_wait + global_wait


rate_limiter = RateLimiter()


async def request_executor():
    await rate_limiter.request_executor()


async def handle_rate_limited_request(bot: Client, message: Message, handler: Callable, *args, **kwargs):
    rl_user_id = kwargs.pop('rl_user_id', None)
    user_id = rl_user_id if rl_user_id is not None else (message.from_user.id if message and message.from_user else None)
    if not isinstance(user_id, int):
        logger.error(f"Invalid user_id provided for rate limiting: {user_id}")
        return

    file_identifier = message.document.file_unique_id if message and message.document else None

    if rate_limiter.is_owner(user_id):
        logger.debug(f"Owner {user_id} bypassing rate limit.")
        await handler(bot, message, *args, **kwargs)
        return

    if await rate_limiter.check_limits(user_id, record=True):
        logger.debug(f"User {user_id} within rate limits, executing immediately.")
        await handler(bot, message, *args, **kwargs)
        return

    is_channel = rl_user_id is not None and rl_user_id < 0

    if not is_channel:
        try:
            user_priority = await rate_limiter.get_user_priority(user_id)
            notification_msg = await send_queue_notification(
                bot, message, is_priority=(user_priority == 'authorized'), file_identifier=file_identifier
            )
            kwargs['notification_msg'] = notification_msg
        except Exception as e:
            logger.error(f"Error sending queue notification for user {user_id}: {e}", exc_info=True)

    try:
        await rate_limiter.add_to_queue(handler, user_id, file_identifier, bot, message, *args, **kwargs)
        logger.debug(f"Request for user {user_id} queued.")
    except QueueFullError:
        logger.warning(f"Queue full, request for user {user_id} rejected.")
        if not is_channel:
            await send_queue_full_message(bot, message, file_identifier)
    except Exception as e:
        logger.error(f"Error adding request to queue for user {user_id}: {e}", exc_info=True)
        if not is_channel:
            await send_queue_full_message(bot, message, file_identifier)


async def _send_notification(bot: Client, message: Message, template: str, file_identifier: Optional[str], **format_kwargs):
    try:
        if message.from_user:
            user_id = message.from_user.id
            wait_seconds = await rate_limiter.estimate_wait_time(user_id, file_identifier)
            wait_estimate = max(1, math.ceil(wait_seconds / 60))

            text = template.format(wait_estimate=wait_estimate, s="s" if wait_estimate > 1 else "", **format_kwargs)

            try:
                return await bot.send_message(
                    chat_id=message.chat.id,
                    text=text,
                    reply_to_message_id=message.id
                )
            except FloodWait as e:
                await asyncio.sleep(e.value)
                return await bot.send_message(
                    chat_id=message.chat.id,
                    text=text,
                    reply_to_message_id=message.id
                )
        else:
            logger.debug("Skipping notification for channel message (no from_user)")
            return None
    except (FloodWait, RPCError) as e:
        user_id = message.from_user.id if message.from_user else "channel"
        logger.warning(f"Error sending notification to user {user_id}: {e}")
    except Exception as e:
        logger.error(f"Unexpected error sending notification: {e}", exc_info=True)
    return None


async def send_queue_notification(bot: Client, message: Message, is_priority: bool, file_identifier: Optional[str]):
    if is_priority:
        template = MSG_RATE_LIMIT_QUEUE_PRIORITY
        params = {}
    else:
        template = MSG_RATE_LIMIT_QUEUE_REGULAR
        time_window = rate_limiter.rate_limit_period_seconds // 60
        params = {
            "max_requests": rate_limiter.max_requests_per_period,
            "time_window": time_window,
            "s1": "s" if rate_limiter.max_requests_per_period > 1 else "",
            "s2": "s" if time_window > 1 else ""
        }
    user_id = message.from_user.id if message.from_user else "channel"
    logger.debug(f"Sending {'priority' if is_priority else 'regular'} queue notification to user {user_id}")
    return await _send_notification(bot, message, template, file_identifier, **params)


async def send_queue_full_message(bot: Client, message: Message, file_identifier: Optional[str]):
    user_id = message.from_user.id if message.from_user else "channel"
    logger.debug(f"Sending queue full message to user {user_id}")
    await _send_notification(bot, message, MSG_RATE_LIMIT_QUEUE_FULL, file_identifier)