from __future__ import annotations

import asyncio

from dataclasses import dataclass

from typing import Any, Protocol

__all__ = [

"InMemoryMessageStateStore",

"IncomingMessage",

"MessageStateStore",

"OutgoingMessage",

]

@dataclass(slots=True)

class OutgoingMessage:

request_id: int

method: str

future: asyncio.Future[Any]

@dataclass(slots=True)

class IncomingMessage:

method: str

params: Any

status: str = "pending"

result: Any = None

error: Any = None

class MessageStateStore(Protocol):

def register_outgoing(self, request_id: int, method: str) -> asyncio.Future[Any]: ...

def resolve_outgoing(self, request_id: int, result: Any) -> None: ...

def reject_outgoing(self, request_id: int, error: Any) -> None: ...

def reject_all_outgoing(self, error: Any) -> None: ...

def begin_incoming(self, method: str, params: Any) -> IncomingMessage: ...

def complete_incoming(self, record: IncomingMessage, result: Any) -> None: ...

def fail_incoming(self, record: IncomingMessage, error: Any) -> None: ...

class InMemoryMessageStateStore(MessageStateStore):

def __init__(self) -> None:

self._outgoing: dict[int, OutgoingMessage] = {}

self._incoming: list[IncomingMessage] = []

def register_outgoing(self, request_id: int, method: str) -> asyncio.Future[Any]:

future: asyncio.Future[Any] = asyncio.get_running_loop().create_future()

self._outgoing[request_id] = OutgoingMessage(request_id, method, future)

return future

def resolve_outgoing(self, request_id: int, result: Any) -> None:

record = self._outgoing.pop(request_id, None)

if record and not record.future.done():

record.future.set_result(result)

def reject_outgoing(self, request_id: int, error: Any) -> None:

record = self._outgoing.pop(request_id, None)

if record and not record.future.done():

record.future.set_exception(error)

def reject_all_outgoing(self, error: Any) -> None:

for record in self._outgoing.values():

if not record.future.done():

record.future.set_exception(error)

self._outgoing.clear()

def begin_incoming(self, method: str, params: Any) -> IncomingMessage:

record = IncomingMessage(method=method, params=params)

self._incoming.append(record)

return record

def complete_incoming(self, record: IncomingMessage, result: Any) -> None:

record.status = "completed"

record.result = result

def fail_incoming(self, record: IncomingMessage, error: Any) -> None:

record.status = "failed"

record.error = error