"""The IPython Controller Hub with 0MQ
This is the master object that handles connections from engines and clients,
and monitors traffic through the various queues.

Authors:

* Min RK
"""
#-----------------------------------------------------------------------------
#  Copyright (C) 2010-2011  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 time
from datetime import datetime

import zmq
from zmq.eventloop import ioloop
from zmq.eventloop.zmqstream import ZMQStream

# internal:
from IPython.utils.importstring import import_item
from IPython.utils.py3compat import cast_bytes
from IPython.utils.traitlets import (
        HasTraits, Instance, Integer, Unicode, Dict, Set, Tuple, CBytes, DottedObjectName
        )

from IPython.parallel import error, util
from IPython.parallel.factory import RegistrationFactory

from IPython.zmq.session import SessionFactory

from .heartmonitor import HeartMonitor

#-----------------------------------------------------------------------------
# Code
#-----------------------------------------------------------------------------

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

def _printer(*args, **kwargs):
    print (args)
    print (kwargs)

def empty_record():
    """Return an empty dict with all record keys."""
    return {
        'msg_id' : None,
        'header' : None,
        'content': None,
        'buffers': None,
        'submitted': None,
        'client_uuid' : None,
        'engine_uuid' : None,
        'started': None,
        'completed': None,
        'resubmitted': None,
        'received': None,
        'result_header' : None,
        'result_content' : None,
        'result_buffers' : None,
        'queue' : None,
        'pyin' : None,
        'pyout': None,
        'pyerr': None,
        'stdout': '',
        'stderr': '',
    }

def init_record(msg):
    """Initialize a TaskRecord based on a request."""
    header = msg['header']
    return {
        'msg_id' : header['msg_id'],
        'header' : header,
        'content': msg['content'],
        'buffers': msg['buffers'],
        'submitted': header['date'],
        'client_uuid' : None,
        'engine_uuid' : None,
        'started': None,
        'completed': None,
        'resubmitted': None,
        'received': None,
        'result_header' : None,
        'result_content' : None,
        'result_buffers' : None,
        'queue' : None,
        'pyin' : None,
        'pyout': None,
        'pyerr': None,
        'stdout': '',
        'stderr': '',
    }

class EngineConnector(HasTraits):
    """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 DEALER socket
    registration (str): identity of registration DEALER socket
    heartbeat (str): identity of heartbeat DEALER socket
    """
    id=Integer(0)
    queue=CBytes()
    control=CBytes()
    registration=CBytes()
    heartbeat=CBytes()
    pending=Set()

_db_shortcuts = {
    'sqlitedb' : 'IPython.parallel.controller.sqlitedb.SQLiteDB',
    'mongodb'  : 'IPython.parallel.controller.mongodb.MongoDB',
    'dictdb'   : 'IPython.parallel.controller.dictdb.DictDB',
    'nodb'     : 'IPython.parallel.controller.dictdb.NoDB',
}

class HubFactory(RegistrationFactory):
    """The Configurable for setting up a Hub."""

    # port-pairs for monitoredqueues:
    hb = Tuple(Integer,Integer,config=True,
        help="""DEALER/SUB Port pair for Engine heartbeats""")
    def _hb_default(self):
        return tuple(util.select_random_ports(2))

    mux = Tuple(Integer,Integer,config=True,
        help="""Engine/Client Port pair for MUX queue""")

    def _mux_default(self):
        return tuple(util.select_random_ports(2))

    task = Tuple(Integer,Integer,config=True,
        help="""Engine/Client Port pair for Task queue""")
    def _task_default(self):
        return tuple(util.select_random_ports(2))

    control = Tuple(Integer,Integer,config=True,
        help="""Engine/Client Port pair for Control queue""")

    def _control_default(self):
        return tuple(util.select_random_ports(2))

    iopub = Tuple(Integer,Integer,config=True,
        help="""Engine/Client Port pair for IOPub relay""")

    def _iopub_default(self):
        return tuple(util.select_random_ports(2))

    # single ports:
    mon_port = Integer(config=True,
        help="""Monitor (SUB) port for queue traffic""")

    def _mon_port_default(self):
        return util.select_random_ports(1)[0]

    notifier_port = Integer(config=True,
        help="""PUB port for sending engine status notifications""")

    def _notifier_port_default(self):
        return util.select_random_ports(1)[0]

    engine_ip = Unicode('127.0.0.1', config=True,
        help="IP on which to listen for engine connections. [default: loopback]")
    engine_transport = Unicode('tcp', config=True,
        help="0MQ transport for engine connections. [default: tcp]")

    client_ip = Unicode('127.0.0.1', config=True,
        help="IP on which to listen for client connections. [default: loopback]")
    client_transport = Unicode('tcp', config=True,
        help="0MQ transport for client connections. [default : tcp]")

    monitor_ip = Unicode('127.0.0.1', config=True,
        help="IP on which to listen for monitor messages. [default: loopback]")
    monitor_transport = Unicode('tcp', config=True,
        help="0MQ transport for monitor messages. [default : tcp]")

    monitor_url = Unicode('')

    db_class = DottedObjectName('NoDB',
        config=True, help="""The class to use for the DB backend

        Options include:

        SQLiteDB: SQLite
        MongoDB : use MongoDB
        DictDB  : in-memory storage (fastest, but be mindful of memory growth of the Hub)
        NoDB    : disable database altogether (default)

        """)

    # not configurable
    db = Instance('IPython.parallel.controller.dictdb.BaseDB')
    heartmonitor = Instance('IPython.parallel.controller.heartmonitor.HeartMonitor')

    def _ip_changed(self, name, old, new):
        self.engine_ip = new
        self.client_ip = new
        self.monitor_ip = new
        self._update_monitor_url()

    def _update_monitor_url(self):
        self.monitor_url = "%s://%s:%i" % (self.monitor_transport, self.monitor_ip, self.mon_port)

    def _transport_changed(self, name, old, new):
        self.engine_transport = new
        self.client_transport = new
        self.monitor_transport = new
        self._update_monitor_url()

    def __init__(self, **kwargs):
        super(HubFactory, self).__init__(**kwargs)
        self._update_monitor_url()

    def construct(self):
        self.init_hub()

    def start(self):
        self.heartmonitor.start()
        self.log.info("Heartmonitor started")

    def init_hub(self):
        """construct"""
        client_iface = "%s://%s:" % (self.client_transport, self.client_ip) + "%i"
        engine_iface = "%s://%s:" % (self.engine_transport, self.engine_ip) + "%i"

        ctx = self.context
        loop = self.loop

        # Registrar socket
        q = ZMQStream(ctx.socket(zmq.ROUTER), loop)
        q.bind(client_iface % self.regport)
        self.log.info("Hub listening on %s for registration.", client_iface % self.regport)
        if self.client_ip != self.engine_ip:
            q.bind(engine_iface % self.regport)
            self.log.info("Hub listening on %s for registration.", engine_iface % self.regport)

        ### Engine connections ###

        # heartbeat
        hpub = ctx.socket(zmq.PUB)
        hpub.bind(engine_iface % self.hb[0])
        hrep = ctx.socket(zmq.ROUTER)
        hrep.bind(engine_iface % self.hb[1])
        self.heartmonitor = HeartMonitor(loop=loop, config=self.config, log=self.log,
                                pingstream=ZMQStream(hpub,loop),
                                pongstream=ZMQStream(hrep,loop)
                            )

        ### Client connections ###
        # Notifier socket
        n = ZMQStream(ctx.socket(zmq.PUB), loop)
        n.bind(client_iface%self.notifier_port)

        ### build and launch the queues ###

        # monitor socket
        sub = ctx.socket(zmq.SUB)
        sub.setsockopt(zmq.SUBSCRIBE, b"")
        sub.bind(self.monitor_url)
        sub.bind('inproc://monitor')
        sub = ZMQStream(sub, loop)

        # connect the db
        db_class = _db_shortcuts.get(self.db_class.lower(), self.db_class)
        self.log.info('Hub using DB backend: %r', (db_class.split('.')[-1]))
        self.db = import_item(str(db_class))(session=self.session.session,
                                            config=self.config, log=self.log)
        time.sleep(.25)
        try:
            scheme = self.config.TaskScheduler.scheme_name
        except AttributeError:
            from .scheduler import TaskScheduler
            scheme = TaskScheduler.scheme_name.get_default_value()
        # build connection dicts
        self.engine_info = {
            'control' : engine_iface%self.control[1],
            'mux': engine_iface%self.mux[1],
            'heartbeat': (engine_iface%self.hb[0], engine_iface%self.hb[1]),
            'task' : engine_iface%self.task[1],
            'iopub' : engine_iface%self.iopub[1],
            # 'monitor' : engine_iface%self.mon_port,
            }

        self.client_info = {
            'control' : client_iface%self.control[0],
            'mux': client_iface%self.mux[0],
            'task' : (scheme, client_iface%self.task[0]),
            'iopub' : client_iface%self.iopub[0],
            'notification': client_iface%self.notifier_port
            }
        self.log.debug("Hub engine addrs: %s", self.engine_info)
        self.log.debug("Hub client addrs: %s", self.client_info)

        # resubmit stream
        r = ZMQStream(ctx.socket(zmq.DEALER), loop)
        url = util.disambiguate_url(self.client_info['task'][-1])
        r.setsockopt(zmq.IDENTITY, self.session.bsession)
        r.connect(url)

        self.hub = Hub(loop=loop, session=self.session, monitor=sub, heartmonitor=self.heartmonitor,
                query=q, notifier=n, resubmit=r, db=self.db,
                engine_info=self.engine_info, client_info=self.client_info,
                log=self.log)

class Hub(SessionFactory):
    """The IPython Controller Hub with 0MQ connections

    Parameters
    ==========
    loop: zmq IOLoop instance
    session: Session object
    <removed> context: zmq context for creating new connections (?)
    queue: ZMQStream for monitoring the command queue (SUB)
    query: ZMQStream for engine registration and client queries requests (ROUTER)
    heartbeat: HeartMonitor object checking the pulse of the engines
    notifier: ZMQStream for broadcasting engine registration changes (PUB)
    db: connection to db for out of memory logging of commands
                NotImplemented
    engine_info: dict of zmq connection information for engines to connect
                to the queues.
    client_info: dict of zmq connection information for engines to connect
                to the queues.
    """
    # internal data structures:
    ids=Set() # engine IDs
    keytable=Dict()
    by_ident=Dict()
    engines=Dict()
    clients=Dict()
    hearts=Dict()
    pending=Set()
    queues=Dict()  # pending msg_ids keyed by engine_id
    tasks=Dict() # pending msg_ids submitted as tasks, keyed by client_id
    completed=Dict() # completed msg_ids keyed by engine_id
    all_completed=Set() # completed msg_ids keyed by engine_id
    dead_engines=Set() # completed msg_ids keyed by engine_id
    unassigned=Set() # set of task msg_ds not yet assigned a destination
    incoming_registrations=Dict()
    registration_timeout=Integer()
    _idcounter=Integer(0)

    # objects from constructor:
    query=Instance(ZMQStream)
    monitor=Instance(ZMQStream)
    notifier=Instance(ZMQStream)
    resubmit=Instance(ZMQStream)
    heartmonitor=Instance(HeartMonitor)
    db=Instance(object)
    client_info=Dict()
    engine_info=Dict()

    def __init__(self, **kwargs):
        """
        # universal:
        loop: IOLoop for creating future connections
        session: streamsession for sending serialized data
        # engine:
        queue: ZMQStream for monitoring queue messages
        query: ZMQStream for engine+client registration and client requests
        heartbeat: HeartMonitor object for tracking engines
        # extra:
        db: ZMQStream for db connection (NotImplemented)
        engine_info: zmq address/protocol dict for engine connections
        client_info: zmq address/protocol dict for client connections
        """

        super(Hub, self).__init__(**kwargs)
        self.registration_timeout = max(5000, 2*self.heartmonitor.period)

        # validate connection dicts:
        for k,v in self.client_info.iteritems():
            if k == 'task':
                util.validate_url_container(v[1])
            else:
                util.validate_url_container(v)
        # util.validate_url_container(self.client_info)
        util.validate_url_container(self.engine_info)

        # register our callbacks
        self.query.on_recv(self.dispatch_query)
        self.monitor.on_recv(self.dispatch_monitor_traffic)

        self.heartmonitor.add_heart_failure_handler(self.handle_heart_failure)
        self.heartmonitor.add_new_heart_handler(self.handle_new_heart)

        self.monitor_handlers = {b'in' : self.save_queue_request,
                                b'out': self.save_queue_result,
                                b'intask': self.save_task_request,
                                b'outtask': self.save_task_result,
                                b'tracktask': self.save_task_destination,
                                b'incontrol': _passer,
                                b'outcontrol': _passer,
                                b'iopub': self.save_iopub_message,
        }

        self.query_handlers = {'queue_request': self.queue_status,
                                'result_request': self.get_results,
                                'history_request': self.get_history,
                                'db_request': self.db_query,
                                'purge_request': self.purge_results,
                                'load_request': self.check_load,
                                'resubmit_request': self.resubmit_task,
                                'shutdown_request': self.shutdown_request,
                                'registration_request' : self.register_engine,
                                'unregistration_request' : self.unregister_engine,
                                'connection_request': self.connection_request,
        }

        # ignore resubmit replies
        self.resubmit.on_recv(lambda msg: None, copy=False)

        self.log.info("hub::created hub")

    @property
    def _next_id(self):
        """gemerate a new ID.

        No longer reuse old ids, just count from 0."""
        newid = self._idcounter
        self._idcounter += 1
        return newid
        # 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
            return 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(cast_bytes(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

    #-----------------------------------------------------------------------------
    # dispatch methods (1 per stream)
    #-----------------------------------------------------------------------------

    @util.log_errors
    def dispatch_monitor_traffic(self, msg):
        """all ME and Task queue messages come through here, as well as
        IOPub traffic."""
        self.log.debug("monitor traffic: %r", msg[0])
        switch = msg[0]
        try:
            idents, msg = self.session.feed_identities(msg[1:])
        except ValueError:
            idents=[]
        if not idents:
            self.log.error("Monitor message without topic: %r", msg)
            return
        handler = self.monitor_handlers.get(switch, None)
        if handler is not None:
            handler(idents, msg)
        else:
            self.log.error("Unrecognized monitor topic: %r", switch)

    @util.log_errors
    def dispatch_query(self, msg):
        """Route registration requests and queries from clients."""
        try:
            idents, msg = self.session.feed_identities(msg)
        except ValueError:
            idents = []
        if not idents:
            self.log.error("Bad Query Message: %r", msg)
            return
        client_id = idents[0]
        try:
            msg = self.session.unserialize(msg, content=True)
        except Exception:
            content = error.wrap_exception()
            self.log.error("Bad Query Message: %r", msg, exc_info=True)
            self.session.send(self.query, "hub_error", ident=client_id,
                    content=content)
            return
        # print client_id, header, parent, content
        #switch on message type:
        msg_type = msg['header']['msg_type']
        self.log.info("client::client %r requested %r", client_id, msg_type)
        handler = self.query_handlers.get(msg_type, None)
        try:
            assert handler is not None, "Bad Message Type: %r" % msg_type
        except:
            content = error.wrap_exception()
            self.log.error("Bad Message Type: %r", msg_type, exc_info=True)
            self.session.send(self.query, "hub_error", ident=client_id,
                    content=content)
            return

        else:
            handler(idents, 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."""
        self.log.debug("heartbeat::handle_new_heart(%r)", heart)
        if heart not in self.incoming_registrations:
            self.log.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"""
        self.log.debug("heartbeat::handle_heart_failure(%r)", heart)
        eid = self.hearts.get(heart, None)
        queue = self.engines[eid].queue
        if eid is None or self.keytable[eid] in self.dead_engines:
            self.log.info("heartbeat::ignoring heart failure %r (not an engine or already dead)", 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:
            self.log.error("invalid identity prefix: %r", idents)
            return
        queue_id, client_id = idents[:2]
        try:
            msg = self.session.unserialize(msg)
        except Exception:
            self.log.error("queue::client %r sent invalid message to %r: %r", client_id, queue_id, msg, exc_info=True)
            return

        eid = self.by_ident.get(queue_id, None)
        if eid is None:
            self.log.error("queue::target %r not registered", queue_id)
            self.log.debug("queue::    valid are: %r", self.by_ident.keys())
            return
        record = init_record(msg)
        msg_id = record['msg_id']
        self.log.info("queue::client %r submitted request %r to %s", client_id, msg_id, eid)
        # Unicode in records
        record['engine_uuid'] = queue_id.decode('ascii')
        record['client_uuid'] = client_id.decode('ascii')
        record['queue'] = 'mux'

        try:
            # it's posible iopub arrived first:
            existing = self.db.get_record(msg_id)
            for key,evalue in existing.iteritems():
                rvalue = record.get(key, None)
                if evalue and rvalue and evalue != rvalue:
                    self.log.warn("conflicting initial state for record: %r:%r <%r> %r", msg_id, rvalue, key, evalue)
                elif evalue and not rvalue:
                    record[key] = evalue
            try:
                self.db.update_record(msg_id, record)
            except Exception:
                self.log.error("DB Error updating record %r", msg_id, exc_info=True)
        except KeyError:
            try:
                self.db.add_record(msg_id, record)
            except Exception:
                self.log.error("DB Error adding record %r", msg_id, exc_info=True)

        self.pending.add(msg_id)
        self.queues[eid].append(msg_id)

    def save_queue_result(self, idents, msg):
        if len(idents) < 2:
            self.log.error("invalid identity prefix: %r", idents)
            return

        client_id, queue_id = idents[:2]
        try:
            msg = self.session.unserialize(msg)
        except Exception:
            self.log.error("queue::engine %r sent invalid message to %r: %r",
                    queue_id, client_id, msg, exc_info=True)
            return

        eid = self.by_ident.get(queue_id, None)
        if eid is None:
            self.log.error("queue::unknown engine %r is sending a reply: ", queue_id)
            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)
            self.log.info("queue::request %r completed on %s", msg_id, eid)
        elif msg_id not in self.all_completed:
            # it could be a result from a dead engine that died before delivering the
            # result
            self.log.warn("queue:: unknown msg finished %r", msg_id)
            return
        # update record anyway, because the unregistration could have been premature
        rheader = msg['header']
        completed = rheader['date']
        started = rheader.get('started', None)
        result = {
            'result_header' : rheader,
            'result_content': msg['content'],
            'received': datetime.now(),
            'started' : started,
            'completed' : completed
        }

        result['result_buffers'] = msg['buffers']
        try:
            self.db.update_record(msg_id, result)
        except Exception:
            self.log.error("DB Error updating record %r", msg_id, exc_info=True)

    #--------------------- Task Queue Traffic ------------------------------

    def save_task_request(self, idents, msg):
        """Save the submission of a task."""
        client_id = idents[0]

        try:
            msg = self.session.unserialize(msg)
        except Exception:
            self.log.error("task::client %r sent invalid task message: %r",
                    client_id, msg, exc_info=True)
            return
        record = init_record(msg)

        record['client_uuid'] = client_id.decode('ascii')
        record['queue'] = 'task'
        header = msg['header']
        msg_id = header['msg_id']
        self.pending.add(msg_id)
        self.unassigned.add(msg_id)
        try:
            # it's posible iopub arrived first:
            existing = self.db.get_record(msg_id)
            if existing['resubmitted']:
                for key in ('submitted', 'client_uuid', 'buffers'):
                    # don't clobber these keys on resubmit
                    # submitted and client_uuid should be different
                    # and buffers might be big, and shouldn't have changed
                    record.pop(key)
                    # still check content,header which should not change
                    # but are not expensive to compare as buffers

            for key,evalue in existing.iteritems():
                if key.endswith('buffers'):
                    # don't compare buffers
                    continue
                rvalue = record.get(key, None)
                if evalue and rvalue and evalue != rvalue:
                    self.log.warn("conflicting initial state for record: %r:%r <%r> %r", msg_id, rvalue, key, evalue)
                elif evalue and not rvalue:
                    record[key] = evalue
            try:
                self.db.update_record(msg_id, record)
            except Exception:
                self.log.error("DB Error updating record %r", msg_id, exc_info=True)
        except KeyError:
            try:
                self.db.add_record(msg_id, record)
            except Exception:
                self.log.error("DB Error adding record %r", msg_id, exc_info=True)
        except Exception:
            self.log.error("DB Error saving task request %r", msg_id, exc_info=True)

    def save_task_result(self, idents, msg):
        """save the result of a completed task."""
        client_id = idents[0]
        try:
            msg = self.session.unserialize(msg)
        except Exception:
            self.log.error("task::invalid task result message send to %r: %r",
                    client_id, msg, exc_info=True)
            return

        parent = msg['parent_header']
        if not parent:
            # print msg
            self.log.warn("Task %r had no parent!", msg)
            return
        msg_id = parent['msg_id']
        if msg_id in self.unassigned:
            self.unassigned.remove(msg_id)

        header = msg['header']
        engine_uuid = header.get('engine', u'')
        eid = self.by_ident.get(cast_bytes(engine_uuid), None)

        status = header.get('status', None)

        if msg_id in self.pending:
            self.log.info("task::task %r finished on %s", msg_id, eid)
            self.pending.remove(msg_id)
            self.all_completed.add(msg_id)
            if eid is not None:
                if status != 'aborted':
                    self.completed[eid].append(msg_id)
                if msg_id in self.tasks[eid]:
                    self.tasks[eid].remove(msg_id)
            completed = header['date']
            started = header.get('started', None)
            result = {
                'result_header' : header,
                'result_content': msg['content'],
                'started' : started,
                'completed' : completed,
                'received' : datetime.now(),
                'engine_uuid': engine_uuid,
            }

            result['result_buffers'] = msg['buffers']
            try:
                self.db.update_record(msg_id, result)
            except Exception:
                self.log.error("DB Error saving task request %r", msg_id, exc_info=True)

        else:
            self.log.debug("task::unknown task %r finished", msg_id)

    def save_task_destination(self, idents, msg):
        try:
            msg = self.session.unserialize(msg, content=True)
        except Exception:
            self.log.error("task::invalid task tracking message", exc_info=True)
            return
        content = msg['content']
        # print (content)
        msg_id = content['msg_id']
        engine_uuid = content['engine_id']
        eid = self.by_ident[cast_bytes(engine_uuid)]

        self.log.info("task::task %r arrived on %r", msg_id, eid)
        if msg_id in self.unassigned:
            self.unassigned.remove(msg_id)
        # else:
        #     self.log.debug("task::task %r 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))
        try:
            self.db.update_record(msg_id, dict(engine_uuid=engine_uuid))
        except Exception:
            self.log.error("DB Error saving task destination %r", msg_id, exc_info=True)

    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)

    #--------------------- IOPub Traffic ------------------------------

    def save_iopub_message(self, topics, msg):
        """save an iopub message into the db"""
        # print (topics)
        try:
            msg = self.session.unserialize(msg, content=True)
        except Exception:
            self.log.error("iopub::invalid IOPub message", exc_info=True)
            return

        parent = msg['parent_header']
        if not parent:
            self.log.warn("iopub::IOPub message lacks parent: %r", msg)
            return
        msg_id = parent['msg_id']
        msg_type = msg['header']['msg_type']
        content = msg['content']

        # ensure msg_id is in db
        try:
            rec = self.db.get_record(msg_id)
        except KeyError:
            rec = empty_record()
            rec['msg_id'] = msg_id
            self.db.add_record(msg_id, rec)
        # stream
        d = {}
        if msg_type == 'stream':
            name = content['name']
            s = rec[name] or ''
            d[name] = s + content['data']

        elif msg_type == 'pyerr':
            d['pyerr'] = content
        elif msg_type == 'pyin':
            d['pyin'] = content['code']
        elif msg_type in ('display_data', 'pyout'):
            d[msg_type] = content
        elif msg_type == 'status':
            pass
        else:
            self.log.warn("unhandled iopub msg_type: %r", msg_type)

        if not d:
            return

        try:
            self.db.update_record(msg_id, d)
        except Exception:
            self.log.error("DB Error saving iopub message %r", msg_id, exc_info=True)

    #-------------------------------------------------------------------------
    # Registration requests
    #-------------------------------------------------------------------------

    def connection_request(self, client_id, msg):
        """Reply with connection addresses for clients."""
        self.log.info("client::client %r connected", client_id)
        content = dict(status='ok')
        content.update(self.client_info)
        jsonable = {}
        for k,v in self.keytable.iteritems():
            if v not in self.dead_engines:
                jsonable[str(k)] = v.decode('ascii')
        content['engines'] = jsonable
        self.session.send(self.query, 'connection_reply', content, parent=msg, ident=client_id)

    def register_engine(self, reg, msg):
        """Register a new engine."""
        content = msg['content']
        try:
            queue = cast_bytes(content['queue'])
        except KeyError:
            self.log.error("registration::queue not specified", exc_info=True)
            return
        heart = content.get('heartbeat', None)
        if heart:
            heart = cast_bytes(heart)
        """register a new engine, and create the socket(s) necessary"""
        eid = self._next_id
        # print (eid, queue, reg, heart)

        self.log.debug("registration::register_engine(%i, %r, %r, %r)", eid, queue, reg, heart)

        content = dict(id=eid,status='ok')
        content.update(self.engine_info)
        # check if requesting available IDs:
        if queue in self.by_ident:
            try:
                raise KeyError("queue_id %r in use" % queue)
            except:
                content = error.wrap_exception()
                self.log.error("queue_id %r in use", queue, exc_info=True)
        elif heart in self.hearts: # need to check unique hearts?
            try:
                raise KeyError("heart_id %r in use" % heart)
            except:
                self.log.error("heart_id %r in use", heart, exc_info=True)
                content = error.wrap_exception()
        else:
            for h, pack in self.incoming_registrations.iteritems():
                if heart == h:
                    try:
                        raise KeyError("heart_id %r in use" % heart)
                    except:
                        self.log.error("heart_id %r in use", heart, exc_info=True)
                        content = error.wrap_exception()
                    break
                elif queue == pack[1]:
                    try:
                        raise KeyError("queue_id %r in use" % queue)
                    except:
                        self.log.error("queue_id %r in use", queue, exc_info=True)
                        content = error.wrap_exception()
                    break

        msg = self.session.send(self.query, "registration_reply",
                content=content,
                ident=reg)

        if content['status'] == 'ok':
            if heart in self.heartmonitor.hearts:
                # already beating
                self.incoming_registrations[heart] = (eid,queue,reg[0],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[0],dc)
        else:
            self.log.error("registration::registration %i failed: %r", 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:
            self.log.error("registration::bad engine id for unregistration: %r", ident, exc_info=True)
            return
        self.log.info("registration::unregister_engine(%r)", eid)
        # print (eid)
        uuid = self.keytable[eid]
        content=dict(id=eid, queue=uuid.decode('ascii'))
        self.dead_engines.add(uuid)
        # self.ids.remove(eid)
        # uuid = self.keytable.pop(eid)
        #
        # ec = self.engines.pop(eid)
        # self.hearts.pop(ec.heartbeat)
        # self.by_ident.pop(ec.queue)
        # self.completed.pop(eid)
        handleit = lambda : self._handle_stranded_msgs(eid, uuid)
        dc = ioloop.DelayedCallback(handleit, self.registration_timeout, self.loop)
        dc.start()
        ############## TODO: HANDLE IT ################

        if self.notifier:
            self.session.send(self.notifier, "unregistration_notification", content=content)

    def _handle_stranded_msgs(self, eid, uuid):
        """Handle messages known to be on an engine when the engine unregisters.

        It is possible that this will fire prematurely - that is, an engine will
        go down after completing a result, and the client will be notified
        that the result failed and later receive the actual result.
        """

        outstanding = self.queues[eid]

        for msg_id in outstanding:
            self.pending.remove(msg_id)
            self.all_completed.add(msg_id)
            try:
                raise error.EngineError("Engine %r died while running task %r" % (eid, msg_id))
            except:
                content = error.wrap_exception()
            # build a fake header:
            header = {}
            header['engine'] = uuid
            header['date'] = datetime.now()
            rec = dict(result_content=content, result_header=header, result_buffers=[])
            rec['completed'] = header['date']
            rec['engine_uuid'] = uuid
            try:
                self.db.update_record(msg_id, rec)
            except Exception:
                self.log.error("DB Error handling stranded msg %r", msg_id, exc_info=True)

    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:
            self.log.error("registration::tried to finish nonexistant registration", exc_info=True)
            return
        self.log.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(id=eid, queue=queue, registration=reg,
                                    control=control, heartbeat=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.decode('ascii'))
        if self.notifier:
            self.session.send(self.notifier, "registration_notification", content=content)
        self.log.info("engine::Engine Connected: %i", eid)

    def _purge_stalled_registration(self, heart):
        if heart in self.incoming_registrations:
            eid = self.incoming_registrations.pop(heart)[0]
            self.log.info("registration::purging stalled registration: %i", eid)
        else:
            pass

    #-------------------------------------------------------------------------
    # Client Requests
    #-------------------------------------------------------------------------

    def shutdown_request(self, client_id, msg):
        """handle shutdown request."""
        self.session.send(self.query, 'shutdown_reply', content={'status': 'ok'}, ident=client_id)
        # also notify other clients of shutdown
        self.session.send(self.notifier, 'shutdown_notice', content={'status': 'ok'})
        dc = ioloop.DelayedCallback(lambda : self._shutdown(), 1000, self.loop)
        dc.start()

    def _shutdown(self):
        self.log.info("hub::hub 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 = error.wrap_exception()
            self.session.send(self.query, "hub_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.query, "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 = error.wrap_exception()
            self.session.send(self.query, "hub_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[str(t)] = {'queue': queue, 'completed': completed , 'tasks': tasks}
        content['unassigned'] = list(self.unassigned) if verbose else len(self.unassigned)
        # print (content)
        self.session.send(self.query, "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']
        self.log.info("Dropping records with %s", content)
        msg_ids = content.get('msg_ids', [])
        reply = dict(status='ok')
        if msg_ids == 'all':
            try:
                self.db.drop_matching_records(dict(completed={'$ne':None}))
            except Exception:
                reply = error.wrap_exception()
        else:
            pending = filter(lambda m: m in self.pending, msg_ids)
            if pending:
                try:
                    raise IndexError("msg pending: %r" % pending[0])
                except:
                    reply = error.wrap_exception()
            else:
                try:
                    self.db.drop_matching_records(dict(msg_id={'$in':msg_ids}))
                except Exception:
                    reply = error.wrap_exception()

            if reply['status'] == 'ok':
                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 = error.wrap_exception()
                        break
                    uid = self.engines[eid].queue
                    try:
                        self.db.drop_matching_records(dict(engine_uuid=uid, completed={'$ne':None}))
                    except Exception:
                        reply = error.wrap_exception()
                        break

        self.session.send(self.query, 'purge_reply', content=reply, ident=client_id)

    def resubmit_task(self, client_id, msg):
        """Resubmit one or more tasks."""
        def finish(reply):
            self.session.send(self.query, 'resubmit_reply', content=reply, ident=client_id)

        content = msg['content']
        msg_ids = content['msg_ids']
        reply = dict(status='ok')
        try:
            records = self.db.find_records({'msg_id' : {'$in' : msg_ids}}, keys=[
                'header', 'content', 'buffers'])
        except Exception:
            self.log.error('db::db error finding tasks to resubmit', exc_info=True)
            return finish(error.wrap_exception())

        # validate msg_ids
        found_ids = [ rec['msg_id'] for rec in records ]
        pending_ids = [ msg_id for msg_id in found_ids if msg_id in self.pending ]
        if len(records) > len(msg_ids):
            try:
                raise RuntimeError("DB appears to be in an inconsistent state."
                    "More matching records were found than should exist")
            except Exception:
                return finish(error.wrap_exception())
        elif len(records) < len(msg_ids):
            missing = [ m for m in msg_ids if m not in found_ids ]
            try:
                raise KeyError("No such msg(s): %r" % missing)
            except KeyError:
                return finish(error.wrap_exception())
        elif pending_ids:
            pass
            # no need to raise on resubmit of pending task, now that we
            # resubmit under new ID, but do we want to raise anyway?
            # msg_id = invalid_ids[0]
            # try:
            #     raise ValueError("Task(s) %r appears to be inflight" % )
            # except Exception:
            #     return finish(error.wrap_exception())

        # mapping of original IDs to resubmitted IDs
        resubmitted = {}

        # send the messages
        for rec in records:
            header = rec['header']
            msg = self.session.msg(header['msg_type'], parent=header)
            msg_id = msg['msg_id']
            msg['content'] = rec['content']

            # use the old header, but update msg_id and timestamp
            fresh = msg['header']
            header['msg_id'] = fresh['msg_id']
            header['date'] = fresh['date']
            msg['header'] = header

            self.session.send(self.resubmit, msg, buffers=rec['buffers'])

            resubmitted[rec['msg_id']] = msg_id
            self.pending.add(msg_id)
            msg['buffers'] = rec['buffers']
            try:
                self.db.add_record(msg_id, init_record(msg))
            except Exception:
                self.log.error("db::DB Error updating record: %s", msg_id, exc_info=True)

        finish(dict(status='ok', resubmitted=resubmitted))

        # store the new IDs in the Task DB
        for msg_id, resubmit_id in resubmitted.iteritems():
            try:
                self.db.update_record(msg_id, {'resubmitted' : resubmit_id})
            except Exception:
                self.log.error("db::DB Error updating record: %s", msg_id, exc_info=True)

    def _extract_record(self, rec):
        """decompose a TaskRecord dict into subsection of reply for get_result"""
        io_dict = {}
        for key in ('pyin', 'pyout', 'pyerr', 'stdout', 'stderr'):
                io_dict[key] = rec[key]
        content = { 'result_content': rec['result_content'],
                            'header': rec['header'],
                            'result_header' : rec['result_header'],
                            'received' : rec['received'],
                            'io' : io_dict,
                          }
        if rec['result_buffers']:
            buffers = map(bytes, rec['result_buffers'])
        else:
            buffers = []

        return content, buffers

    def get_results(self, client_id, msg):
        """Get the result of 1 or more messages."""
        content = msg['content']
        msg_ids = sorted(set(content['msg_ids']))
        statusonly = content.get('status_only', False)
        pending = []
        completed = []
        content = dict(status='ok')
        content['pending'] = pending
        content['completed'] = completed
        buffers = []
        if not statusonly:
            try:
                matches = self.db.find_records(dict(msg_id={'$in':msg_ids}))
                # turn match list into dict, for faster lookup
                records = {}
                for rec in matches:
                    records[rec['msg_id']] = rec
            except Exception:
                content = error.wrap_exception()
                self.session.send(self.query, "result_reply", content=content,
                                                    parent=msg, ident=client_id)
                return
        else:
            records = {}
        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:
                    c,bufs = self._extract_record(records[msg_id])
                    content[msg_id] = c
                    buffers.extend(bufs)
            elif msg_id in records:
                if rec['completed']:
                    completed.append(msg_id)
                    c,bufs = self._extract_record(records[msg_id])
                    content[msg_id] = c
                    buffers.extend(bufs)
                else:
                    pending.append(msg_id)
            else:
                try:
                    raise KeyError('No such message: '+msg_id)
                except:
                    content = error.wrap_exception()
                break
        self.session.send(self.query, "result_reply", content=content,
                                            parent=msg, ident=client_id,
                                            buffers=buffers)

    def get_history(self, client_id, msg):
        """Get a list of all msg_ids in our DB records"""
        try:
            msg_ids = self.db.get_history()
        except Exception as e:
            content = error.wrap_exception()
        else:
            content = dict(status='ok', history=msg_ids)

        self.session.send(self.query, "history_reply", content=content,
                                            parent=msg, ident=client_id)

    def db_query(self, client_id, msg):
        """Perform a raw query on the task record database."""
        content = msg['content']
        query = content.get('query', {})
        keys = content.get('keys', None)
        buffers = []
        empty = list()
        try:
            records = self.db.find_records(query, keys)
        except Exception as e:
            content = error.wrap_exception()
        else:
            # extract buffers from reply content:
            if keys is not None:
                buffer_lens = [] if 'buffers' in keys else None
                result_buffer_lens = [] if 'result_buffers' in keys else None
            else:
                buffer_lens = None
                result_buffer_lens = None

            for rec in records:
                # buffers may be None, so double check
                b = rec.pop('buffers', empty) or empty
                if buffer_lens is not None:
                    buffer_lens.append(len(b))
                    buffers.extend(b)
                rb = rec.pop('result_buffers', empty) or empty
                if result_buffer_lens is not None:
                    result_buffer_lens.append(len(rb))
                    buffers.extend(rb)
            content = dict(status='ok', records=records, buffer_lens=buffer_lens,
                                    result_buffer_lens=result_buffer_lens)
        # self.log.debug (content)
        self.session.send(self.query, "db_reply", content=content,
                                            parent=msg, ident=client_id,
                                            buffers=buffers)


# -*- coding:utf-8 -*-
import subprocess as sp
import os
import time

from contextlib import contextmanager
from collections import OrderedDict, namedtuple

ParsedCompletedCommand = namedtuple(
    'ParsedCompletedCommand',
    ['returncode', 'args', 'stdout', 'stderr']
)

class Rscript:
    """
    Using subprocess to call Rscript execute R program
    :type file: str, target R file
    :type cmd: str, parameter of target R file
    :type folder: str, default '', only support the sub folder

    example:
    >>> import os
    >>> from pyRscript import pyRscript
    >>> RPATH = '/home/pyRscript' # also ex: os.path.join(os.path.dirname(os.path.abspath('test.R')), 'R')
    >>> cmd = '-d,SN1234,-t,2017-07-13 00:00:00'
    >>> r = pyRscript.Rscript(path=RPATH, file='test.R', cmd=cmd)
    >>> print(r)
    >>> <Parameter test.R /home/pyRscript ['-d', 'SN1234', '-t', '2017-07-13 00:00:00']>
    >>> ret = r.execute()
    >>> print(ret)
    >>> ParsedCompletedCommand(returncode=0, args=['Rscript', 'test.R', '-d', 'SN1234', '-t', '2017-07-13 00:00:00'], stdout='2017-12-03 11:35:38 INFO::Execute R\n2017-12-03 11:35:39 INFO::id: SN1234, time: 2017-07-13 00:00:00', stderr='Loading required package: methods')
    >>> print(ret.returncode)
    >>> 0
    >>> print(ret.stdout)
    >>> '2017-12-03 11:37:59 INFO::Execute R\n2017-12-03 11:37:59 INFO::id: SN1234, time: 2017-07-13 00:00:00'
    >>> print(ret.stderr)
    >>> 'Loading required package: methods'
    """

    def __init__(self, path, file, cmd=None):
        super(Rscript, self).__init__()
        self.file = file
        self.path = path
        if cmd is None:
            self.cmd = []
        else:
            self.cmd = cmd.split(',')

    @contextmanager
    def cd(self, newdir):
        """
        go to the path
        """
        prevdir = os.getcwd()
        os.chdir(newdir)
        try:
            yield
        finally:
            os.chdir(prevdir)

    def decode_cmd_out(self, completed_cmd):
        """
        return a standard message
        """
        try:
            stdout = completed_cmd.stdout.encode('utf-8').decode()
        except AttributeError:
            try:
                stdout = str(bytes(completed_cmd.stdout), 'big5').strip()
            except AttributeError:
                stdout = str(bytes(completed_cmd.stdout).decode('utf-8')).strip()
        try:
            stderr = completed_cmd.stderr.encode('utf-8').decode()
        except AttributeError:
            try:
                stderr = str(bytes(completed_cmd.stderr), 'big5').strip()
            except AttributeError:
                stderr = str(bytes(completed_cmd.stderr).decode('utf-8')).strip()
        return ParsedCompletedCommand(
            completed_cmd.returncode,
            completed_cmd.args,
            stdout,
            stderr
        )

    def run_command_under_r_root(self, cmd, catched=True):
        """
        subprocess run on here
        """
        RPATH = self.path
        with self.cd(newdir=RPATH):
            if catched:
                process = sp.run(cmd, stdout=sp.PIPE, stderr=sp.PIPE)
            else:
                process = sp.run(cmd)
            return process

    def execute(self):
        """
        Execute R script
        """
        rprocess = OrderedDict()
        commands = OrderedDict([
            (self.file, ['Rscript', self.file] + self.cmd),
        ])
        for cmd_name, cmd in commands.items():
            rprocess[cmd_name] = self.run_command_under_r_root(cmd)

        return self.decode_cmd_out(completed_cmd=rprocess[self.file])

    def __repr__(self):
        fmt = f'<Parameter {self.file} {self.path} {self.cmd}>'
        return fmt


""" Defines a convenient mix-in class for implementing Qt frontends.
"""

class BaseFrontendMixin(object):
    """ A mix-in class for implementing Qt frontends.

    To handle messages of a particular type, frontends need only define an
    appropriate handler method. For example, to handle 'stream' messaged, define
    a '_handle_stream(msg)' method.
    """

    #---------------------------------------------------------------------------
    # 'BaseFrontendMixin' concrete interface
    #---------------------------------------------------------------------------

    def _get_kernel_manager(self):
        """ Returns the current kernel manager.
        """
        return self._kernel_manager

    def _set_kernel_manager(self, kernel_manager):
        """ Disconnect from the current kernel manager (if any) and set a new
            kernel manager.
        """
        # Disconnect the old kernel manager, if necessary.
        old_manager = self._kernel_manager
        if old_manager is not None:
            old_manager.started_kernel.disconnect(self._started_kernel)
            old_manager.started_channels.disconnect(self._started_channels)
            old_manager.stopped_channels.disconnect(self._stopped_channels)

            # Disconnect the old kernel manager's channels.
            old_manager.sub_channel.message_received.disconnect(self._dispatch)
            old_manager.shell_channel.message_received.disconnect(self._dispatch)
            old_manager.stdin_channel.message_received.disconnect(self._dispatch)
            old_manager.hb_channel.kernel_died.disconnect(
                self._handle_kernel_died)

            # Handle the case where the old kernel manager is still listening.
            if old_manager.channels_running:
                self._stopped_channels()

        # Set the new kernel manager.
        self._kernel_manager = kernel_manager
        if kernel_manager is None:
            return

        # Connect the new kernel manager.
        kernel_manager.started_kernel.connect(self._started_kernel)
        kernel_manager.started_channels.connect(self._started_channels)
        kernel_manager.stopped_channels.connect(self._stopped_channels)

        # Connect the new kernel manager's channels.
        kernel_manager.sub_channel.message_received.connect(self._dispatch)
        kernel_manager.shell_channel.message_received.connect(self._dispatch)
        kernel_manager.stdin_channel.message_received.connect(self._dispatch)
        kernel_manager.hb_channel.kernel_died.connect(self._handle_kernel_died)

        # Handle the case where the kernel manager started channels before
        # we connected.
        if kernel_manager.channels_running:
            self._started_channels()

    kernel_manager = property(_get_kernel_manager, _set_kernel_manager)

    #---------------------------------------------------------------------------
    # 'BaseFrontendMixin' abstract interface
    #---------------------------------------------------------------------------

    def _handle_kernel_died(self, since_last_heartbeat):
        """ This is called when the ``kernel_died`` signal is emitted.

        This method is called when the kernel heartbeat has not been
        active for a certain amount of time. The typical action will be to
        give the user the option of restarting the kernel.

        Parameters
        ----------
        since_last_heartbeat : float
            The time since the heartbeat was last received.
        """

    def _started_kernel(self):
        """Called when the KernelManager starts (or restarts) the kernel subprocess.
        Channels may or may not be running at this point.
        """

    def _started_channels(self):
        """ Called when the KernelManager channels have started listening or
            when the frontend is assigned an already listening KernelManager.
        """

    def _stopped_channels(self):
        """ Called when the KernelManager channels have stopped listening or
            when a listening KernelManager is removed from the frontend.
        """

    #---------------------------------------------------------------------------
    # 'BaseFrontendMixin' protected interface
    #---------------------------------------------------------------------------

    def _dispatch(self, msg):
        """ Calls the frontend handler associated with the message type of the
            given message.
        """
        msg_type = msg['header']['msg_type']
        handler = getattr(self, '_handle_' + msg_type, None)
        if handler:
            handler(msg)

    def _is_from_this_session(self, msg):
        """ Returns whether a reply from the kernel originated from a request
            from this frontend.
        """
        session = self._kernel_manager.session.session
        parent = msg['parent_header']
        if not parent:
            # if the message has no parent, assume it is meant for all frontends
            return True
        else:
            return parent.get('session') == session


"""Source file annotation for Coverage."""

import os, re

from coverage.backward import sorted                    # pylint: disable=W0622
from coverage.report import Reporter

class AnnotateReporter(Reporter):
    """Generate annotated source files showing line coverage.

    This reporter creates annotated copies of the measured source files. Each
    .py file is copied as a .py,cover file, with a left-hand margin annotating
    each line::

        > def h(x):
        -     if 0:   #pragma: no cover
        -         pass
        >     if x == 1:
        !         a = 1
        >     else:
        >         a = 2

        > h(2)

    Executed lines use '>', lines not executed use '!', lines excluded from
    consideration use '-'.

    """

    def __init__(self, coverage, config):
        super(AnnotateReporter, self).__init__(coverage, config)
        self.directory = None

    blank_re = re.compile(r"\s*(#|$)")
    else_re = re.compile(r"\s*else\s*:\s*(#|$)")

    def report(self, morfs, directory=None):
        """Run the report.

        See `coverage.report()` for arguments.

        """
        self.report_files(self.annotate_file, morfs, directory)

    def annotate_file(self, cu, analysis):
        """Annotate a single file.

        `cu` is the CodeUnit for the file to annotate.

        """
        if not cu.relative:
            return

        filename = cu.filename
        source = cu.source_file()
        if self.directory:
            dest_file = os.path.join(self.directory, cu.flat_rootname())
            dest_file += ".py,cover"
        else:
            dest_file = filename + ",cover"
        dest = open(dest_file, 'w')

        statements = sorted(analysis.statements)
        missing = sorted(analysis.missing)
        excluded = sorted(analysis.excluded)

        lineno = 0
        i = 0
        j = 0
        covered = True
        while True:
            line = source.readline()
            if line == '':
                break
            lineno += 1
            while i < len(statements) and statements[i] < lineno:
                i += 1
            while j < len(missing) and missing[j] < lineno:
                j += 1
            if i < len(statements) and statements[i] == lineno:
                covered = j >= len(missing) or missing[j] > lineno
            if self.blank_re.match(line):
                dest.write('  ')
            elif self.else_re.match(line):
                # Special logic for lines containing only 'else:'.
                if i >= len(statements) and j >= len(missing):
                    dest.write('! ')
                elif i >= len(statements) or j >= len(missing):
                    dest.write('> ')
                elif statements[i] == missing[j]:
                    dest.write('! ')
                else:
                    dest.write('> ')
            elif lineno in excluded:
                dest.write('- ')
            elif covered:
                dest.write('> ')
            else:
                dest.write('! ')
            dest.write(line)
        source.close()
        dest.close()


import subprocess
import shlex, re
import os
from subprocess import CalledProcessError
from datetime import *

def find(name):
    '''
        returns a list of tuples (package_name, description) for apt-cache search results
    '''
    cmd = 'apt-cache search %s' % name
    args = shlex.split(cmd)
    try:
        output = subprocess.check_output(args)
    except CalledProcessError:
        return []
    lines = output.splitlines()
    packages = []
    for line in lines:
        package, _, desc = line.partition(' - ')
        packages.append((package, desc))
    return packages

def is_installed(name):
    '''
        returns True if package (name) is installed
    '''
    return True if get_installed_version(name) else False

def get_installed_version(name):
    '''
        returns installed package version and None if package is not installed
    '''
    pattern = re.compile(r'''Installed:\s+(?P<version>.*)''')
    cmd = 'apt-cache policy %s' % name
    args = shlex.split(cmd)
    try:
        output = subprocess.check_output(args)
        if not output:
            return None
    except CalledProcessError:
        return None
    # check output
    match = pattern.search(output)
    if match:
        version = match.groupdict()['version']
        if version == '(none)':
            return None
        else:
            return version

def is_partially_installed(filename, final_destination):
    try:
        files = os.popen("ls /var/cache/apt/archives/%s" % filename).readlines()
    except e:
        return False
    if len(files)>0:
        file_path = files[0].rstrip()
        file_installed = os.path.exists(final_destination)
        return os.path.exists(file_path) and not file_installed
    else:
        return False

def check_in_source_list(source):
    file_path = '/etc/apt/sources.list'
    lines=os.popen("grep %s %s" %(source, file_path)).readlines()
    if len(lines) > 0: # grep has found the source
        for line in lines:
            if not line.startswith("#"):
                return True
    else:
        return False

def is_recently_updated():
    latest_modification = os.popen("stat /var/lib/apt/lists/ | grep Modify | cut -d: -f2| cut -d \" \" -f2").readline().strip()
    formated_date = latest_modification.split('-')
    modification_datetime = datetime(int(formated_date[0]), int(formated_date[1]), int(formated_date[2]))
    delta = datetime.now() - modification_datetime
    if delta.days > 1:
        return False
    else:
        return True


"""Session object for building, serializing, sending, and receiving messages in
IPython. The Session object supports serialization, HMAC signatures, and
metadata on messages.

Also defined here are utilities for working with Sessions:
* A SessionFactory to be used as a base class for configurables that work with
Sessions.
* A Message object for convenience that allows attribute-access to the msg dict.

Authors:

* Min RK
* Brian Granger
* Fernando Perez
"""
#-----------------------------------------------------------------------------
#  Copyright (C) 2010-2011  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
#-----------------------------------------------------------------------------

import hmac
import logging
import os
import pprint
import uuid
from datetime import datetime

try:
    import cPickle
    pickle = cPickle
except:
    cPickle = None
    import pickle

import zmq
from zmq.utils import jsonapi
from zmq.eventloop.ioloop import IOLoop
from zmq.eventloop.zmqstream import ZMQStream

from IPython.config.application import Application, boolean_flag
from IPython.config.configurable import Configurable, LoggingConfigurable
from IPython.utils.importstring import import_item
from IPython.utils.jsonutil import extract_dates, squash_dates, date_default
from IPython.utils.py3compat import str_to_bytes
from IPython.utils.traitlets import (CBytes, Unicode, Bool, Any, Instance, Set,
                                        DottedObjectName, CUnicode)

#-----------------------------------------------------------------------------
# utility functions
#-----------------------------------------------------------------------------

def squash_unicode(obj):
    """coerce unicode back to bytestrings."""
    if isinstance(obj,dict):
        for key in obj.keys():
            obj[key] = squash_unicode(obj[key])
            if isinstance(key, unicode):
                obj[squash_unicode(key)] = obj.pop(key)
    elif isinstance(obj, list):
        for i,v in enumerate(obj):
            obj[i] = squash_unicode(v)
    elif isinstance(obj, unicode):
        obj = obj.encode('utf8')
    return obj

#-----------------------------------------------------------------------------
# globals and defaults
#-----------------------------------------------------------------------------

# ISO8601-ify datetime objects
json_packer = lambda obj: jsonapi.dumps(obj, default=date_default)
json_unpacker = lambda s: extract_dates(jsonapi.loads(s))

pickle_packer = lambda o: pickle.dumps(o,-1)
pickle_unpacker = pickle.loads

default_packer = json_packer
default_unpacker = json_unpacker

DELIM=b"<IDS|MSG>"

#-----------------------------------------------------------------------------
# Mixin tools for apps that use Sessions
#-----------------------------------------------------------------------------

session_aliases = dict(
    ident = 'Session.session',
    user = 'Session.username',
    keyfile = 'Session.keyfile',
)

session_flags  = {
    'secure' : ({'Session' : { 'key' : str_to_bytes(str(uuid.uuid4())),
                            'keyfile' : '' }},
        """Use HMAC digests for authentication of messages.
        Setting this flag will generate a new UUID to use as the HMAC key.
        """),
    'no-secure' : ({'Session' : { 'key' : b'', 'keyfile' : '' }},
        """Don't authenticate messages."""),
}

def default_secure(cfg):
    """Set the default behavior for a config environment to be secure.

    If Session.key/keyfile have not been set, set Session.key to
    a new random UUID.
    """

    if 'Session' in cfg:
        if 'key' in cfg.Session or 'keyfile' in cfg.Session:
            return
    # key/keyfile not specified, generate new UUID:
    cfg.Session.key = str_to_bytes(str(uuid.uuid4()))

#-----------------------------------------------------------------------------
# Classes
#-----------------------------------------------------------------------------

class SessionFactory(LoggingConfigurable):
    """The Base class for configurables that have a Session, Context, logger,
    and IOLoop.
    """

    logname = Unicode('')
    def _logname_changed(self, name, old, new):
        self.log = logging.getLogger(new)

    # not configurable:
    context = Instance('zmq.Context')
    def _context_default(self):
        return zmq.Context.instance()

    session = Instance('IPython.zmq.session.Session')

    loop = Instance('zmq.eventloop.ioloop.IOLoop', allow_none=False)
    def _loop_default(self):
        return IOLoop.instance()

    def __init__(self, **kwargs):
        super(SessionFactory, self).__init__(**kwargs)

        if self.session is None:
            # construct the session
            self.session = Session(**kwargs)

class Message(object):
    """A simple message object that maps dict keys to attributes.

    A Message can be created from a dict and a dict from a Message instance
    simply by calling dict(msg_obj)."""

    def __init__(self, msg_dict):
        dct = self.__dict__
        for k, v in dict(msg_dict).iteritems():
            if isinstance(v, dict):
                v = Message(v)
            dct[k] = v

    # Having this iterator lets dict(msg_obj) work out of the box.
    def __iter__(self):
        return iter(self.__dict__.iteritems())

    def __repr__(self):
        return repr(self.__dict__)

    def __str__(self):
        return pprint.pformat(self.__dict__)

    def __contains__(self, k):
        return k in self.__dict__

    def __getitem__(self, k):
        return self.__dict__[k]

def msg_header(msg_id, msg_type, username, session):
    date = datetime.now()
    return locals()

def extract_header(msg_or_header):
    """Given a message or header, return the header."""
    if not msg_or_header:
        return {}
    try:
        # See if msg_or_header is the entire message.
        h = msg_or_header['header']
    except KeyError:
        try:
            # See if msg_or_header is just the header
            h = msg_or_header['msg_id']
        except KeyError:
            raise
        else:
            h = msg_or_header
    if not isinstance(h, dict):
        h = dict(h)
    return h

class Session(Configurable):
    """Object for handling serialization and sending of messages.

    The Session object handles building messages and sending them
    with ZMQ sockets or ZMQStream objects.  Objects can communicate with each
    other over the network via Session objects, and only need to work with the
    dict-based IPython message spec. The Session will handle
    serialization/deserialization, security, and metadata.

    Sessions support configurable serialiization via packer/unpacker traits,
    and signing with HMAC digests via the key/keyfile traits.

    Parameters
    ----------

    debug : bool
        whether to trigger extra debugging statements
    packer/unpacker : str : 'json', 'pickle' or import_string
        importstrings for methods to serialize message parts.  If just
        'json' or 'pickle', predefined JSON and pickle packers will be used.
        Otherwise, the entire importstring must be used.

        The functions must accept at least valid JSON input, and output *bytes*.

        For example, to use msgpack:
        packer = 'msgpack.packb', unpacker='msgpack.unpackb'
    pack/unpack : callables
        You can also set the pack/unpack callables for serialization directly.
    session : bytes
        the ID of this Session object.  The default is to generate a new UUID.
    username : unicode
        username added to message headers.  The default is to ask the OS.
    key : bytes
        The key used to initialize an HMAC signature.  If unset, messages
        will not be signed or checked.
    keyfile : filepath
        The file containing a key.  If this is set, `key` will be initialized
        to the contents of the file.

    """

    debug=Bool(False, config=True, help="""Debug output in the Session""")

    packer = DottedObjectName('json',config=True,
            help="""The name of the packer for serializing messages.
            Should be one of 'json', 'pickle', or an import name
            for a custom callable serializer.""")
    def _packer_changed(self, name, old, new):
        if new.lower() == 'json':
            self.pack = json_packer
            self.unpack = json_unpacker
        elif new.lower() == 'pickle':
            self.pack = pickle_packer
            self.unpack = pickle_unpacker
        else:
            self.pack = import_item(str(new))

    unpacker = DottedObjectName('json', config=True,
        help="""The name of the unpacker for unserializing messages.
        Only used with custom functions for `packer`.""")
    def _unpacker_changed(self, name, old, new):
        if new.lower() == 'json':
            self.pack = json_packer
            self.unpack = json_unpacker
        elif new.lower() == 'pickle':
            self.pack = pickle_packer
            self.unpack = pickle_unpacker
        else:
            self.unpack = import_item(str(new))

    session = CUnicode(u'', config=True,
        help="""The UUID identifying this session.""")
    def _session_default(self):
        u = unicode(uuid.uuid4())
        self.bsession = u.encode('ascii')
        return u

    def _session_changed(self, name, old, new):
        self.bsession = self.session.encode('ascii')

    # bsession is the session as bytes
    bsession = CBytes(b'')

    username = Unicode(os.environ.get('USER',u'username'), config=True,
        help="""Username for the Session. Default is your system username.""")

    # message signature related traits:

    key = CBytes(b'', config=True,
        help="""execution key, for extra authentication.""")
    def _key_changed(self, name, old, new):
        if new:
            self.auth = hmac.HMAC(new)
        else:
            self.auth = None
    auth = Instance(hmac.HMAC)
    digest_history = Set()

    keyfile = Unicode('', config=True,
        help="""path to file containing execution key.""")
    def _keyfile_changed(self, name, old, new):
        with open(new, 'rb') as f:
            self.key = f.read().strip()

    # serialization traits:

    pack = Any(default_packer) # the actual packer function
    def _pack_changed(self, name, old, new):
        if not callable(new):
            raise TypeError("packer must be callable, not %s"%type(new))

    unpack = Any(default_unpacker) # the actual packer function
    def _unpack_changed(self, name, old, new):
        # unpacker is not checked - it is assumed to be
        if not callable(new):
            raise TypeError("unpacker must be callable, not %s"%type(new))

    def __init__(self, **kwargs):
        """create a Session object

        Parameters
        ----------

        debug : bool
            whether to trigger extra debugging statements
        packer/unpacker : str : 'json', 'pickle' or import_string
            importstrings for methods to serialize message parts.  If just
            'json' or 'pickle', predefined JSON and pickle packers will be used.
            Otherwise, the entire importstring must be used.

            The functions must accept at least valid JSON input, and output
            *bytes*.

            For example, to use msgpack:
            packer = 'msgpack.packb', unpacker='msgpack.unpackb'
        pack/unpack : callables
            You can also set the pack/unpack callables for serialization
            directly.
        session : unicode (must be ascii)
            the ID of this Session object.  The default is to generate a new
            UUID.
        bsession : bytes
            The session as bytes
        username : unicode
            username added to message headers.  The default is to ask the OS.
        key : bytes
            The key used to initialize an HMAC signature.  If unset, messages
            will not be signed or checked.
        keyfile : filepath
            The file containing a key.  If this is set, `key` will be
            initialized to the contents of the file.
        """
        super(Session, self).__init__(**kwargs)
        self._check_packers()
        self.none = self.pack({})
        # ensure self._session_default() if necessary, so bsession is defined:
        self.session

    @property
    def msg_id(self):
        """always return new uuid"""
        return str(uuid.uuid4())

    def _check_packers(self):
        """check packers for binary data and datetime support."""
        pack = self.pack
        unpack = self.unpack

        # check simple serialization
        msg = dict(a=[1,'hi'])
        try:
            packed = pack(msg)
        except Exception:
            raise ValueError("packer could not serialize a simple message")

        # ensure packed message is bytes
        if not isinstance(packed, bytes):
            raise ValueError("message packed to %r, but bytes are required"%type(packed))

        # check that unpack is pack's inverse
        try:
            unpacked = unpack(packed)
        except Exception:
            raise ValueError("unpacker could not handle the packer's output")

        # check datetime support
        msg = dict(t=datetime.now())
        try:
            unpacked = unpack(pack(msg))
        except Exception:
            self.pack = lambda o: pack(squash_dates(o))
            self.unpack = lambda s: extract_dates(unpack(s))

    def msg_header(self, msg_type):
        return msg_header(self.msg_id, msg_type, self.username, self.session)

    def msg(self, msg_type, content=None, parent=None, subheader=None, header=None):
        """Return the nested message dict.

        This format is different from what is sent over the wire. The
        serialize/unserialize methods converts this nested message dict to the wire
        format, which is a list of message parts.
        """
        msg = {}
        header = self.msg_header(msg_type) if header is None else header
        msg['header'] = header
        msg['msg_id'] = header['msg_id']
        msg['msg_type'] = header['msg_type']
        msg['parent_header'] = {} if parent is None else extract_header(parent)
        msg['content'] = {} if content is None else content
        sub = {} if subheader is None else subheader
        msg['header'].update(sub)
        return msg

    def sign(self, msg_list):
        """Sign a message with HMAC digest. If no auth, return b''.

        Parameters
        ----------
        msg_list : list
            The [p_header,p_parent,p_content] part of the message list.
        """
        if self.auth is None:
            return b''
        h = self.auth.copy()
        for m in msg_list:
            h.update(m)
        return str_to_bytes(h.hexdigest())

    def serialize(self, msg, ident=None):
        """Serialize the message components to bytes.

        This is roughly the inverse of unserialize. The serialize/unserialize
        methods work with full message lists, whereas pack/unpack work with
        the individual message parts in the message list.

        Parameters
        ----------
        msg : dict or Message
            The nexted message dict as returned by the self.msg method.

        Returns
        -------
        msg_list : list
            The list of bytes objects to be sent with the format:
            [ident1,ident2,...,DELIM,HMAC,p_header,p_parent,p_content,
             buffer1,buffer2,...]. In this list, the p_* entities are
            the packed or serialized versions, so if JSON is used, these
            are utf8 encoded JSON strings.
        """
        content = msg.get('content', {})
        if content is None:
            content = self.none
        elif isinstance(content, dict):
            content = self.pack(content)
        elif isinstance(content, bytes):
            # content is already packed, as in a relayed message
            pass
        elif isinstance(content, unicode):
            # should be bytes, but JSON often spits out unicode
            content = content.encode('utf8')
        else:
            raise TypeError("Content incorrect type: %s"%type(content))

        real_message = [self.pack(msg['header']),
                        self.pack(msg['parent_header']),
                        content
        ]

        to_send = []

        if isinstance(ident, list):
            # accept list of idents
            to_send.extend(ident)
        elif ident is not None:
            to_send.append(ident)
        to_send.append(DELIM)

        signature = self.sign(real_message)
        to_send.append(signature)

        to_send.extend(real_message)

        return to_send

    def send(self, stream, msg_or_type, content=None, parent=None, ident=None,
             buffers=None, subheader=None, track=False, header=None):
        """Build and send a message via stream or socket.

        The message format used by this function internally is as follows:

        [ident1,ident2,...,DELIM,HMAC,p_header,p_parent,p_content,
         buffer1,buffer2,...]

        The serialize/unserialize methods convert the nested message dict into this
        format.

        Parameters
        ----------

        stream : zmq.Socket or ZMQStream
            The socket-like object used to send the data.
        msg_or_type : str or Message/dict
            Normally, msg_or_type will be a msg_type unless a message is being
            sent more than once. If a header is supplied, this can be set to
            None and the msg_type will be pulled from the header.

        content : dict or None
            The content of the message (ignored if msg_or_type is a message).
        header : dict or None
            The header dict for the message (ignores if msg_to_type is a message).
        parent : Message or dict or None
            The parent or parent header describing the parent of this message
            (ignored if msg_or_type is a message).
        ident : bytes or list of bytes
            The zmq.IDENTITY routing path.
        subheader : dict or None
            Extra header keys for this message's header (ignored if msg_or_type
            is a message).
        buffers : list or None
            The already-serialized buffers to be appended to the message.
        track : bool
            Whether to track.  Only for use with Sockets, because ZMQStream
            objects cannot track messages.

        Returns
        -------
        msg : dict
            The constructed message.
        (msg,tracker) : (dict, MessageTracker)
            if track=True, then a 2-tuple will be returned,
            the first element being the constructed
            message, and the second being the MessageTracker

        """

        if not isinstance(stream, (zmq.Socket, ZMQStream)):
            raise TypeError("stream must be Socket or ZMQStream, not %r"%type(stream))
        elif track and isinstance(stream, ZMQStream):
            raise TypeError("ZMQStream cannot track messages")

        if isinstance(msg_or_type, (Message, dict)):
            # We got a Message or message dict, not a msg_type so don't
            # build a new Message.
            msg = msg_or_type
        else:
            msg = self.msg(msg_or_type, content=content, parent=parent,
                           subheader=subheader, header=header)

        buffers = [] if buffers is None else buffers
        to_send = self.serialize(msg, ident)
        flag = 0
        if buffers:
            flag = zmq.SNDMORE
            _track = False
        else:
            _track=track
        if track:
            tracker = stream.send_multipart(to_send, flag, copy=False, track=_track)
        else:
            tracker = stream.send_multipart(to_send, flag, copy=False)
        for b in buffers[:-1]:
            stream.send(b, flag, copy=False)
        if buffers:
            if track:
                tracker = stream.send(buffers[-1], copy=False, track=track)
            else:
                tracker = stream.send(buffers[-1], copy=False)

        # omsg = Message(msg)
        if self.debug:
            pprint.pprint(msg)
            pprint.pprint(to_send)
            pprint.pprint(buffers)

        msg['tracker'] = tracker

        return msg

    def send_raw(self, stream, msg_list, flags=0, copy=True, ident=None):
        """Send a raw message via ident path.

        This method is used to send a already serialized message.

        Parameters
        ----------
        stream : ZMQStream or Socket
            The ZMQ stream or socket to use for sending the message.
        msg_list : list
            The serialized list of messages to send. This only includes the
            [p_header,p_parent,p_content,buffer1,buffer2,...] portion of
            the message.
        ident : ident or list
            A single ident or a list of idents to use in sending.
        """
        to_send = []
        if isinstance(ident, bytes):
            ident = [ident]
        if ident is not None:
            to_send.extend(ident)

        to_send.append(DELIM)
        to_send.append(self.sign(msg_list))
        to_send.extend(msg_list)
        stream.send_multipart(msg_list, flags, copy=copy)

    def recv(self, socket, mode=zmq.NOBLOCK, content=True, copy=True):
        """Receive and unpack a message.

        Parameters
        ----------
        socket : ZMQStream or Socket
            The socket or stream to use in receiving.

        Returns
        -------
        [idents], msg
            [idents] is a list of idents and msg is a nested message dict of
            same format as self.msg returns.
        """
        if isinstance(socket, ZMQStream):
            socket = socket.socket
        try:
            msg_list = socket.recv_multipart(mode, copy=copy)
        except zmq.ZMQError as e:
            if e.errno == zmq.EAGAIN:
                # We can convert EAGAIN to None as we know in this case
                # recv_multipart won't return None.
                return None,None
            else:
                raise
        # split multipart message into identity list and message dict
        # invalid large messages can cause very expensive string comparisons
        idents, msg_list = self.feed_identities(msg_list, copy)
        try:
            return idents, self.unserialize(msg_list, content=content, copy=copy)
        except Exception as e:
            # TODO: handle it
            raise e

    def feed_identities(self, msg_list, copy=True):
        """Split the identities from the rest of the message.

        Feed until DELIM is reached, then return the prefix as idents and
        remainder as msg_list. This is easily broken by setting an IDENT to DELIM,
        but that would be silly.

        Parameters
        ----------
        msg_list : a list of Message or bytes objects
            The message to be split.
        copy : bool
            flag determining whether the arguments are bytes or Messages

        Returns
        -------
        (idents, msg_list) : two lists
            idents will always be a list of bytes, each of which is a ZMQ
            identity. msg_list will be a list of bytes or zmq.Messages of the
            form [HMAC,p_header,p_parent,p_content,buffer1,buffer2,...] and
            should be unpackable/unserializable via self.unserialize at this
            point.
        """
        if copy:
            idx = msg_list.index(DELIM)
            return msg_list[:idx], msg_list[idx+1:]
        else:
            failed = True
            for idx,m in enumerate(msg_list):
                if m.bytes == DELIM:
                    failed = False
                    break
            if failed:
                raise ValueError("DELIM not in msg_list")
            idents, msg_list = msg_list[:idx], msg_list[idx+1:]
            return [m.bytes for m in idents], msg_list

    def unserialize(self, msg_list, content=True, copy=True):
        """Unserialize a msg_list to a nested message dict.

        This is roughly the inverse of serialize. The serialize/unserialize
        methods work with full message lists, whereas pack/unpack work with
        the individual message parts in the message list.

        Parameters:
        -----------
        msg_list : list of bytes or Message objects
            The list of message parts of the form [HMAC,p_header,p_parent,
            p_content,buffer1,buffer2,...].
        content : bool (True)
            Whether to unpack the content dict (True), or leave it packed
            (False).
        copy : bool (True)
            Whether to return the bytes (True), or the non-copying Message
            object in each place (False).

        Returns
        -------
        msg : dict
            The nested message dict with top-level keys [header, parent_header,
            content, buffers].
        """
        minlen = 4
        message = {}
        if not copy:
            for i in range(minlen):
                msg_list[i] = msg_list[i].bytes
        if self.auth is not None:
            signature = msg_list[0]
            if not signature:
                raise ValueError("Unsigned Message")
            if signature in self.digest_history:
                raise ValueError("Duplicate Signature: %r"%signature)
            self.digest_history.add(signature)
            check = self.sign(msg_list[1:4])
            if not signature == check:
                raise ValueError("Invalid Signature: %r"%signature)
        if not len(msg_list) >= minlen:
            raise TypeError("malformed message, must have at least %i elements"%minlen)
        header = self.unpack(msg_list[1])
        message['header'] = header
        message['msg_id'] = header['msg_id']
        message['msg_type'] = header['msg_type']
        message['parent_header'] = self.unpack(msg_list[2])
        if content:
            message['content'] = self.unpack(msg_list[3])
        else:
            message['content'] = msg_list[3]

        message['buffers'] = msg_list[4:]
        return message

def test_msg2obj():
    am = dict(x=1)
    ao = Message(am)
    assert ao.x == am['x']

    am['y'] = dict(z=1)
    ao = Message(am)
    assert ao.y.z == am['y']['z']

    k1, k2 = 'y', 'z'
    assert ao[k1][k2] == am[k1][k2]

    am2 = dict(ao)
    assert am['x'] == am2['x']
    assert am['y']['z'] == am2['y']['z']


""" Defines utility functions for working with SVG documents in Qt.
"""

# System library imports.
from IPython.external.qt import QtCore, QtGui, QtSvg

def save_svg(string, parent=None):
    """ Prompts the user to save an SVG document to disk.

    Parameters:
    -----------
    string : basestring
        A Python string containing a SVG document.

    parent : QWidget, optional
        The parent to use for the file dialog.

    Returns:
    --------
    The name of the file to which the document was saved, or None if the save
    was cancelled.
    """
    if isinstance(string, unicode):
        string = string.encode('utf-8')

    dialog = QtGui.QFileDialog(parent, 'Save SVG Document')
    dialog.setAcceptMode(QtGui.QFileDialog.AcceptSave)
    dialog.setDefaultSuffix('svg')
    dialog.setNameFilter('SVG document (*.svg)')
    if dialog.exec_():
        filename = dialog.selectedFiles()[0]
        f = open(filename, 'w')
        try:
            f.write(string)
        finally:
            f.close()
        return filename
    return None

def svg_to_clipboard(string):
    """ Copy a SVG document to the clipboard.

    Parameters:
    -----------
    string : basestring
        A Python string containing a SVG document.
    """
    if isinstance(string, unicode):
        string = string.encode('utf-8')

    mime_data = QtCore.QMimeData()
    mime_data.setData('image/svg+xml', string)
    QtGui.QApplication.clipboard().setMimeData(mime_data)

def svg_to_image(string, size=None):
    """ Convert a SVG document to a QImage.

    Parameters:
    -----------
    string : basestring
        A Python string containing a SVG document.

    size : QSize, optional
        The size of the image that is produced. If not specified, the SVG
        document's default size is used.

    Raises:
    -------
    ValueError
        If an invalid SVG string is provided.

    Returns:
    --------
    A QImage of format QImage.Format_ARGB32.
    """
    if isinstance(string, unicode):
        string = string.encode('utf-8')

    renderer = QtSvg.QSvgRenderer(QtCore.QByteArray(string))
    if not renderer.isValid():
        raise ValueError('Invalid SVG data.')

    if size is None:
        size = renderer.defaultSize()
    image = QtGui.QImage(size, QtGui.QImage.Format_ARGB32)
    painter = QtGui.QPainter(image)
    renderer.render(painter)
    return image