import zmq
import time
from threading import Thread, local
from queue import Queue

from ._msg import MSG, MSGType
from ._decorators import timeout
from ._general import logger
from .exceptions import ZMQPairTimeout


_ctx = local()


class ZMQPair(object):
    def __init__(self, dest_ip=None, port=8001):
        self._context = _get_context()
        self._socket = self._context.socket(zmq.PAIR)
        self._socket.setsockopt(zmq.LINGER, 0)
        self._dest_ip = dest_ip
        self._port = port
        self._connected = False
        self.__queue = Queue()
        self.__msg_backlog = {}
        self.__bg_worker = Thread(target=ZMQPair.__bg_receiver, args=(self.__queue, self._socket), daemon=True)
        self.__bg_worker.start()

    @property
    def context(self):
        return self._context

    @property
    def socket(self):
        return self._socket

    @property
    def dest_ip(self):
        return self._dest_ip

    @property
    def port(self):
        return self._port

    @property
    def connected(self):
        return self._connected

    def close(self):
        if(self._connected and self._dest_ip):
            self.disconnect()
        self._socket.close()

    def bind(self):
        self._connected = self.__heartbeat(0.1)
        if(not(self._connected)):
            self._socket.bind("tcp://*:{}".format(self._port))
            self._connected = True

    def unbind(self):
        if(self._connected):
            self._socket.unbind()
            self._connected = False

    def connect(self):
        self._connected = self.__heartbeat(0.1)
        if(not(self._connected)):
            self._socket.connect("tcp://{}:{}".format(self._dest_ip, self._port))
            self._connected = True
            logger.debug("Main:: Connected: {}:{}".format(self._dest_ip, self._port))

    def disconnect(self):
        self._socket.disconnect("tcp://{}:{}".format(self._dest_ip, self._port))
        self._connected = False
        logger.debug("Main:: Disconnect: {}:{}".format(self._dest_ip, self._port))

    def send_msg(self, msg, timeout=None):
        assert type(msg) == MSG, "ZMQPair only communicates through MSG class"
        if(self._connected):
            logger.debug("Main:: Sending Message {}".format(str(msg)))
            self.__send_process(msg, timeout=timeout, exception=ZMQPairTimeout)
            return True
        else:
            self._connected = False
            return False

    def recv_msg(self, timeout=None, msg_id=None):
        if(msg_id is not None and msg_id in self.__msg_backlog and self.__msg_backlog[msg_id]):
            return self.__msg_backlog[msg_id].pop(0)
        self.__check_mail(timeout=timeout, exception=ZMQPairTimeout)
        mail = self.__queue.get(timeout=10)
        logger.debug("Main:: Retrieved Message: {}".format(str(mail)))
        return mail

    def requeue_msg(self, msg):
        if(not(msg.type == MSGType.ACKNOWLEDGE and msg.has_request_id())):
            self.__msg_backlog.setdefault(msg.request_id, list()).append(msg)
            logger.debug("Putting MSG to backlog {}".format(msg))
        else:
            self.__queue.put(msg)

    def pulse(self, timeout=None):
        return self.__heartbeat(timeout=timeout)

    @timeout(name="Trasmitter")
    def __send_process(self, msg, **kwargs):
        tracker = self._socket.send_pyobj(msg, flags=zmq.NOBLOCK, track=True, copy=False)
        while(not(tracker.done)):
            yield None

    @timeout(name="Receiver")
    def __check_mail(self, **kwargs):
        if(self.__bg_worker.is_alive()):
            while(self.__queue.empty()):
                yield None
        else:
            raise RuntimeError("Recv thread died?")

    def __heartbeat(self, timeout):
        if(self._connected):
            self.send_msg(MSG(MSGType.HEARTBEAT, request=True), timeout=timeout)
            try:
                while(True):
                    msg = self.recv_msg(timeout=timeout)
                    if(msg.type == MSGType.HEARTBEAT):
                        return True
                    else:
                        self.__queue.put(msg)
            except ZMQPairTimeout:
                return False
        else:
            return False

    def __bg_receiver(queue, socket):
        logger.debug("Thread:: Starting")
        try:
            while(True):
                while(not(socket.poll(timeout=1, flags=zmq.POLLIN))):
                    time.sleep(0)
                msg = socket.recv_pyobj()
                if(msg.type == MSGType.ACKNOWLEDGE and msg.request_id == 0):  # if server reset, clear out older msgs
                    queue_return = []
                    while(not(queue.empty())):
                        msg = queue.get()
                        if(msg.type == MSGType.HEARTBEAT or msg.type == MSGType.SYNC):
                            queue_return.append(msg)
                    while(queue_return):
                        queue.put(queue_return.pop(0))
                logger.debug("Thread:: {}".format(str(msg)))
                if(msg.request and msg.type == MSGType.HEARTBEAT):
                    socket.send_pyobj(MSG(MSGType.HEARTBEAT, request=False))
                else:
                    queue.put(msg)
        except zmq.ZMQError:
            pass


def _get_context():
    try:
        return getattr(_ctx, "zmq_context")
    except (AttributeError, IndexError):
        _ctx.__dict__["zmq_context"] = zmq.Context()
        return _ctx.zmq_context


def _close_context():
    try:
        getattr(_ctx, "zmq_context").term()
    except (AttributeError, IndexError):
        pass  # No need to close context?
