import asyncio
from queue import Queue, Empty
from typing import Sequence

import zmq

from clarabot_nano.common.zmq.constants import QUEUE_PUSH_NOTIFICATION


class InprocClient(object):
    """
    Inproc communication client with a zmq-connection handler thread. In additional to the socket connection a
    notification-bound queue is also prepared. This can be used to avoid copying over massive amounts of data
    between threads.
    The send behaviour is to not block, because that state would indicate thread failure and waiting is no good.
    """

    def __init__(self, queue: Queue, inproc_connect: str, auto_queue_above: int = zmq.COPY_THRESHOLD):
        self._auto_queue_above = auto_queue_above
        self._queue = queue
        self.socket = zmq.Context.instance().socket(zmq.PAIR)
        self.socket.connect(inproc_connect)
        self.pollin = zmq.Poller()
        self.pollin.register(self.socket, zmq.POLLIN)

    def send_multipart(self, multipart_message: Sequence[bytes]):
        if self._auto_queue_above and sum(len(p) for p in multipart_message) > self._auto_queue_above:
            self.send_queued_multipart(multipart_message)
        else:
            self.socket.send_multipart(multipart_message, flags=zmq.NOBLOCK)

    def send_queued_multipart(self, multipart_message: Sequence[bytes]):
        self._queue.put(multipart_message)
        self.socket.send_multipart((QUEUE_PUSH_NOTIFICATION, QUEUE_PUSH_NOTIFICATION), flags=zmq.NOBLOCK)

    def close(self, linger=250):
        self.socket.close(linger)
        self.socket = self.pollin = None


class AioInprocClient(InprocClient):
    """
    Asyncio compatibility.
    """
    def __init__(
        self,
        receive_queue: Queue,
        receive_event: asyncio.Event,
        send_queue: Queue,
        inproc_connect: str,
        auto_queue_above: int = zmq.COPY_THRESHOLD,
    ):
        InprocClient.__init__(self, send_queue, inproc_connect, auto_queue_above=auto_queue_above)
        self._receive_queue = receive_queue  # TODO: queue is consumed whole so it could be a lock-protected list
        self._receive_event = receive_event

    async def recv_multipart_messages(self) -> Sequence[Sequence[bytes]]:
        await self._receive_event.wait()
        # Because the asyncio Event is reset before the processing of the queue, it is guaranteed that the event will
        # again be set *later* than the last message becomes available through the queue. If the last message is
        # missed by a race-condition the event will trigger consumption. The race condition can result in the event
        # being set while the last message get consumed already. So the consume logic here ultimately cannot miss out
        # on a message, in the worst case scenario it may look for a message and find none left.
        self._receive_event.clear()
        messages = []
        while True:
            try:
                messages.append(self._receive_queue.get(block=False))
            except Empty:
                break
        return messages
