import asyncpg
from aiogram import BaseMiddleware
from aiogram.types import Update


def _chat_id(update: Update) -> int | None:
    if update.message is not None:
        return update.message.chat.id
    if update.callback_query is not None and update.callback_query.message is not None:
        return update.callback_query.message.chat.id
    return None


class DedupMiddleware(BaseMiddleware):
    """Один update_id обрабатывается один раз, даже если Telegram прислал повтор.

    Повторная доставка на другую реплику упирается в первичный ключ.
    Сообщения одного чата сериализуются advisory-lock'ом в Postgres.
    """

    def __init__(self, pool: asyncpg.Pool) -> None:
        self.pool = pool

    async def __call__(self, handler, event, data):
        if not isinstance(event, Update):
            return await handler(event, data)

        async with self.pool.acquire() as conn:
            inserted = await conn.fetchrow(
                "INSERT INTO telegram_updates (update_id) VALUES ($1) "
                "ON CONFLICT DO NOTHING RETURNING update_id",
                event.update_id,
            )
            if inserted is None:
                return None
            try:
                async with conn.transaction():
                    chat_id = _chat_id(event)
                    if chat_id is not None:
                        await conn.execute("SELECT pg_advisory_xact_lock($1)", chat_id)
                    return await handler(event, data)
            except Exception:
                await conn.execute("DELETE FROM telegram_updates WHERE update_id = $1", event.update_id)
                raise