from __future__ import annotations

import asyncio
import logging
from collections.abc import Callable, Coroutine
from typing import Protocol, TypeVar

T = TypeVar("T")

_logger = logging.getLogger(__name__)


class KeyedAsyncQueueHooks(Protocol):
    def on_enqueue(self) -> None: ...
    def on_settle(self) -> None: ...


def enqueue_keyed_task(
    tails: dict[str, asyncio.Future],
    key: str,
    task: Callable[[], Coroutine[None, None, T]],
    *,
    hooks: KeyedAsyncQueueHooks | None = None,
) -> Coroutine[None, None, T]:
    async def _run() -> T:
        if hooks:
            hooks.on_enqueue()

        prev = tails.get(key)
        if prev is not None and not prev.done():
            try:
                await prev
            except Exception:
                _logger.debug("Keyed task predecessor failed for key=%s", key, exc_info=True)

        loop = asyncio.get_running_loop()
        tail: asyncio.Future = loop.create_future()
        tails[key] = tail

        try:
            return await task()
        finally:
            if hooks:
                hooks.on_settle()
            tail.set_result(None)
            if tails.get(key) is tail:
                tails.pop(key, None)

    return _run()


class KeyedAsyncQueue:

    def __init__(self) -> None:
        self._tails: dict[str, asyncio.Future] = {}

    def get_tails(self) -> dict[str, asyncio.Future]:
        return self._tails

    async def enqueue(
        self,
        key: str,
        task: Callable[[], Coroutine[None, None, T]],
        *,
        hooks: KeyedAsyncQueueHooks | None = None,
    ) -> T:
        return await enqueue_keyed_task(self._tails, key, task, hooks=hooks)