import asyncio
import traceback
from dataclasses import dataclass
from collections import defaultdict
from typing import Coroutine, TypeVar

from pyrogram import Client, filters
from pyrogram.types import Message

from bot import config
from bot.loader import app1
from bot.sending import forward_post

_tasks = set()
T = TypeVar("T")


def background(coro: Coroutine[None, None, T]) -> asyncio.Task[T]:
    loop = asyncio.get_event_loop()
    task = loop.create_task(coro)
    _tasks.add(task)
    task.add_done_callback(_tasks.remove)
    return task


@dataclass
class Album:
    media_group_id: str
    messages: list[Message]


# chat_id: group_id: album
_albums: defaultdict[int, dict[str, Album]] = defaultdict(dict)


@app1.on_message(filters.media_group)
async def on_media_group(client: Client, message: Message):
    try:
        chat_id = message.chat.id
        media_group_id = message.media_group_id
        if media_group_id is None:
            return

        if media_group_id not in _albums[chat_id]:
            album = Album(messages=[message], media_group_id=media_group_id)
            _albums[chat_id][media_group_id] = album

            async def task():
                await asyncio.sleep(1)
                _albums[chat_id].pop(media_group_id, None)
                try:
                    album.messages.sort(key=lambda m: m.id)
                    await on_album(client, album)
                except Exception:
                    traceback.print_exc()

            background(task())
        else:
            album = _albums[chat_id][media_group_id]
            album.messages.append(message)
    finally:
        message.continue_propagation()


async def on_album(client: Client, album: Album):
    """
    Обрабатывает новые альбомы
    """
    if album.messages[0].chat.id == config.rent_channel_id:
        await forward_post(client, album.messages[0], config.groups_for_rent)
    if album.messages[0].chat.id == config.sell_channel_id:
        await forward_post(client, album.messages[0], config.groups_for_sell)