class DeviceRep(object):
    """
    Helper class to represent a device to a worker
    """

    def __init__(self, device_id, state='unknown'):
        self.id = device_id
        self.state = state
        self.curr_liveness = CLIENT_HB_LIVENESS
        self.hb_timer = PeriodicCallback(self.heartbeat, CLIENT_HB_INTERVAL)
        self.hb_timer.start()
        return

    def heartbeat(self):
        if self.curr_liveness > 0:
            self.curr_liveness -= 1
        if self.curr_liveness == 0:
            self.state = 'dead'
        return

    def on_message_received(self):
        self.curr_liveness = CLIENT_HB_LIVENESS
        return

    def is_alive(self):
        return self.curr_liveness > 0

    def get_state(self):
        return self.state

    def shutdown(self):
        self.hb_timer.stop()
        self.hb_timer = None

def main(pat):
    
    fname = find_connection_file(pat)
    with open(fname) as f:
        cfg = json.load(f)
    
    url = "%s://%s:%s" % (cfg.get('transport', 'tcp'), cfg['ip'], cfg['iopub_port'])
    
    session = Session(key=cfg['key'])
    
    ctx = zmq.Context.instance()
    sub = ctx.socket(zmq.SUB)
    sub.subscribe = b''
    sub.connect(url)
    # import IPython
    # IPython.embed()
    # return
    
    stream = ZMQStream(sub)
    
    stream.on_recv(lambda msg_list: log_msg(session, msg_list))
    
    pc = PeriodicCallback(print_time, 5 * 60 * 1000)
    pc.start()
    IOLoop.instance().start()

    def __init__(self, settings, identity, insocket, outsocket, mgmt, frontier,
            log_handler, log_level, io_loop):
        """
        Initialize the master.
        """
        LoggingMixin.__init__(self, log_handler, log_level)
        self._identity = identity
        self._io_loop = io_loop or IOLoop.instance()

        self._in_stream = ZMQStream(insocket, io_loop)
        self._out_stream = ZMQStream(outsocket, io_loop)

        self._mgmt = mgmt
        self._frontier = frontier

        self._running = False
        self._available_workers = []

        # periodically check if there are pending URIs to crawl
        self._periodic_update = PeriodicCallback(self._send_next_uri,
                settings.MASTER_PERIODIC_UPDATE_INTERVAL, io_loop=io_loop)
        # start this periodic callback when you are waiting for the workers to
        # finish
        self._periodic_shutdown = PeriodicCallback(self._shutdown_wait, 500,
                io_loop=io_loop)
        self._shutdown_counter = 0
        self._logger.debug("zmqmaster::initialized")

    def __init__(self, frontier,
            data_in_sock='ipc:///tmp/robot-data-w2m.sock',
            data_out_sock='ipc:///tmp/robot-data-m2w.sock',
            msg_in_sock='ipc:///tmp/robot-msg-w2m.sock',
            msg_out_sock='ipc:///tmp/robot-msg-m2w.sock',
            io_loop=None):
        self.identity = 'master:%s:%s' % (socket.gethostname(), os.getpid())

        context = zmq.Context()

        self._io_loop = io_loop or IOLoop.instance()

        self._in_socket = context.socket(zmq.SUB)
        self._in_socket.setsockopt(zmq.SUBSCRIBE, '')
        self._in_socket.bind(data_in_sock)
        self._in_stream = ZMQStream(self._in_socket, io_loop)

        self._out_socket = context.socket(zmq.PUSH)
        self._out_socket.bind(data_out_sock)
        self._out_stream = ZMQStream(self._out_socket, io_loop)

        self._online_workers = set()
        self._running = False

        self._updater = PeriodicCallback(self._send_next, 100, io_loop=io_loop)
        self._reloader = PeriodicCallback(self.reload, 1000, io_loop=io_loop)

        self.frontier = frontier
        self.messenger = ServerMessenger(msg_in_sock, msg_out_sock,
                context, io_loop)

class WorkerRep(object):

    """Helper class to represent a worker in the broker.

    Instances of this class are used to track the state of the attached worker
    and carry the timers for incomming and outgoing heartbeats.

    :type wid:       str
    :param wid:      the worker id.
    :param service:  service this worker serves
    :type service:   str
    :param stream:   the ZMQStream used to send messages
    :type stream:    ZMQStream
    """

    def __init__(self, wid, service, stream):
        self.id = wid
        self.service = service
        self.multicasts = []
        self.curr_liveness = HB_LIVENESS
        self.stream = stream
        self.last_hb = 0
        self.hb_out_timer = PeriodicCallback(self.send_hb, HB_INTERVAL)
        self.hb_out_timer.start()
        return

    def send_hb(self):
        """Called on every HB_INTERVAL.

        Decrements the current liveness by one.

        Sends heartbeat to worker.
        """
        self.curr_liveness -= 1
        msg = [ self.id, b'', MDP_WORKER_VERSION, b'\x05' ]
        self.stream.send_multipart(msg)
        return

    def on_heartbeat(self):
        """Called when a heartbeat message from the worker was received.

        Sets current liveness to HB_LIVENESS.
        """
        self.curr_liveness = HB_LIVENESS
        return

    def is_alive(self):
        """Returns True when the worker is considered alive.
        """
        return self.curr_liveness > 0

    def shutdown(self):
        """Cleanup worker.

        Stops timer.
        """
        self.hb_out_timer.stop()
        self.hb_out_timer = None
        self.stream = None
        return

 def run(self):
     _period = 1000 * self.configreader["services"][self.name]["checkmail_interval"]
     beat = PeriodicCallback(self.cmd_processmails,
                             _period,
                             io_loop=self.ioloop)
     beat.start()
     Service.run(self)

class BufferModule(base.ZmqProcess):

    def __init__(self, recv_addr, send_addr, recv_title, send_title, state_handler, period=1000):
        super().__init__()
        self.sub_stream    = None
        self.timer         = None        
        self.recv_addr     = recv_addr
        self.recv_title    = recv_title
        self.state_handler = state_handler
        self.period        = period
        self.callback      = Callback(Publisher(send_addr, send_title), self.state_handler)
        
    def setup(self):
        super().setup() 
        self.sub_stream, _ = self.stream(zmq.SUB, self.recv_addr, bind=False, subscribe=self.recv_title.encode('utf-8'))
        self.sub_stream.on_recv(SubStreamHandler(self.sub_stream, self.stop, self.state_handler))
        self.timer = PeriodicCallback(self.callback, self.period, self.loop)

    def run(self):
        self.setup()
        print('Start loop!')
        self.timer.start()        
        self.loop.start()

    def stop(self):
        self.loop.stop()

class TaskState (object):
    """ Tracks task state (with help of watchdog) """

    log = skytools.getLogger ('d:TaskState')

    def __init__ (self, uid, name, info, ioloop, cc, xtx):
        self.uid = uid
        self.name = name
        self.info = info
        self.pidfile = info['config']['pidfile']
        self.ioloop = ioloop
        self.cc = cc
        self.xtx = xtx
        self.timer = None
        self.timer_tick = 1
        self.heartbeat = False
        self.start_time = None
        self.dead_since = None

    def start (self):
        self.start_time = time.time()
        self.timer = PeriodicCallback (self.watchdog, self.timer_tick * 1000, self.ioloop)
        self.timer.start()

    def stop (self):
        try:
            self.log.info ('Signalling %s', self.name)
            skytools.signal_pidfile (self.pidfile, signal.SIGINT)
        except:
            self.log.exception ('signal_pidfile failed: %s', self.pidfile)

    def watchdog (self):
        live = skytools.signal_pidfile (self.pidfile, 0)
        if live:
            self.log.debug ('%s is alive', self.name)
            if self.heartbeat:
                self.send_reply ('running')
        else:
            self.log.info ('%s is over', self.name)
            self.dead_since = time.time()
            self.timer.stop()
            self.timer = None
            self.send_reply ('stopped')

    def ccpublish (self, msg):
        assert isinstance (msg, TaskReplyMessage)
        cmsg = self.xtx.create_cmsg (msg)
        cmsg.send_to (self.cc)

    def send_reply (self, status, feedback = {}):
        msg = TaskReplyMessage(
                req = 'task.reply.%s' % self.uid,
                handler = self.info['task']['task_handler'],
                task_id = self.info['task']['task_id'],
                status = status,
                feedback = feedback)
        self.ccpublish (msg)

    def loop():
      self.__ioloop = IOLoop()
      queue_socket = self.__context.socket(zmq.PULL)
      queue_socket.connect(self.message_address)
      queue_stream = ZMQStream(queue_socket, self.__ioloop)
      worker_socket = self.__context.socket(zmq.DEALER)
      for address in self.active_connections:
        worker_socket.connect(address)
      worker_stream = ZMQStream(worker_socket, self.__ioloop)

      def receive_response(message, response_override=None):
        self.__queued_messages.pop(message[1], None)
        self.__message_timeouts.pop(message[1], None)
        callback = self.__callbacks.pop(message[1], None)
        if callback:
          try:
            callback(response_override or self.loads(self.decompress(message[2])))
          except Exception as e:
            self.log_error(e)
            callback({'error': e})
      worker_stream.on_recv(receive_response)

      def queue_message(message):
        if message[0]:
          if message[0] == WORKER_SOCKET_CONNECT and message[2] not in self.active_connections:
            self.active_connections.add(message[2])
            worker_stream.socket.connect(message[2])
          elif message[0] == WORKER_SOCKET_DISCONNECT and message[2] in self.active_connections:
            self.active_connections.remove(message[2])
            worker_stream.socket.disconnect(message[2])
          return
        self.__queued_messages[message[1]] = (time(), message)
        try:
          worker_stream.send_multipart(message)
        except Exception as e:
          self.log_error(e)
      queue_stream.on_recv(queue_message)

      def timeout_message():
        now = time()
        for message, retry in [(item[1], self.__message_auto_retry.get(item[1][1], self.__auto_retry)) for item, t in ((i, self.__message_timeouts.get(i[1][1], self.__timeout)) for i in self.__queued_messages.itervalues()) if t >= 0 and (item[0] + t < now)]:
          if retry:
            logging.info('Worker timeout, requeuing ' + message[1])
            queue_message(message)
          else:
            receive_response(('', message[1]), {'error': 'timeout'})
      timeout_callback = PeriodicCallback(timeout_message, int(abs(self.__timeout * 1000.0)), io_loop = self.__ioloop)
      timeout_callback.start()

      self.__ioloop.start()
      self.__thread = None

 def start(self):
     loop = self.loop
     loop.add_handler(self.udp.handle.fileno(), self.handle_beacon, loop.READ)
     stream = ZMQStream(self.pipe, loop)
     stream.on_recv(self.control_message)
     pc = PeriodicCallback(self.send_ping, PING_INTERVAL * 1000, loop)
     pc.start()
     pc = PeriodicCallback(self.reap_peers, PING_INTERVAL * 1000, loop)
     pc.start()
     pc = PeriodicCallback(self.send_ports, PING_INTERVAL * 1000, loop)
     pc.start()
     loop.start()

    def __init__(self, context, main_ep, opt_ep=None):
        """Init MDPBroker instance.
        """
        l = logger.Logger('mq_broker')
        self.log = l.get_logger()
        self.log.info("MDP broker startup...")

        socket = ZmqSocket(context, zmq.ROUTER)
        socket.bind(main_ep)
        self.main_stream = ZMQStream(socket)
        self.main_stream.on_recv(self.on_message)
        if opt_ep:
            socket = ZmqSocket(context, zmq.ROUTER)
            socket.bind(opt_ep)
            self.client_stream = ZMQStream(socket)
            self.client_stream.on_recv(self.on_message)
        else:
            self.client_stream = self.main_stream
        self.log.debug("Socket created...")
        self._workers = {}
        # services contain the worker queue and the request queue
        self._services = {}
        self._worker_cmds = { b'\x01': self.on_ready,
                              b'\x03': self.on_reply,
                              b'\x04': self.on_heartbeat,
                              b'\x05': self.on_disconnect,
                              }
        self.log.debug("Launch the timer...")
        self.hb_check_timer = PeriodicCallback(self.on_timer, HB_INTERVAL)
        self.hb_check_timer.start()
        self.log.info("MDP broker started")
        return

    def __init__(self, broker, service, io_loop=None):
        """Create and setup an MDP worker.
           @param broker A string containing the broker's URL
           @param service A string containing the service name
           @param io_loop An existing I/O loop object. If None, the default will be used.
        """
        self.service=service
        self._broker = broker

        self.ctx = zmq.Context()
        sock = self.ctx.socket(zmq.DEALER)
        ZMQStream.__init__(self, sock, io_loop)
        # last watchdog timer tick
        self.watchdog = 0
        # connection callback one-shot
        self._conncb = DelayedCallback(self.send_ready, 3000, self.io_loop)
        # heartbeat callback..runs continuous when connected
        self._hbcb = PeriodicCallback(self.send_heartbeat, 2000, self.io_loop)
        # number of connection attempts
        self._conn_attempt = 0
        # waiting to connect state
        self._waiting_to_connect = True
        # have we been disconnected? (flags a reconnect attempt)
        self.disconnected = False

        # connect the socket and send a READY when the io_loop starts
        self.connect(self._broker)
        self._conncb.start()

 def __init__(self, device_id, state='unknown'):
     self.id = device_id
     self.state = state
     self.curr_liveness = CLIENT_HB_LIVENESS
     self.hb_timer = PeriodicCallback(self.heartbeat, CLIENT_HB_INTERVAL)
     self.hb_timer.start()
     return

    def __init__(self, port=5556):
        self.port = port
        self.ctx = zmq.Context()
        self.kvmap = {}
        self.loop = IOLoop.instance()

        # Set up our clone server sockets
        self.snapshot  = self.ctx.socket(zmq.ROUTER)
        self.publisher = self.ctx.socket(zmq.PUB)
        self.collector = self.ctx.socket(zmq.PULL)
        self.snapshot.bind("tcp://*:%d" % self.port)
        self.publisher.bind("tcp://*:%d" % (self.port + 1))
        self.collector.bind("tcp://*:%d" % (self.port + 2))

        # Wrap sockets in ZMQStreams for IOLoop handlers
        self.snapshot = ZMQStream(self.snapshot)
        self.publisher = ZMQStream(self.publisher)
        self.collector = ZMQStream(self.collector)

        # Register our handlers with reactor
        self.snapshot.on_recv(self.handle_snapshot)
        self.collector.on_recv(self.handle_collect)
        self.flush_callback = PeriodicCallback(self.flush_ttl, 1000)

        # basic log formatting:
        logging.basicConfig(format="%(asctime)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S",
                level=logging.INFO)

    def __init__(self, primary, local, remote):
        # initialize the Binary Star
        self.ctx = zmq.Context()  # Our private context
        self.loop = IOLoop.instance()  # Reactor loop
        self.state = STATE_PRIMARY if primary else STATE_BACKUP

        self.event = None  # Current event
        self.peer_expiry = 0  # When peer is considered 'dead'
        self.voter_callback = None  # Voting socket handler
        self.master_callback = None  # Call when become master
        self.slave_callback = None  # Call when become slave

        # Create publisher for state going to peer
        self.statepub = self.ctx.socket(zmq.PUB)
        self.statepub.bind(local)

        # Create subscriber for state coming from peer
        self.statesub = self.ctx.socket(zmq.SUB)
        self.statesub.setsockopt_string(zmq.SUBSCRIBE, u'')
        self.statesub.connect(remote)

        # wrap statesub in ZMQStream for event triggers
        self.statesub = ZMQStream(self.statesub, self.loop)

        # setup basic reactor events
        self.heartbeat = PeriodicCallback(self.send_state,
                                          HEARTBEAT, self.loop)
        self.statesub.on_recv(self.recv_state)

 def __init__(self, context, main_ep, opt_ep=None, worker_q=None):
     """Init MDPBroker instance.
     """
     socket = context.socket(zmq.XREP)
     socket.bind(main_ep)
     self.main_stream = ZMQStream(socket)
     self.main_stream.on_recv(self.on_message)
     if opt_ep:
         socket = context.socket(zmq.XREP)
         socket.bind(opt_ep)
         self.client_stream = ZMQStream(socket)
         self.client_stream.on_recv(self.on_message)
     else:
         self.client_stream = self.main_stream
     self._workers = {}
     # services contain the worker queue and the request queue
     self._services = {}
     self._worker_cmds = { '\x01': self.on_ready,
                           '\x03': self.on_reply,
                           '\x04': self.on_heartbeat,
                           '\x05': self.on_disconnect,
                           }
     self.hb_check_timer = PeriodicCallback(self.on_timer, HB_INTERVAL)
     self.hb_check_timer.start()
     return

    def __init__ (self, hname, hcf, ccscript):
        super(TailWriter, self).__init__(hname, hcf, ccscript)

        self.files = {}
        self.workers = []
        self.wparams = {} # passed to workers

        self.wparams['dstdir'] = self.cf.getfile ('dstdir')
        self.wparams['host_subdirs'] = self.cf.getbool ('host-subdirs', 0)
        self.wparams['maint_period'] = self.cf.getint ('maint-period', 3)
        self.wparams['write_compressed'] = self.cf.get ('write-compressed', '')
        assert self.wparams['write_compressed'] in [None, '', 'no', 'keep', 'yes']
        if self.wparams['write_compressed'] in ('keep', 'yes'):
            self.log.info ("position checking not supported for compressed files")
        if self.wparams['write_compressed'] == 'yes':
            self.wparams['compression'] = self.cf.get ('compression', '')
            if self.wparams['compression'] not in ('gzip', 'bzip2'):
                self.log.error ("unsupported compression: %s", self.wparams['compression'])
            self.wparams['compression_level'] = self.cf.getint ('compression-level', '')
            self.wparams['buf_maxbytes'] = cc.util.hsize_to_bytes (self.cf.get ('buffer-bytes', '1 MB'))
            if self.wparams['buf_maxbytes'] < BUF_MINBYTES:
                self.log.info ("buffer-bytes too low, adjusting: %i -> %i", self.wparams['buf_maxbytes'], BUF_MINBYTES)
                self.wparams['buf_maxbytes'] = BUF_MINBYTES

        # initialise sockets for communication with workers
        self.dealer_stream, self.dealer_url = self.init_comm (zmq.XREQ, 'inproc://workers-dealer', self.dealer_on_recv)
        self.router_stream, self.router_url = self.init_comm (zmq.XREP, 'inproc://workers-router', self.router_on_recv)

        self.launch_workers()

        self.timer_maint = PeriodicCallback (self.do_maint, self.wparams['maint_period'] * 1000, self.ioloop)
        self.timer_maint.start()

    def __init__(self, primary=True, ports=(5556,5566)):
        self.primary = primary
        if primary:
            self.port, self.peer = ports
            frontend = "tcp://*:5003"
            backend  = "tcp://localhost:5004"
            self.kvmap = {}
        else:
            self.peer, self.port = ports
            frontend = "tcp://*:5004"
            backend  = "tcp://localhost:5003"
        
        self.ctx = zmq.Context.instance()
        self.pending = []
        self.bstar = BinaryStar(primary, frontend, backend)
        
        self.bstar.register_voter("tcp://*:%i" % self.port, zmq.ROUTER, self.handle_snapshot)
        
        # Set up our clone server sockets
        self.publisher = self.ctx.socket(zmq.PUB)
        self.collector = self.ctx.socket(zmq.SUB)
        self.collector.setsockopt(zmq.SUBSCRIBE, b'')
        self.publisher.bind("tcp://*:%d" % (self.port + 1))
        self.collector.bind("tcp://*:%d" % (self.port + 2))
        
        # Set up our own clone client interface to peer
        self.subscriber = self.ctx.socket(zmq.SUB)
        self.subscriber.setsockopt(zmq.SUBSCRIBE, b'')
        self.subscriber.connect("tcp://localhost:%d" % (self.peer + 1))
        
        # Register state change handlers
        self.bstar.master_callback = self.become_master
        self.bstar.slave_callback = self.become_slave

        # Wrap sockets in ZMQStreams for IOLoop handlers
        self.publisher = ZMQStream(self.publisher)
        self.subscriber = ZMQStream(self.subscriber)
        self.collector = ZMQStream(self.collector)
        
        # Register our handlers with reactor
        self.collector.on_recv(self.handle_collect)
        self.flush_callback = PeriodicCallback(self.flush_ttl, 1000)
        self.hugz_callback = PeriodicCallback(self.send_hugz, 1000)
        
        # basic log formatting:
        logging.basicConfig(format="%(asctime)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S",
                level=logging.INFO)

    def run(self):
        self.received_count = 0
        ctx, io_loop, self.socks, streams = get_run_context(self.sock_configs)
        self.io_loop = io_loop

        if self.request_delay > 0:
            periodic_callback = PeriodicCallback(self.do_request,
                    self.request_delay * 1000., io_loop)
            io_loop.add_timeout(timedelta(seconds=0.1), lambda: periodic_callback.start())
        else:
            for i in range(self.target_count):
                io_loop.add_callback(self.do_request)

        try:
            io_loop.start()
        except KeyboardInterrupt:
            pass

    def setup(self):
        super(DemoApp, self).setup()

        self.pub, self.pub_addr = self.stream(zmq.PUB, 'tcp://127.0.0.1:%(port)s', True)

        self.sub, sub_addr = self.stream(zmq.SUB, self.pub_addr, False,
                callback=DemoHandler())
        self.heartbeat = PeriodicCallback(self.ping, 1000, self.loop)