import logging
from threading import Thread
from typing import (
    Set,
)

import zmq
from messagebus import (
    LowLevelMessageBus,
    Message,
    MessageBusException,
    MessageBusExceptionType,
)
from ._LocalThreadData import LocalThreadData
from ._MessageDecoder import MessageDecoder

logger = logging.getLogger(__name__)


class ZeroMqLowLevelMessageBus(LowLevelMessageBus):
    def __init__(
        self,
        context: zmq.Context,
        port: int,
    ):
        self.context = context
        self.port = port

        self.local_thread_data = LocalThreadData(context=self.context)

        self.front_socket = FrontSocket(
            context=context,
            port=port,
        )
        self.front_socket.start()

    def send_direct_message(self, address: str, message: Message):
        socket = self.local_thread_data.open_socket(
            address=address,
            socket_type=zmq.REQ,
        )
        logger.debug("front_identity: %s", socket.get(zmq.IDENTITY))

        socket.send(MessageDecoder.format_message(message=message))
        if message.timeout >= 0:
            if not socket.poll(timeout=message.timeout, flags=zmq.POLLIN):
                self.local_thread_data.close_socket(socket=socket)
                raise MessageBusException(type=MessageBusExceptionType.timeout)

        data = socket.recv()
        return MessageDecoder.parse_message(data=data)

    def reply_message(self, message, reply_message):
        internal_socket = self.local_thread_data.open_socket(
            address=self.front_socket.address,
            socket_type=zmq.DEALER,
        )
        logger.debug("internal_identity: %s", internal_socket.get(zmq.IDENTITY))

        identity = getattr(message, "_identity")
        internal_socket.send_multipart(
            msg_parts=[b"REPLY", identity, MessageDecoder.format_message(message=reply_message)]
        )
        internal_socket.recv()
        delattr(message, "_identity")

    def receive(self, timeout: int = None):
        internal_socket = self.local_thread_data.open_socket(
            address=self.front_socket.address,
            socket_type=zmq.DEALER,
        )
        logger.debug("internal_identity: %s", internal_socket.get(zmq.IDENTITY))

        internal_socket.send_multipart(msg_parts=[b"RECEIVE"])

        if timeout is not None and not internal_socket.poll(timeout=timeout, flags=zmq.POLLIN):
            internal_socket.send_multipart(msg_parts=[b"TIMEOUT"])
            [identity, data] = internal_socket.recv_multipart()
            if data == b"":
                return None
            msg = MessageDecoder.parse_message(data)
            setattr(msg, "_identity", identity)
            return msg

        [identity, data] = internal_socket.recv_multipart()
        msg = MessageDecoder.parse_message(data)
        setattr(msg, "_identity", identity)
        return msg

    def close(self):
        self.front_socket.stop()
        self.front_socket.close()
        del self.local_thread_data

    def thread_done(self):
        self.local_thread_data.close_all_sockets()

    def _create_internal_socket(self):
        internal_socket: zmq.Socket = self.context.socket(zmq.DEALER)
        internal_socket.connect(addr=self.front_socket.address)
        return internal_socket


class FrontSocket(object):
    def __init__(self, context: zmq.Context, port: int):
        self._in_proc_address = "inproc://internal_socket"
        self.front_socket: zmq.Socket = context.socket(zmq.ROUTER)
        self.bind_address = f"tcp://*:{port}"
        self.front_socket.bind(self.bind_address)

        self.internal_socket: zmq.Socket = context.socket(zmq.ROUTER)
        self.internal_socket.bind(self._in_proc_address)

        self.running = True
        self.thread = Thread(
            target=self._run,
            name=f"FrontSocket_{self.bind_address}"
        )
        self.thread.daemon = True

    def start(self):
        self.running = True
        self.thread.start()

    def stop(self):
        self.running = False
        self.thread.join()
        self.thread = Thread(target=self._run)
        self.thread.daemon = True

    def close(self):
        self.front_socket.close(linger=0)
        self.internal_socket.close(linger=0)

    @property
    def address(self):
        return self._in_proc_address

    def _run(self):
        poller = zmq.Poller()
        poller.register(self.internal_socket, zmq.POLLIN)

        receivers = set()
        while self.running:
            try:
                sockets = poller.poll(timeout=10)

                discarded_identity = None
                if (self.internal_socket, zmq.POLLIN) in sockets:
                    discarded_identity = self._handle_internal_socket_data(
                        poller=poller,
                        receivers=receivers,
                    )

                if (self.front_socket, zmq.POLLIN) in sockets:
                    self._handle_front_socket_data(
                        poller=poller,
                        receivers=receivers,
                        discarded_identity=discarded_identity,
                    )

                elif discarded_identity is not None:
                    self.internal_socket.send_multipart(msg_parts=[discarded_identity, b"", b""])
            except (Exception,):
                logger.exception("Unhandled exception on front socket loop")

    def _handle_internal_socket_data(self, poller, receivers):
        discarded_identity = None
        things = self.internal_socket.recv_multipart()
        internal_identify = things[0]
        command = things[1]
        if command == b"RECEIVE":
            logger.debug("command: RECEIVE, internal_identify: %s", internal_identify)
            receivers.add(internal_identify)
            if len(receivers) == 1:
                poller.register(self.front_socket, zmq.POLLIN)

        elif command == b"REPLY":
            reply_identify = things[2]
            reply = things[3]
            logger.debug(
                "command: REPLY, internal_identity: %s, front_identity %s",
                internal_identify,
                reply_identify
            )
            self.front_socket.send_multipart(msg_parts=[reply_identify, b"", reply])
            self.internal_socket.send_multipart(msg_parts=[internal_identify, b""])

        elif command == b"TIMEOUT":
            logger.debug("command: TIMEOUT, internal_identity: %s", internal_identify)
            discarded_identity = internal_identify
            receivers.discard(internal_identify)
            if len(receivers) == 0:
                poller.unregister(self.front_socket)

        return discarded_identity

    def _handle_front_socket_data(
        self,
        poller: zmq.Poller,
        receivers: Set[bytes],
        discarded_identity: bytes,
    ):
        [identity, _, data] = self.front_socket.recv_multipart()
        if discarded_identity is not None:
            receiver_identify = discarded_identity
        else:
            receiver_identify = receivers.pop()
            if len(receivers) == 0:
                poller.unregister(self.front_socket)
        logger.debug(
            "command: DATA, front_identity: %s, internal_identity %s",
            identity,
            receiver_identify,
        )
        self.internal_socket.send_multipart([receiver_identify, identity, data])

