ForcePilot/backend/package/yuxi/channels/adapters/qqbot/message_queue.py

from __future__ import annotations

import asyncio
import logging

from yuxi.channels.models import ChannelMessage

logger = logging.getLogger(__name__)


class MessageQueue:
    def __init__(
        self,
        global_limit: int = 1000,
        per_user_limit: int = 20,
        per_group_limit: int = 50,
        max_concurrency: int = 10,
    ):
        self._queue: asyncio.Queue[ChannelMessage] = asyncio.Queue(maxsize=global_limit)
        self._global_limit = global_limit
        self._per_user_limit = per_user_limit
        self._per_group_limit = per_group_limit
        self._max_concurrency = max_concurrency

        self._per_user_counts: dict[str, int] = {}
        self._per_group_counts: dict[str, int] = {}

        self._semaphore = asyncio.Semaphore(max_concurrency)
        self._lock = asyncio.Lock()

    async def enqueue(self, msg: ChannelMessage) -> bool:
        user_id = msg.identity.channel_user_id
        chat_id = msg.identity.channel_chat_id

        async with self._lock:
            if self._queue.qsize() >= self._global_limit:
                try:
                    self._queue.get_nowait()
                    self._queue.task_done()
                except asyncio.QueueEmpty:
                    pass

            if user_id:
                user_count = self._per_user_counts.get(user_id, 0)
                if user_count >= self._per_user_limit:
                    logger.debug("MessageQueue: per-user limit reached for %s", user_id)
                    return False

            if chat_id and chat_id.startswith("group_"):
                group_count = self._per_group_counts.get(chat_id, 0)
                if group_count >= self._per_group_limit:
                    logger.debug("MessageQueue: per-group limit reached for %s", chat_id)
                    return False

            try:
                self._queue.put_nowait(msg)
                if user_id:
                    self._per_user_counts[user_id] = self._per_user_counts.get(user_id, 0) + 1
                if chat_id and chat_id.startswith("group_"):
                    self._per_group_counts[chat_id] = self._per_group_counts.get(chat_id, 0) + 1
                return True
            except asyncio.QueueFull:
                return False

    async def dequeue(self) -> ChannelMessage:
        msg = await self._queue.get()
        user_id = msg.identity.channel_user_id
        chat_id = msg.identity.channel_chat_id

        async with self._lock:
            if user_id:
                self._per_user_counts[user_id] = max(0, self._per_user_counts.get(user_id, 1) - 1)
            if chat_id and chat_id.startswith("group_"):
                self._per_group_counts[chat_id] = max(0, self._per_group_counts.get(chat_id, 1) - 1)

        return msg

    def task_done(self) -> None:
        self._queue.task_done()

    @property
    def size(self) -> int:
        return self._queue.qsize()

    @property
    def snapshot(self) -> dict:
        return {
            "queue_size": self._queue.qsize(),
            "global_limit": self._global_limit,
            "per_user_counts": dict(self._per_user_counts),
            "per_group_counts": dict(self._per_group_counts),
        }

    async def acquire(self) -> bool:
        try:
            await self._semaphore.acquire()
            return True
        except Exception:
            return False

    def release(self) -> None:
        try:
            self._semaphore.release()
        except ValueError:
            pass

    async def drain(self) -> list[ChannelMessage]:
        messages: list[ChannelMessage] = []
        async with self._lock:
            while not self._queue.empty():
                try:
                    messages.append(self._queue.get_nowait())
                except asyncio.QueueEmpty:
                    break
            self._per_user_counts.clear()
            self._per_group_counts.clear()
        return messages