import logging
import time
from threading import Thread

import zmq
from zmq import Context

logger = logging.getLogger("ZeroMqEventManager.StreamerDevice")


class StreamerDevice(object):
    def __init__(
        self,
        context: Context,
        sub_address: str,
        push_address: str,
        min_port: int,
        max_port: int,
        streamer_poll_timeout: int,
    ) -> None:
        self.context = context
        self.sub_address = sub_address
        self.push_address = push_address
        self.min_port = min_port
        self.max_port = max_port
        self.streamer_poll_timeout = streamer_poll_timeout
        self.port = None
        self.sub_socket = None
        self.push_socket = None
        self.running = False
        self.thread = None
        self.workers = None

    def start(self):
        self.sub_socket = self.context.socket(zmq.SUB)
        self.sub_socket.set_hwm(100000)
        self.sub_socket.connect(self.sub_address)

        self.push_socket = self.context.socket(zmq.ROUTER)
        self.push_socket.set_hwm(100000)
        self.port = self.push_socket.bind_to_random_port(self.push_address, self.min_port, self.max_port)
        logger.info("StreamerPort: {}".format(self.port))

        self.workers = []
        self.running = True
        thread_name = "StreamerDevice_{}:{}".format(self.push_address, self.port).format()
        self.thread = Thread(name=thread_name, target=self.execute)
        self.thread.daemon = True
        self.thread.start()

        return self.port

    def stop(self):
        self.running = False
        self.thread.join()
        self.push_socket.close()
        self.sub_socket.close()

    def add_subjects(self, subjects):
        for subject in subjects:
            self.sub_socket.setsockopt_string(zmq.SUBSCRIBE, subject + " ")

    def remove_subjects(self, subjects):
        for subject in subjects:
            self.sub_socket.setsockopt_string(zmq.UNSUBSCRIBE, subject + " ")

    def execute(self):
        logger.info("Starting streamer device: %s %s", self.sub_address, self.push_address)

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

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

                if self.push_socket in sockets:
                    # Handle worker activity on the backend
                    request = self.push_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.sub_socket)
                        else:
                            logger.debug("Worker {} DIED and was not in the list".format(worker))
                    else:
                        if client == b"DYING":
                            logger.debug("Dying worker: {}".format(worker))
                        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.sub_socket, zmq.POLLIN)
                            self.workers.append(worker)

                if self.sub_socket in sockets:
                    # Get next client request, route to last-used worker
                    request = self.sub_socket.recv()
                    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.push_socket.send_multipart([worker, b"", request])
                    if len(self.workers) == 0:
                        # Don't poll clients if no workers are available
                        logger.debug("No more workers, unregistering")
                        poller.unregister(self.sub_socket)
            except (Exception,):
                logger.error("Error on StreamingDevice")
                raise
