import asyncio
from typing import Tuple, Optional, Callable, Sequence, Dict, Coroutine


class ZmqMessageHandler(object):
    def __init__(
        self,
        handler_map: Optional[Dict[bytes, Callable[[Sequence[bytes]], None]]] = None,
        fallback_handler: Optional[Callable[[Sequence[bytes]], None]] = None,
        type_index: int = 0,
    ):
        self._hm = handler_map or dict()
        self._fallback_h = fallback_handler or self._strict_fallback
        self._type_index = type_index

    @property
    def type_index(self):
        return self._type_index

    def add_handler(self, type_key: bytes, h: Callable[[Sequence[bytes]], None], allow_override: bool = False):
        if type_key in self._hm and not allow_override:
            raise Exception(f"Handler has a conflicting entry for {type_key}")
        self._hm[type_key] = h

    def handle(self, multipart_message: Tuple[bytes]):
        if multipart_message:
            handler = self._hm.get(multipart_message[self._type_index], self._fallback_h)
            handler(multipart_message)

    def _strict_fallback(self, multipart_message: Tuple[bytes]):
        raise Exception(f"Unknown message type for handler: {multipart_message[self._type_index]}")


class AioZmqMessageHandler(ZmqMessageHandler):
    def __init__(
        self,
        handler_map: Optional[Dict[bytes, Callable[[Sequence[bytes]], Optional[Coroutine]]]] = None,
        fallback_handler: Optional[Callable[[Sequence[bytes]], Optional[Coroutine]]] = None,
        type_index: int = 0,
    ):
        ZmqMessageHandler.__init__(self, handler_map, fallback_handler, type_index)

    async def handle(self, multipart_message: Sequence[bytes]):
        if multipart_message:
            handler = self._hm.get(multipart_message[self._type_index], self._fallback_h)
            if asyncio.iscoroutinefunction(handler):
                await handler(multipart_message)
            else:
                handler(multipart_message)


class QueuedAioZmqMessageHandler(AioZmqMessageHandler):
    """
    Queue coroutine handlers while executing function handlers as they come. This allows for coroutine handlers to
    wait on requests which response will arrive in the same handler loop.
    """
    def __init__(self, *args, **kwargs):
        AioZmqMessageHandler.__init__(self, *args, **kwargs)
        self._coro_q = asyncio.Queue()
        self._exec_t = None

    async def __aenter__(self):
        if self._exec_t:
            raise Exception("Queued executor is already running")
        self._exec_t = asyncio.create_task(self._queue_executor())

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        t = self._exec_t
        self._exec_t = None
        if t.done() and (e := t.exception()):
            if exc_val:
                raise e from exc_val
            raise e
        t.cancel()
        await asyncio.gather(t, return_exceptions=True)

    async def _queue_executor(self):
        while True:
            coro_handler = await self._coro_q.get()
            await coro_handler

    async def handle(self, multipart_message: Sequence[bytes]):
        if not self._exec_t or self._exec_t.done():
            raise Exception("Queued executor is not running")
        if multipart_message:
            handler = self._hm.get(multipart_message[self._type_index], self._fallback_h)
            if asyncio.iscoroutinefunction(handler):
                await self._coro_q.put(handler(multipart_message))
            else:
                handler(multipart_message)
