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):
        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:
        if not event.media_group_id:
            return await handler(event, data)

        group_id = event.media_group_id

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

                # Проверяем, что ключ все еще существует
                if group_id in self.album_data:
                    # Передаем собранный альбом в хендлер
                    # Сортируем по message_id, чтобы порядок был верным
                    current_album = self.album_data[group_id]
                    current_album.sort(key=lambda x: x.message_id)
                    data["album"] = current_album
                    return await handler(event, data)

            finally:
                # ЧИСТКА: Удаляем запись после обработки или таймаута
                # Используем pop() с дефолтом, чтобы избежать KeyError
                self.album_data.pop(group_id, None)

        else:
            # Если группа уже собирается - просто добавляем и выходим
            self.album_data[group_id].append(event)
            return