from __future__ import annotations

import asyncio

from collections.abc import AsyncIterator

from contextlib import suppress

from typing import Protocol

from . import RpcTask

__all__ = ["InMemoryMessageQueue", "MessageQueue"]

class MessageQueue(Protocol):

async def publish(self, task: RpcTask) -> None: ...

async def close(self) -> None: ...

def task_done(self) -> None: ...

async def join(self) -> None: ...

def __aiter__(self) -> AsyncIterator[RpcTask]: ...

class InMemoryMessageQueue:

"""Simple in-memory broker for RPC task dispatch."""

def __init__(self, *, maxsize: int = 0) -> None:

self._queue: asyncio.Queue[RpcTask | None] = asyncio.Queue(maxsize=maxsize)

self._closed = False

async def publish(self, task: RpcTask) -> None:

if self._closed:

msg = "mssage queue already closed"

raise RuntimeError(msg)

await self._queue.put(task)

async def close(self) -> None:

if self._closed:

return

self._closed = True

await self._queue.put(None)

async def join(self) -> None:

await self._queue.join()

def task_done(self) -> None:

with suppress(ValueError):

self._queue.task_done()

def __aiter__(self) -> AsyncIterator[RpcTask]:

return _QueueIterator(self)

class _QueueIterator:

def __init__(self, queue: InMemoryMessageQueue) -> None:

self._queue = queue

def __aiter__(self) -> _QueueIterator:

return self

async def __anext__(self) -> RpcTask:

item = await self._queue._queue.get()

if item is None:

self._queue.task_done()

raise StopAsyncIteration

return item