import logging
from threading import Thread

import zmq
import zmq.devices

logger = logging.getLogger("ZeroMqMessageHandler.QueueDevice")


class QueueDevice(object):
    def __init__(self,
                 context: zmq.Context,
                 front_address: str,
                 back_address: str,
                 min_port: int,
                 max_port: int):
        self.context = context
        self.front_address = front_address
        self.back_address = back_address
        self.min_port = min_port
        self.max_port = max_port

        self.running = False
        self.main_thread = None
        self.front_port = None
        self.front_socket = None
        self.back_socket = None
        self.workers = []

    def start(self):
        self.front_socket = self.context.socket(zmq.ROUTER)
        if self.front_address.startswith("tcp://") and self.front_address.find(":", 6) < 0:
            self.front_port = self.front_socket.bind_to_random_port(self.front_address, self.min_port, self.max_port)
        else:
            self.front_socket.bind(self.front_address)

        self.back_socket = self.context.socket(zmq.ROUTER)
        self.back_socket.bind(self.back_address)
        self.running = True
        self.main_thread = Thread(name="QueueDevice-{}".format(self.front_address), target=self.execute)
        self.main_thread.daemon = True
        self.main_thread.start()
        return self.front_port

    def stop(self):
        self.running = False
        self.main_thread.join()
        self.main_thread = None

    def execute(self):
        poller = zmq.Poller()
        # Only poll for requests from backend until workers are available
        poller.register(self.back_socket, zmq.POLLIN)

        while self.running:
            try:
                sockets = dict(poller.poll(1000))

                if self.back_socket in sockets:
                    # Handle worker activity on the backend
                    request = self.back_socket.recv_multipart()
                    worker, client = request[:2]
                    if client == b"DIED":
                        logger.debug("Worker {} DIED".format(worker))
                        if worker in self.workers:
                            self.workers.remove(worker)
                            if len(self.workers) == 0:
                                # Don't poll clients if no workers are available
                                logger.debug("No more workers, unregistering")
                                poller.unregister(self.front_socket)
                        else:
                            logger.debug("Worker {} DIED and was not in the list".format(worker))
                    else:
                        if client == b"DYING":
                            logger.debug("Reply from dying worker: {}".format(worker))
                            real_client = request[2]
                            reply = request[3]
                            self.front_socket.send_multipart([real_client, b"", reply])
                        else:
                            if client == b"READY":
                                logger.debug("New worker: {}".format(worker))
                            if len(self.workers) == 0:
                                # Poll for clients now that a worker is available
                                logger.debug("New worker available, registering again")
                                poller.register(self.front_socket, zmq.POLLIN)
                            self.workers.append(worker)
                            if client != b"READY" and len(request) >= 3:
                                # If client reply, send rest back to frontend
                                logger.debug("Reply from worker: {}".format(worker))
                                reply = request[2]
                                self.front_socket.send_multipart([client, b"", reply])

                if self.front_socket in sockets:
                    # Get next client request, route to last-used worker
                    client, empty, request = self.front_socket.recv_multipart()
                    try:
                        worker = self.workers.pop(0)
                    except IndexError:
                        # In exceptional cases, both sockets can have data and the back_end is a
                        # DIED message of the last worker
                        continue
                    logger.debug("Using worker: {}".format(worker))
                    self.back_socket.send_multipart([worker, client, request])
                    if len(self.workers) == 0:
                        # Don't poll clients if no workers are available
                        logger.debug("No more workers, unregistering")
                        poller.unregister(self.front_socket)
            except:
                logger.exception("Error on QueueDevice")
                raise
