import zmq
import zmq.auth
import threading

from pymoku.network import Socket, SocketFactory, FrameSocket
from pymoku.network import NetworkError, FrameTimeout

SERVER_KEY = b"G8&2Q3?P%G96BqbQ=H>87cfHia4R{nK%/NoZjC8L"


class ZMQReq(Socket):
    def __init__(self, conn):
        self.conn = conn
        self._lock = threading.RLock()

    def transfer(self, data, timeout=3000):
        self._cancel = False
        with self._lock:
            self.conn.send(data)
            for _ in range(timeout // 100):
                if self.conn.poll(100):
                    return self.conn.recv()
                if self._cancel:
                    return
            else:
                raise NetworkError('Socket timeout.')

    def close(self):
        try:
            self._cancel = True
            with self._lock:
                self.conn.close()
        except Exception:
            pass


class ZMQSub(FrameSocket):
    def __init__(self, conn):
        self.conn = conn

    def get_data(self, timeout=2.0):
        if self.conn.poll(timeout * 1000):
            data = self.conn.recv_multipart()[-1]
        else:
            raise FrameTimeout("Timed out waiting on instrument data.")

        return data

    def close(self):
        self.conn.close()


class ZMQCurveSocketFactory(SocketFactory):
    def __init__(self, ip_addr):
        self.ip_addr = ip_addr

    def _socket(self, port, curve=False):
        ctx = zmq.Context.instance()
        conn = ctx.socket(zmq.REQ)
        if curve:
            conn.curve_publickey, \
                conn.curve_secretkey = zmq.curve_keypair()
            conn.curve_serverkey = SERVER_KEY
        conn.setsockopt(zmq.IPV6, True)
        conn.setsockopt(zmq.LINGER, 5000)
        conn.setsockopt(zmq.REQ_RELAXED, 1)
        conn.setsockopt(zmq.CONNECT_TIMEOUT, 5000)
        conn.setsockopt(zmq.TCP_KEEPALIVE, 1)
        conn.setsockopt(zmq.TCP_KEEPALIVE_IDLE, 5)
        conn.setsockopt(zmq.TCP_KEEPALIVE_CNT, 2)
        conn.setsockopt(zmq.TCP_KEEPALIVE_INTVL, 5)

        conn.connect(f'tcp://{self.ip_addr}:{port:d}')
        return ZMQReq(conn)

    def get_control_socket(self):
        return self._socket(27184, curve=True)

    def get_render_socket(self):
        return self._socket(27189)

    def get_frame_socket(self, name):
        ctx = zmq.Context.instance()
        conn = ctx.socket(zmq.SUB)
        conn.setsockopt(zmq.IPV6, True)
        conn.setsockopt(zmq.LINGER, 0)
        conn.setsockopt(zmq.RCVHWM, 1)
        conn.setsockopt(zmq.CONNECT_TIMEOUT, 5000)
        conn.setsockopt(zmq.TCP_KEEPALIVE, 1)
        conn.setsockopt(zmq.TCP_KEEPALIVE_IDLE, 1)
        conn.setsockopt(zmq.TCP_KEEPALIVE_CNT, 10)
        conn.setsockopt(zmq.TCP_KEEPALIVE_INTVL, 1)

        conn.connect(f'tcp://{self.ip_addr}:27185')
        conn.subscribe(name.encode())
        return ZMQSub(conn)

    def get_notification_socket(self):
        ctx = zmq.Context.instance()
        conn = ctx.socket(zmq.SUB)
        conn.setsockopt(zmq.IPV6, True)
        conn.setsockopt(zmq.LINGER, 0)
        conn.setsockopt(zmq.CONNECT_TIMEOUT, 5000)
        conn.setsockopt(zmq.TCP_KEEPALIVE, 1)
        conn.setsockopt(zmq.TCP_KEEPALIVE_IDLE, 1)
        conn.setsockopt(zmq.TCP_KEEPALIVE_CNT, 10)
        conn.setsockopt(zmq.TCP_KEEPALIVE_INTVL, 1)

        conn.connect(f'tcp://{self.ip_addr}:27183')
        conn.subscribe(b'')
        return ZMQSub(conn)


class ZMQIPCSocketFactory(SocketFactory):
    def __init__(self, *args):
        pass

    def _ipc_socket(self, path, curve=False):
        ctx = zmq.Context.instance()
        conn = ctx.socket(zmq.REQ)
        if curve:  # TODO not for IPC
            conn.curve_publickey, \
                conn.curve_secretkey = zmq.curve_keypair()
            conn.curve_serverkey = SERVER_KEY
        conn.setsockopt(zmq.REQ_RELAXED, 1)
        conn.setsockopt(zmq.CONNECT_TIMEOUT, 5000)
        conn.connect(path)
        return ZMQReq(conn)

    def get_control_socket(self):
        return self._ipc_socket('ipc:///var/run/control_req')

    def get_render_socket(self):
        return self._ipc_socket('ipc:///var/run/render')

    def get_frame_socket(self, name):
        ctx = zmq.Context.instance()
        conn = ctx.socket(zmq.SUB)
        conn.setsockopt(zmq.RCVTIMEO, 3000)
        conn.setsockopt(zmq.RCVHWM, 1)
        conn.connect('ipc:///var/run/frame_pub')
        conn.subscribe(name.encode())
        return ZMQSub(conn)
