import asyncio
from typing import Any, Dict, List, Union, Callable, Awaitable
from aiogram import BaseMiddleware
from aiogram.types import Message


class AlbumMiddleware(BaseMiddleware):
    def __init__(self, latency: float = 0.5):
        # latency - задержка в секундах для сбора частей альбома
        self.latency = latency
        self.album_data: Dict[str, List[Message]] = {}

    async def __call__(
            self,
            handler: Callable[[Message, Dict[str, Any]], Awaitable[Any]],
            event: Message,
            data: Dict[str, Any]
    ) -> Any:
        # Если у сообщения нет media_group_id, это не альбом -> пропускаем дальше как обычно
        if not event.media_group_id:
            return await handler(event, data)

        group_id = event.media_group_id

        try:
            # Если этот альбом мы еще не видели (первое сообщение из пачки)
            if group_id not in self.album_data:
                self.album_data[group_id] = [event]  # Создаем список
                await asyncio.sleep(self.latency)  # Ждем остальные части

                # После ожидания кладем собранный список в data
                # Теперь в хендлере будет доступен аргумент 'album'
                data["album"] = self.album_data[group_id]

                # Вызываем хендлер ОДИН раз
                return await handler(event, data)

            else:
                # Если альбом уже собирается, просто добавляем сообщение в список
                # и НЕ вызываем хендлер (прерываем цепочку для этого сообщения)
                self.album_data[group_id].append(event)
                return

        finally:
            # Чистим память после обработки, если это был "главный" поток обработки
            if group_id in self.album_data and len(self.album_data[group_id]) > 1:
                # Маленький хак: удаляем только если обработчик завершился
                # Проверка len нужна, чтобы не удалить раньше времени в параллельных тасках,
                # но корректнее просто удалять в блоке первого сообщения.
                if event == self.album_data[group_id][0]:
                    del self.album_data[group_id]