import logging
import time

import zmq
from abc import ABCMeta, abstractmethod
from threading import Thread


class WorkerThread(object):
    __meta__ = ABCMeta

    def __init__(self, context, worker_address, idle_timeout, identity):
        # type: (zmq.Context, str, int, str) -> None
        self.context = context
        self.worker_address = worker_address
        self.idle_timeout = idle_timeout
        self.identity = identity

        self.logger = logging.getLogger(self.get_logger_name())

        self.is_working = False
        self.is_idle = False

        self.running = False
        self.inactive_time = 0
        self.thread = None  # type: Thread

    def start(self):
        if self.running:
            return

        self.running = True
        self.inactive_time = 0

        self.thread = Thread(target=self.run, name=self.identity)
        self.thread.daemon = True
        self.thread.start()

    def stop(self):
        self.running = False

    def join(self):
        self.thread.join()

    def get_id(self):
        if self.thread is None:
            return None
        return self.thread.ident

    def run(self):
        socket = None
        flush_count = 0

        try:
            self.logger.debug("Creating socket to address: %s", self.worker_address)
            socket = self.create_socket(self.worker_address)
            socket.send(b"READY")

            receiving = False
            while self.running or flush_count < 3:
                if not self.running:
                    if flush_count == 0:
                        socket.send(b"DIED")
                    flush_count += 1

                if receiving:
                    try:
                        address, data = socket.recv_multipart(flags=zmq.NOBLOCK)
                    except zmq.ZMQError:
                        receiving = False
                        continue
                else:
                    ready = socket.poll(1000, zmq.POLLIN)
                    if ready & zmq.POLLIN == 0:
                        self.is_working = False
                        self.inactive_time += 1000
                        if self.inactive_time > self.idle_timeout:
                            self.is_idle = True

                        continue
                    else:
                        self.logger.debug("Data available, start receiving: %s", type(self))
                        receiving = True
                        address, data = socket.recv_multipart()
                        self.logger.debug("Data received")

                self.inactive_time = 0
                self.is_idle = False
                self.is_working = True

                self.logger.debug("Doing work")
                resp = self.do_work(data)
                if resp is None:
                    resp = b""
                self.logger.debug("Work done")

                self.logger.debug("Sending reply")
                if self.running:
                    socket.send_multipart([address, resp])
                else:
                    socket.send_multipart([b"DYING", address, resp])
                self.logger.debug("Reply sent")
        except:
            self.logger.exception("Error handling work")
            raise
        finally:
            self.running = False
            try:
                if flush_count == 0 and socket is not None:
                    try:
                        socket.send(b"DIED")
                    except:
                        pass
                if socket is not None:
                    socket.close(linger=1000)
            except:
                pass

    @abstractmethod
    def do_work(self, data):
        # type: (bytes) -> bytes
        raise NotImplementedError()

    @abstractmethod
    def dispose(self):
        raise NotImplementedError()

    def create_socket(self, address):
        socket = self.context.socket(zmq.DEALER)
        socket.identity = "Worker-{}".format(self.identity).encode("utf-8")
        socket.connect(address)
        time.sleep(0.1)
        return socket

    def get_logger_name(self):
        return "ZeroMqThreadPool.WorkerThread"
