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, чтобы порядок был верным
                    self.album_data[group_id].sort(key=lambda x: x.message_id)
                    data["album"] = self.album_data[group_id]
                    return await handler(event, data)

            finally:
                # ЧИСТКА: Удаляем всегда, если это "головной" поток, который создал запись
                # Проверяем, что мы удаляем именно то, что создали, и ключ существует
                if group_id in self.album_data and self.album_data[group_id][0] == event:
                    del self.album_data[group_id]

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