from __future__ import annotations

import asyncio
from collections.abc import Callable, Awaitable
from typing import Any

from yuxi.channels.models import ChannelMessage

QUEUE_DEBOUNCE_MS = 500


class WeChatQueueDebounce:
    def __init__(self, debounce_ms: int = QUEUE_DEBOUNCE_MS):
        self._debounce_s = debounce_ms / 1000.0
        self._pending: dict[str, asyncio.Task] = {}
        self._queues: dict[str, list[ChannelMessage]] = {}
        self._lock = asyncio.Lock()

    async def enqueue(
        self,
        session_key: str,
        message: ChannelMessage,
        handler: Callable[[list[ChannelMessage]], Awaitable[Any]],
    ) -> None:
        async with self._lock:
            if session_key not in self._queues:
                self._queues[session_key] = []
            self._queues[session_key].append(message)

            if session_key in self._pending:
                self._pending[session_key].cancel()
            self._pending[session_key] = asyncio.create_task(self._flush_after_delay(session_key, handler))

    async def _flush_after_delay(
        self,
        session_key: str,
        handler: Callable[[list[ChannelMessage]], Awaitable[Any]],
    ) -> None:
        await asyncio.sleep(self._debounce_s)
        async with self._lock:
            messages = self._queues.pop(session_key, [])
            self._pending.pop(session_key, None)
        if messages:
            await handler(messages)

    async def flush_all(self) -> None:
        async with self._lock:
            keys = list(self._queues.keys())
        for key in keys:
            async with self._lock:
                if key in self._pending:
                    self._pending[key].cancel()
                    self._pending.pop(key, None)

    @property
    def debounce_ms(self) -> int:
        return int(self._debounce_s * 1000)