telegram-helper-bot/helper_bot/middlewares/album_middleware.py

import asyncio
from typing import Any, Dict, List, Optional, Union

from aiogram import BaseMiddleware
from aiogram.types import Message


class AlbumGetter:
    """Вспомогательный класс для получения полной медиагруппы из middleware"""

    def __init__(
        self, album_data: Dict[str, Any], media_group_id: str, event: asyncio.Event
    ):
        self.album_data = album_data
        self.media_group_id = media_group_id
        self.event = event

    async def get_album(self, timeout: float = 10.0) -> Optional[List[Message]]:
        """
        Ждет полную медиагруппу и возвращает ее.

        Args:
            timeout: Максимальное время ожидания в секундах

        Returns:
            Список сообщений медиагруппы или None при таймауте
        """
        try:
            await asyncio.wait_for(self.event.wait(), timeout=timeout)
            if self.media_group_id in self.album_data:
                return self.album_data[self.media_group_id].get("collected_album")
            return None
        except asyncio.TimeoutError:
            return None


class AlbumMiddleware(BaseMiddleware):
    """
    Middleware для обработки медиа групп в Telegram.
    Собирает все сообщения одной медиа группы и передает их как album в data.
    Не блокирует handler - сразу вызывает его, а полную медиагруппу передает через Event.
    """

    def __init__(self, latency: Union[int, float] = 5.0):
        """
        Инициализация middleware.

        Args:
            latency: Задержка в секундах для сбора всех сообщений медиа группы
        """
        super().__init__()
        self.latency = latency
        # Храним данные медиагруппы: messages, event для уведомления, task для сбора
        self.album_data: Dict[str, Dict[str, Any]] = {}

    def collect_album_messages(self, event: Message) -> int:
        """
        Собирает сообщения одной медиа группы.

        Args:
            event: Сообщение для обработки

        Returns:
            Количество сообщений в текущей медиа группе
        """
        if not event.media_group_id:
            return 0

        if event.media_group_id not in self.album_data:
            self.album_data[event.media_group_id] = {"messages": []}

        self.album_data[event.media_group_id]["messages"].append(event)
        return len(self.album_data[event.media_group_id]["messages"])

    async def _collect_album_background(self, media_group_id: str) -> None:
        """
        Фоновая задача для сбора всех сообщений медиагруппы.

        Args:
            media_group_id: ID медиагруппы для сбора
        """
        try:
            await asyncio.sleep(self.latency)

            if media_group_id not in self.album_data:
                return

            # Получаем текущий список сообщений
            album_messages = self.album_data[media_group_id]["messages"].copy()
            album_messages.sort(key=lambda x: x.message_id)

            # Сохраняем собранную медиагруппу и уведомляем через Event
            self.album_data[media_group_id]["collected_album"] = album_messages
            self.album_data[media_group_id]["event"].set()

            # Очищаем данные после небольшой задержки (чтобы handler успел получить album)
            await asyncio.sleep(1.0)
            if media_group_id in self.album_data:
                task = self.album_data[media_group_id].get("task")
                if task and not task.done():
                    task.cancel()
                del self.album_data[media_group_id]
        except Exception:
            # В случае ошибки все равно уведомляем, чтобы handler не завис
            if media_group_id in self.album_data:
                self.album_data[media_group_id]["event"].set()
                # Очищаем данные даже при ошибке
                try:
                    task = self.album_data[media_group_id].get("task")
                    if task and not task.done():
                        task.cancel()
                    del self.album_data[media_group_id]
                except Exception:
                    pass

    async def __call__(self, handler, event: Message, data: Dict[str, Any]) -> Any:
        """
        Основная логика middleware.

        Для медиагрупп: сразу вызывает handler, передавая Event для получения полной медиагруппы.
        Для обычных сообщений: сразу вызывает handler.

        Args:
            handler: Обработчик события
            event: Событие (сообщение)
            data: Данные для передачи в обработчик

        Returns:
            Результат выполнения обработчика
        """
        if not event.media_group_id:
            return await handler(event, data)

        media_group_id = event.media_group_id
        message_id = event.message_id

        # Если это первое сообщение медиагруппы - создаем структуру данных
        is_first_message = False
        if media_group_id not in self.album_data:
            is_first_message = True
            album_event = asyncio.Event()
            self.album_data[media_group_id] = {
                "messages": [],
                "event": album_event,
                "task": None,
                "first_message_id": message_id,
            }
            # Запускаем фоновую задачу для сбора медиагруппы
            task = asyncio.create_task(self._collect_album_background(media_group_id))
            self.album_data[media_group_id]["task"] = task

        # Добавляем сообщение в медиагруппу
        self.album_data[media_group_id]["messages"].append(event)

        # Обрабатываем только первое сообщение медиагруппы
        if not is_first_message:
            # Для остальных сообщений просто возвращаемся, не вызывая handler
            return

        # Передаем объект-геттер в data, чтобы handler мог получить полную медиагруппу
        album_getter = AlbumGetter(
            self.album_data, media_group_id, self.album_data[media_group_id]["event"]
        )
        data["album_getter"] = album_getter

        # Сразу вызываем handler только для первого сообщения (не блокируем)
        return await handler(event, data)