#!/usr/bin/env python
"""The IPython Controller with 0MQ
This is the master object that handles connections from engines and clients,
and monitors traffic through the various queues.
"""
#-----------------------------------------------------------------------------
#  Copyright (C) 2010  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------
from __future__ import print_function

import sys
import os
from datetime import datetime
import logging
import time
import uuid

import zmq
from zmq.eventloop import zmqstream, ioloop

# internal:
from IPython.zmq.log import logger # a Logger object
from IPython.zmq.entry_point import bind_port

from streamsession import Message, wrap_exception, ISO8601
from entry_point import (make_base_argument_parser, select_random_ports, split_ports,
                        connect_logger, parse_url, signal_children, generate_exec_key)
from dictdb import DictDB
try:
    from pymongo.binary import Binary
except ImportError:
    MongoDB=None
else:
    from mongodb import MongoDB
    
#-----------------------------------------------------------------------------
# Code
#-----------------------------------------------------------------------------

def _passer(*args, **kwargs):
    return

class ReverseDict(dict):
    """simple double-keyed subset of dict methods."""
    
    def __init__(self, *args, **kwargs):
        dict.__init__(self, *args, **kwargs)
        self.reverse = dict()
        for key, value in self.iteritems():
            self.reverse[value] = key
    
    def __getitem__(self, key):
        try:
            return dict.__getitem__(self, key)
        except KeyError:
            return self.reverse[key]
    
    def __setitem__(self, key, value):
        if key in self.reverse:
            raise KeyError("Can't have key %r on both sides!"%key)
        dict.__setitem__(self, key, value)
        self.reverse[value] = key
    
    def pop(self, key):
        value = dict.pop(self, key)
        self.d1.pop(value)
        return value
    

def init_record(msg):
    """return an empty TaskRecord dict, with all keys initialized with None."""
    header = msg['header']
    return {
        'msg_id' : header['msg_id'],
        'header' : header,
        'content': msg['content'],
        'buffers': msg['buffers'],
        'submitted': datetime.strptime(header['date'], ISO8601),
        'client_uuid' : None,
        'engine_uuid' : None,
        'started': None,
        'completed': None,
        'resubmitted': None,
        'result_header' : None,
        'result_content' : None,
        'result_buffers' : None,
        'queue' : None
    }


class EngineConnector(object):
    """A simple object for accessing the various zmq connections of an object.
    Attributes are:
    id (int): engine ID
    uuid (str): uuid (unused?)
    queue (str): identity of queue's XREQ socket
    registration (str): identity of registration XREQ socket
    heartbeat (str): identity of heartbeat XREQ socket
    """
    id=0
    queue=None
    control=None
    registration=None
    heartbeat=None
    pending=None
    
    def __init__(self, id, queue, registration, control, heartbeat=None):
        logger.info("engine::Engine Connected: %i"%id)
        self.id = id
        self.queue = queue
        self.registration = registration
        self.control = control
        self.heartbeat = heartbeat
        
class Controller(object):
    """The IPython Controller with 0MQ connections
    
    Parameters
    ==========
    loop: zmq IOLoop instance
    session: StreamSession object
    <removed> context: zmq context for creating new connections (?)
    queue: ZMQStream for monitoring the command queue (SUB)
    registrar: ZMQStream for engine registration requests (XREP)
    heartbeat: HeartMonitor object checking the pulse of the engines
    clientele: ZMQStream for client connections (XREP)
                not used for jobs, only query/control commands
    notifier: ZMQStream for broadcasting engine registration changes (PUB)
    db: connection to db for out of memory logging of commands
                NotImplemented
    engine_addrs: dict of zmq connection information for engines to connect
                to the queues.
    client_addrs: dict of zmq connection information for engines to connect
                to the queues.
    """
    # internal data structures:
    ids=None # engine IDs
    keytable=None
    engines=None
    clients=None
    hearts=None
    pending=None
    results=None
    tasks=None
    completed=None
    mia=None
    incoming_registrations=None
    registration_timeout=None
    
    #objects from constructor:
    loop=None
    registrar=None
    clientelle=None
    queue=None
    heartbeat=None
    notifier=None
    db=None
    client_addr=None
    engine_addrs=None
    
    
    def __init__(self, loop, session, queue, registrar, heartbeat, clientele, notifier, db, engine_addrs, client_addrs):
        """
        # universal:
        loop: IOLoop for creating future connections
        session: streamsession for sending serialized data
        # engine:
        queue: ZMQStream for monitoring queue messages
        registrar: ZMQStream for engine registration
        heartbeat: HeartMonitor object for tracking engines
        # client:
        clientele: ZMQStream for client connections
        # extra:
        db: ZMQStream for db connection (NotImplemented)
        engine_addrs: zmq address/protocol dict for engine connections
        client_addrs: zmq address/protocol dict for client connections
        """
        self.ids = set()
        self.keytable={}
        self.incoming_registrations={}
        self.engines = {}
        self.by_ident = {}
        self.clients = {}
        self.hearts = {}
        # self.mia = set()
        
        # self.sockets = {}
        self.loop = loop
        self.session = session
        self.registrar = registrar
        self.clientele = clientele
        self.queue = queue
        self.heartbeat = heartbeat
        self.notifier = notifier
        self.db = db
        
        # validate connection dicts:
        self.client_addrs = client_addrs
        assert isinstance(client_addrs['queue'], str)
        assert isinstance(client_addrs['control'], str)
        # self.hb_addrs = hb_addrs
        self.engine_addrs = engine_addrs
        assert isinstance(engine_addrs['queue'], str)
        assert isinstance(client_addrs['control'], str)
        assert len(engine_addrs['heartbeat']) == 2
        
        # register our callbacks
        self.registrar.on_recv(self.dispatch_register_request)
        self.clientele.on_recv(self.dispatch_client_msg)
        self.queue.on_recv(self.dispatch_queue_traffic)
        
        if heartbeat is not None:
            heartbeat.add_heart_failure_handler(self.handle_heart_failure)
            heartbeat.add_new_heart_handler(self.handle_new_heart)
        
        self.queue_handlers = { 'in' : self.save_queue_request,
                                'out': self.save_queue_result,
                                'intask': self.save_task_request,
                                'outtask': self.save_task_result,
                                'tracktask': self.save_task_destination,
                                'incontrol': _passer,
                                'outcontrol': _passer,
        }
        
        self.client_handlers = {'queue_request': self.queue_status,
                                'result_request': self.get_results,
                                'purge_request': self.purge_results,
                                'load_request': self.check_load,
                                'resubmit_request': self.resubmit_task,
                                'shutdown_request': self.shutdown_request,
                                }
        
        self.registrar_handlers = {'registration_request' : self.register_engine,
                                'unregistration_request' : self.unregister_engine,
                                'connection_request': self.connection_request,
        }
        self.registration_timeout = max(5000, 2*self.heartbeat.period)
        # this is the stuff that will move to DB:
        # self.results = {} # completed results
        self.pending = set() # pending messages, keyed by msg_id
        self.queues = {} # pending msg_ids keyed by engine_id
        self.tasks = {} # pending msg_ids submitted as tasks, keyed by client_id
        self.completed = {} # completed msg_ids keyed by engine_id
        self.all_completed = set()
        
        logger.info("controller::created controller")
    
    def _new_id(self):
        """gemerate a new ID"""
        newid = 0
        incoming = [id[0] for id in self.incoming_registrations.itervalues()]
        # print newid, self.ids, self.incoming_registrations
        while newid in self.ids or newid in incoming:
            newid += 1
        return newid
    
    #-----------------------------------------------------------------------------
    # message validation
    #-----------------------------------------------------------------------------
    
    def _validate_targets(self, targets):
        """turn any valid targets argument into a list of integer ids"""
        if targets is None:
            # default to all
            targets = self.ids
            
        if isinstance(targets, (int,str,unicode)):
            # only one target specified
            targets = [targets]
        _targets = []
        for t in targets:
            # map raw identities to ids
            if isinstance(t, (str,unicode)):
                t = self.by_ident.get(t, t)
            _targets.append(t)
        targets = _targets
        bad_targets = [ t for t in targets if t not in self.ids ]
        if bad_targets:
            raise IndexError("No Such Engine: %r"%bad_targets)
        if not targets:
            raise IndexError("No Engines Registered")
        return targets
    
    def _validate_client_msg(self, msg):
        """validates and unpacks headers of a message. Returns False if invalid,
        (ident, header, parent, content)"""
        client_id = msg[0]
        try:
            msg = self.session.unpack_message(msg[1:], content=True)
        except:
            logger.error("client::Invalid Message %s"%msg)
            return False
        
        msg_type = msg.get('msg_type', None)
        if msg_type is None:
            return False
        header = msg.get('header')
        # session doesn't handle split content for now:
        return client_id, msg
        
    
    #-----------------------------------------------------------------------------
    # dispatch methods (1 per stream)
    #-----------------------------------------------------------------------------
    
    def dispatch_register_request(self, msg):
        """"""
        logger.debug("registration::dispatch_register_request(%s)"%msg)
        idents,msg = self.session.feed_identities(msg)
        if not idents:
            logger.error("Bad Queue Message: %s"%msg, exc_info=True)
            return
        try:
            msg = self.session.unpack_message(msg,content=True)
        except:
            logger.error("registration::got bad registration message: %s"%msg, exc_info=True)
            return
        
        msg_type = msg['msg_type']
        content = msg['content']
        
        handler = self.registrar_handlers.get(msg_type, None)
        if handler is None:
            logger.error("registration::got bad registration message: %s"%msg)
        else:
            handler(idents, msg)
    
    def dispatch_queue_traffic(self, msg):
        """all ME and Task queue messages come through here"""
        logger.debug("queue traffic: %s"%msg[:2])
        switch = msg[0]
        idents, msg = self.session.feed_identities(msg[1:])
        if not idents:
            logger.error("Bad Queue Message: %s"%msg)
            return
        handler = self.queue_handlers.get(switch, None)
        if handler is not None:
            handler(idents, msg)
        else:
            logger.error("Invalid message topic: %s"%switch)
        
    
    def dispatch_client_msg(self, msg):
        """Route messages from clients"""
        idents, msg = self.session.feed_identities(msg)
        if not idents:
            logger.error("Bad Client Message: %s"%msg)
            return
        client_id = idents[0]
        try:
            msg = self.session.unpack_message(msg, content=True)
        except:
            content = wrap_exception()
            logger.error("Bad Client Message: %s"%msg, exc_info=True)
            self.session.send(self.clientele, "controller_error", ident=client_id, 
                    content=content)
            return
        
        # print client_id, header, parent, content
        #switch on message type:
        msg_type = msg['msg_type']
        logger.info("client:: client %s requested %s"%(client_id, msg_type))
        handler = self.client_handlers.get(msg_type, None)
        try:
            assert handler is not None, "Bad Message Type: %s"%msg_type
        except:
            content = wrap_exception()
            logger.error("Bad Message Type: %s"%msg_type, exc_info=True)
            self.session.send(self.clientele, "controller_error", ident=client_id, 
                    content=content)
            return
        else:
            handler(client_id, msg)
            
    def dispatch_db(self, msg):
        """"""
        raise NotImplementedError
    
    #---------------------------------------------------------------------------
    # handler methods (1 per event)
    #---------------------------------------------------------------------------
    
    #----------------------- Heartbeat --------------------------------------
    
    def handle_new_heart(self, heart):
        """handler to attach to heartbeater.
        Called when a new heart starts to beat.
        Triggers completion of registration."""
        logger.debug("heartbeat::handle_new_heart(%r)"%heart)
        if heart not in self.incoming_registrations:
            logger.info("heartbeat::ignoring new heart: %r"%heart)
        else:
            self.finish_registration(heart)
        
    
    def handle_heart_failure(self, heart):
        """handler to attach to heartbeater.
        called when a previously registered heart fails to respond to beat request.
        triggers unregistration"""
        logger.debug("heartbeat::handle_heart_failure(%r)"%heart)
        eid = self.hearts.get(heart, None)
        queue = self.engines[eid].queue
        if eid is None:
            logger.info("heartbeat::ignoring heart failure %r"%heart)
        else:
            self.unregister_engine(heart, dict(content=dict(id=eid, queue=queue)))
    
    #----------------------- MUX Queue Traffic ------------------------------
    
    def save_queue_request(self, idents, msg):
        if len(idents) < 2:
            logger.error("invalid identity prefix: %s"%idents)
            return
        queue_id, client_id = idents[:2]
        try:
            msg = self.session.unpack_message(msg, content=False)
        except:
            logger.error("queue::client %r sent invalid message to %r: %s"%(client_id, queue_id, msg), exc_info=True)
            return
        
        eid = self.by_ident.get(queue_id, None)
        if eid is None:
            logger.error("queue::target %r not registered"%queue_id)
            logger.debug("queue::    valid are: %s"%(self.by_ident.keys()))
            return
            
        header = msg['header']
        msg_id = header['msg_id']
        record = init_record(msg)
        record['engine_uuid'] = queue_id
        record['client_uuid'] = client_id
        record['queue'] = 'mux'
        if MongoDB is not None and isinstance(self.db, MongoDB):
            record['buffers'] = map(Binary, record['buffers'])
        self.pending.add(msg_id)
        self.queues[eid].append(msg_id)
        self.db.add_record(msg_id, record)
    
    def save_queue_result(self, idents, msg):
        if len(idents) < 2:
            logger.error("invalid identity prefix: %s"%idents)
            return
            
        client_id, queue_id = idents[:2]
        try:
            msg = self.session.unpack_message(msg, content=False)
        except:
            logger.error("queue::engine %r sent invalid message to %r: %s"%(
                    queue_id,client_id, msg), exc_info=True)
            return
        
        eid = self.by_ident.get(queue_id, None)
        if eid is None:
            logger.error("queue::unknown engine %r is sending a reply: "%queue_id)
            logger.debug("queue::       %s"%msg[2:])
            return
        
        parent = msg['parent_header']
        if not parent:
            return
        msg_id = parent['msg_id']
        if msg_id in self.pending:
            self.pending.remove(msg_id)
            self.all_completed.add(msg_id)
            self.queues[eid].remove(msg_id)
            self.completed[eid].append(msg_id)
            rheader = msg['header']
            completed = datetime.strptime(rheader['date'], ISO8601)
            started = rheader.get('started', None)
            if started is not None:
                started = datetime.strptime(started, ISO8601)
            result = {
                'result_header' : rheader,
                'result_content': msg['content'],
                'started' : started,
                'completed' : completed
            }
            if MongoDB is not None and isinstance(self.db, MongoDB):
                result['result_buffers'] = map(Binary, msg['buffers'])
            self.db.update_record(msg_id, result)
        else:
            logger.debug("queue:: unknown msg finished %s"%msg_id)
            
    #--------------------- Task Queue Traffic ------------------------------
    
    def save_task_request(self, idents, msg):
        """Save the submission of a task."""
        client_id = idents[0]
        
        try:
            msg = self.session.unpack_message(msg, content=False)
        except:
            logger.error("task::client %r sent invalid task message: %s"%(
                    client_id, msg), exc_info=True)
            return
        record = init_record(msg)
        if MongoDB is not None and isinstance(self.db, MongoDB):
            record['buffers'] = map(Binary, record['buffers'])
        record['client_uuid'] = client_id
        record['queue'] = 'task'
        header = msg['header']
        msg_id = header['msg_id']
        self.pending.add(msg_id)
        self.db.add_record(msg_id, record)
    
    def save_task_result(self, idents, msg):
        """save the result of a completed task."""
        client_id = idents[0]
        try:
            msg = self.session.unpack_message(msg, content=False)
        except:
            logger.error("task::invalid task result message send to %r: %s"%(
                    client_id, msg))
            raise
            return
        
        parent = msg['parent_header']
        if not parent:
            # print msg
            logger.warn("Task %r had no parent!"%msg)
            return
        msg_id = parent['msg_id']
        
        header = msg['header']
        engine_uuid = header.get('engine', None)
        eid = self.by_ident.get(engine_uuid, None)
        
        if msg_id in self.pending:
            self.pending.remove(msg_id)
            self.all_completed.add(msg_id)
            if eid is not None:
                self.completed[eid].append(msg_id)
                if msg_id in self.tasks[eid]:
                    self.tasks[eid].remove(msg_id)
            completed = datetime.strptime(header['date'], ISO8601)
            started = header.get('started', None)
            if started is not None:
                started = datetime.strptime(started, ISO8601)
            result = {
                'result_header' : header,
                'result_content': msg['content'],
                'started' : started,
                'completed' : completed,
                'engine_uuid': engine_uuid
            }
            if MongoDB is not None and isinstance(self.db, MongoDB):
                result['result_buffers'] = map(Binary, msg['buffers'])
            self.db.update_record(msg_id, result)
            
        else:
            logger.debug("task::unknown task %s finished"%msg_id)
    
    def save_task_destination(self, idents, msg):
        try:
            msg = self.session.unpack_message(msg, content=True)
        except:
            logger.error("task::invalid task tracking message")
            return
        content = msg['content']
        print (content)
        msg_id = content['msg_id']
        engine_uuid = content['engine_id']
        eid = self.by_ident[engine_uuid]
        
        logger.info("task::task %s arrived on %s"%(msg_id, eid))
        # if msg_id in self.mia:
        #     self.mia.remove(msg_id)
        # else:
        #     logger.debug("task::task %s not listed as MIA?!"%(msg_id))
        
        self.tasks[eid].append(msg_id)
        # self.pending[msg_id][1].update(received=datetime.now(),engine=(eid,engine_uuid))
        self.db.update_record(msg_id, dict(engine_uuid=engine_uuid))
    
    def mia_task_request(self, idents, msg):
        raise NotImplementedError
        client_id = idents[0]
        # content = dict(mia=self.mia,status='ok')
        # self.session.send('mia_reply', content=content, idents=client_id)
        
        
            
    #-------------------------------------------------------------------------
    # Registration requests
    #-------------------------------------------------------------------------
        
    def connection_request(self, client_id, msg):
        """Reply with connection addresses for clients."""
        logger.info("client::client %s connected"%client_id)
        content = dict(status='ok')
        content.update(self.client_addrs)
        jsonable = {}
        for k,v in self.keytable.iteritems():
            jsonable[str(k)] = v
        content['engines'] = jsonable
        self.session.send(self.registrar, 'connection_reply', content, parent=msg, ident=client_id)
    
    def register_engine(self, reg, msg):
        """Register a new engine."""
        content = msg['content']
        try:
            queue = content['queue']
        except KeyError:
            logger.error("registration::queue not specified")
            return
        heart = content.get('heartbeat', None)
        """register a new engine, and create the socket(s) necessary"""
        eid = self._new_id()
        # print (eid, queue, reg, heart)
        
        logger.debug("registration::register_engine(%i, %r, %r, %r)"%(eid, queue, reg, heart))
        
        content = dict(id=eid,status='ok')
        content.update(self.engine_addrs)
        # check if requesting available IDs:
        if queue in self.by_ident:
            try:
                raise KeyError("queue_id %r in use"%queue)
            except:
                content = wrap_exception()
        elif heart in self.hearts: # need to check unique hearts?
            try:
                raise KeyError("heart_id %r in use"%heart)
            except:
                content = wrap_exception()
        else:
            for h, pack in self.incoming_registrations.iteritems():
                if heart == h:
                    try:
                        raise KeyError("heart_id %r in use"%heart)
                    except:
                        content = wrap_exception()
                    break
                elif queue == pack[1]:
                    try:
                        raise KeyError("queue_id %r in use"%queue)
                    except:
                        content = wrap_exception()
                    break
        
        msg = self.session.send(self.registrar, "registration_reply", 
                content=content, 
                ident=reg)
        
        if content['status'] == 'ok':
            if heart in self.heartbeat.hearts:
                # already beating
                self.incoming_registrations[heart] = (eid,queue,reg,None)
                self.finish_registration(heart)
            else:
                purge = lambda : self._purge_stalled_registration(heart)
                dc = ioloop.DelayedCallback(purge, self.registration_timeout, self.loop)
                dc.start()
                self.incoming_registrations[heart] = (eid,queue,reg,dc)
        else:
            logger.error("registration::registration %i failed: %s"%(eid, content['evalue']))
        return eid
    
    def unregister_engine(self, ident, msg):
        """Unregister an engine that explicitly requested to leave."""
        try:
            eid = msg['content']['id']
        except:
            logger.error("registration::bad engine id for unregistration: %s"%ident)
            return
        logger.info("registration::unregister_engine(%s)"%eid)
        content=dict(id=eid, queue=self.engines[eid].queue)
        self.ids.remove(eid)
        self.keytable.pop(eid)
        ec = self.engines.pop(eid)
        self.hearts.pop(ec.heartbeat)
        self.by_ident.pop(ec.queue)
        self.completed.pop(eid)
        for msg_id in self.queues.pop(eid):
            msg = self.pending.remove(msg_id)
            ############## TODO: HANDLE IT ################
        
        if self.notifier:
            self.session.send(self.notifier, "unregistration_notification", content=content)
    
    def finish_registration(self, heart):
        """Second half of engine registration, called after our HeartMonitor
        has received a beat from the Engine's Heart."""
        try: 
            (eid,queue,reg,purge) = self.incoming_registrations.pop(heart)
        except KeyError:
            logger.error("registration::tried to finish nonexistant registration")
            return
        logger.info("registration::finished registering engine %i:%r"%(eid,queue))
        if purge is not None:
            purge.stop()
        control = queue
        self.ids.add(eid)
        self.keytable[eid] = queue
        self.engines[eid] = EngineConnector(eid, queue, reg, control, heart)
        self.by_ident[queue] = eid
        self.queues[eid] = list()
        self.tasks[eid] = list()
        self.completed[eid] = list()
        self.hearts[heart] = eid
        content = dict(id=eid, queue=self.engines[eid].queue)
        if self.notifier:
            self.session.send(self.notifier, "registration_notification", content=content)
    
    def _purge_stalled_registration(self, heart):
        if heart in self.incoming_registrations:
            eid = self.incoming_registrations.pop(heart)[0]
            logger.info("registration::purging stalled registration: %i"%eid)
        else:
            pass
            
    #-------------------------------------------------------------------------
    # Client Requests
    #-------------------------------------------------------------------------
    
    def shutdown_request(self, client_id, msg):
        """handle shutdown request."""
        # s = self.context.socket(zmq.XREQ)
        # s.connect(self.client_connections['mux'])
        # time.sleep(0.1)
        # for eid,ec in self.engines.iteritems():
        #     self.session.send(s, 'shutdown_request', content=dict(restart=False), ident=ec.queue)
        # time.sleep(1)
        self.session.send(self.clientele, 'shutdown_reply', content={'status': 'ok'}, ident=client_id)
        dc = ioloop.DelayedCallback(lambda : self._shutdown(), 1000, self.loop)
        dc.start()
    
    def _shutdown(self):
        logger.info("controller::controller shutting down.")
        time.sleep(0.1)
        sys.exit(0)
        
    
    def check_load(self, client_id, msg):
        content = msg['content']
        try:
            targets = content['targets']
            targets = self._validate_targets(targets)
        except:
            content = wrap_exception()
            self.session.send(self.clientele, "controller_error", 
                    content=content, ident=client_id)
            return
        
        content = dict(status='ok')
        # loads = {}
        for t in targets:
            content[bytes(t)] = len(self.queues[t])+len(self.tasks[t])
        self.session.send(self.clientele, "load_reply", content=content, ident=client_id)
            
    
    def queue_status(self, client_id, msg):
        """Return the Queue status of one or more targets.
        if verbose: return the msg_ids
        else: return len of each type.
        keys: queue (pending MUX jobs)
            tasks (pending Task jobs)
            completed (finished jobs from both queues)"""
        content = msg['content']
        targets = content['targets']
        try:
            targets = self._validate_targets(targets)
        except:
            content = wrap_exception()
            self.session.send(self.clientele, "controller_error", 
                    content=content, ident=client_id)
            return
        verbose = content.get('verbose', False)
        content = dict(status='ok')
        for t in targets:
            queue = self.queues[t]
            completed = self.completed[t]
            tasks = self.tasks[t]
            if not verbose:
                queue = len(queue)
                completed = len(completed)
                tasks = len(tasks)
            content[bytes(t)] = {'queue': queue, 'completed': completed , 'tasks': tasks}
            # pending
        self.session.send(self.clientele, "queue_reply", content=content, ident=client_id)
    
    def purge_results(self, client_id, msg):
        """Purge results from memory. This method is more valuable before we move
        to a DB based message storage mechanism."""
        content = msg['content']
        msg_ids = content.get('msg_ids', [])
        reply = dict(status='ok')
        if msg_ids == 'all':
            self.db.drop_matching_records(dict(completed={'$ne':None}))
        else:
            for msg_id in msg_ids:
                if msg_id in self.all_completed:
                    self.db.drop_record(msg_id)
                else:
                    if msg_id in self.pending:
                        try:
                            raise IndexError("msg pending: %r"%msg_id)
                        except:
                            reply = wrap_exception()
                    else:
                        try:
                            raise IndexError("No such msg: %r"%msg_id)
                        except:
                            reply = wrap_exception()
                    break
            eids = content.get('engine_ids', [])
            for eid in eids:
                if eid not in self.engines:
                    try:
                        raise IndexError("No such engine: %i"%eid)
                    except:
                        reply = wrap_exception()
                    break
                msg_ids = self.completed.pop(eid)
                uid = self.engines[eid].queue
                self.db.drop_matching_records(dict(engine_uuid=uid, completed={'$ne':None}))
        
        self.session.send(self.clientele, 'purge_reply', content=reply, ident=client_id)
    
    def resubmit_task(self, client_id, msg, buffers):
        """Resubmit a task."""
        raise NotImplementedError
    
    def get_results(self, client_id, msg):
        """Get the result of 1 or more messages."""
        content = msg['content']
        msg_ids = set(content['msg_ids'])
        statusonly = content.get('status_only', False)
        pending = []
        completed = []
        content = dict(status='ok')
        content['pending'] = pending
        content['completed'] = completed
        if not statusonly:
            records = self.db.find_records(dict(msg_id={'$in':msg_ids}))
        for msg_id in msg_ids:
            if msg_id in self.pending:
                pending.append(msg_id)
            elif msg_id in self.all_completed:
                completed.append(msg_id)
                if not statusonly:
                    content[msg_id] = records[msg_id]['result_content']
            else:
                try:
                    raise KeyError('No such message: '+msg_id)
                except:
                    content = wrap_exception()
                break
        self.session.send(self.clientele, "result_reply", content=content, 
                                            parent=msg, ident=client_id)


#-------------------------------------------------------------------------
# Entry Point
#-------------------------------------------------------------------------

def make_argument_parser():
    """Make an argument parser"""
    parser = make_base_argument_parser()
    
    parser.add_argument('--client', type=int, metavar='PORT', default=0,
                        help='set the XREP port for clients [default: random]')
    parser.add_argument('--notice', type=int, metavar='PORT', default=0,
                        help='set the PUB socket for registration notification [default: random]')
    parser.add_argument('--hb', type=str, metavar='PORTS',
                        help='set the 2 ports for heartbeats [default: random]')
    parser.add_argument('--ping', type=int, default=3000,
                        help='set the heartbeat period in ms [default: 3000]')
    parser.add_argument('--monitor', type=int, metavar='PORT', default=0,
                        help='set the SUB port for queue monitoring [default: random]')
    parser.add_argument('--mux', type=str, metavar='PORTS',
                        help='set the XREP ports for the MUX queue [default: random]')
    parser.add_argument('--task', type=str, metavar='PORTS',
                        help='set the XREP/XREQ ports for the task queue [default: random]')
    parser.add_argument('--control', type=str, metavar='PORTS',
                        help='set the XREP ports for the control queue [default: random]')
    parser.add_argument('--scheduler', type=str, default='pure',
                        choices = ['pure', 'lru', 'plainrandom', 'weighted', 'twobin','leastload'],
                        help='select the task scheduler  [default: pure ZMQ]')
    parser.add_argument('--mongodb', action='store_true',
                        help='Use MongoDB task storage [default: in-memory]')
    
    return parser
    
def main(argv=None):
    import time
    from multiprocessing import Process
    
    from zmq.eventloop.zmqstream import ZMQStream
    from zmq.devices import ProcessMonitoredQueue
    from zmq.log import handlers
    
    import streamsession as session
    import heartmonitor
    from scheduler import launch_scheduler
    
    parser = make_argument_parser()
    
    args = parser.parse_args(argv)
    parse_url(args)
    
    iface="%s://%s"%(args.transport,args.ip)+':%i'
    
    random_ports = 0
    if args.hb:
        hb = split_ports(args.hb, 2)
    else:
        hb = select_random_ports(2)
    if args.mux:
        mux = split_ports(args.mux, 2)
    else:
        mux = None
        random_ports += 2
    if args.task:
        task = split_ports(args.task, 2)
    else:
        task = None
        random_ports += 2
    if args.control:
        control = split_ports(args.control, 2)
    else:
        control = None
        random_ports += 2
    
    ctx = zmq.Context()
    loop = ioloop.IOLoop.instance()
    
    # setup logging
    connect_logger(ctx, iface%args.logport, root="controller", loglevel=args.loglevel)
    
    # Registrar socket
    reg = ZMQStream(ctx.socket(zmq.XREP), loop)
    regport = bind_port(reg, args.ip, args.regport)
    
    ### Engine connections ###
    
    # heartbeat
    hpub = ctx.socket(zmq.PUB)
    bind_port(hpub, args.ip, hb[0])
    hrep = ctx.socket(zmq.XREP)
    bind_port(hrep, args.ip, hb[1])
    
    hmon = heartmonitor.HeartMonitor(loop, ZMQStream(hpub,loop), ZMQStream(hrep,loop),args.ping)
    hmon.start()
    
    ### Client connections ###
    # Clientele socket
    c = ZMQStream(ctx.socket(zmq.XREP), loop)
    cport = bind_port(c, args.ip, args.client)
    # Notifier socket
    n = ZMQStream(ctx.socket(zmq.PUB), loop)
    nport = bind_port(n, args.ip, args.notice)
    
    ### Key File ###
    if args.execkey and not os.path.isfile(args.execkey):
            generate_exec_key(args.execkey)
    
    thesession = session.StreamSession(username=args.ident or "controller", keyfile=args.execkey)
    
    ### build and launch the queues ###
    
    # monitor socket
    sub = ctx.socket(zmq.SUB)
    sub.setsockopt(zmq.SUBSCRIBE, "")
    monport = bind_port(sub, args.ip, args.monitor)
    sub = ZMQStream(sub, loop)
    
    ports = select_random_ports(random_ports)
    children = []
    # Multiplexer Queue (in a Process)
    if not mux:
        mux = (ports.pop(),ports.pop())
    q = ProcessMonitoredQueue(zmq.XREP, zmq.XREP, zmq.PUB, 'in', 'out')
    q.bind_in(iface%mux[0])
    q.bind_out(iface%mux[1])
    q.connect_mon(iface%monport)
    q.daemon=True
    q.start()
    children.append(q.launcher)
    
    # Control Queue (in a Process)
    if not control:
        control = (ports.pop(),ports.pop())
    q = ProcessMonitoredQueue(zmq.XREP, zmq.XREP, zmq.PUB, 'incontrol', 'outcontrol')
    q.bind_in(iface%control[0])
    q.bind_out(iface%control[1])
    q.connect_mon(iface%monport)
    q.daemon=True
    q.start()
    children.append(q.launcher)
    # Task Queue (in a Process)
    if not task:
        task = (ports.pop(),ports.pop())
    if args.scheduler == 'pure':
        q = ProcessMonitoredQueue(zmq.XREP, zmq.XREQ, zmq.PUB, 'intask', 'outtask')
        q.bind_in(iface%task[0])
        q.bind_out(iface%task[1])
        q.connect_mon(iface%monport)
        q.daemon=True
        q.start()
        children.append(q.launcher)
    else:
        sargs = (iface%task[0],iface%task[1],iface%monport,iface%nport,args.scheduler)
        print (sargs)
        q = Process(target=launch_scheduler, args=sargs)
        q.daemon=True
        q.start()
        children.append(q)
    
    if args.mongodb:
        from mongodb import MongoDB
        db = MongoDB(thesession.session)
    else:
        db = DictDB()
    time.sleep(.25)
    
    # build connection dicts
    engine_addrs = {
        'control' : iface%control[1],
        'queue': iface%mux[1],
        'heartbeat': (iface%hb[0], iface%hb[1]),
        'task' : iface%task[1],
        'monitor' : iface%monport,
        }
    
    client_addrs = {
        'control' : iface%control[0],
        'query': iface%cport,
        'queue': iface%mux[0],
        'task' : iface%task[0],
        'notification': iface%nport
        }
    signal_children(children)
    con = Controller(loop, thesession, sub, reg, hmon, c, n, db, engine_addrs, client_addrs)
    dc = ioloop.DelayedCallback(lambda : print("Controller started..."), 100, loop)
    dc.start()
    loop.start()
    
    
    

if __name__ == '__main__':
    main()