"""
Based on
https://github.com/jupyter/notebook/blob/master/notebook/static/services/kernels/comm.js
https://github.com/ipython/ipykernel/blob/master/ipykernel/comm/manager.py
https://github.com/ipython/ipykernel/blob/master/ipykernel/comm/comm.py


Which are distributed under the terms of the Modified BSD License.
"""
import logging

from traitlets.config import LoggingConfigurable

from ipython_genutils.importstring import import_item
from ipython_genutils.py3compat import string_types

import uuid

from qtconsole.qt import QtCore
from qtconsole.util import MetaQObjectHasTraits, SuperQObject

class CommManager(MetaQObjectHasTraits(
        'NewBase', (LoggingConfigurable, SuperQObject), {})):
    """
    Manager for Comms in the Frontend
    """
    def __init__(self, kernel_client, *args, **kwargs):
        super(CommManager, self).__init__(*args, **kwargs)
        self.comms = {}
        self.targets = {}
        if kernel_client:
            self.init_kernel_client(kernel_client)

    def init_kernel_client(self, kernel_client):
        """
        connect the kernel, and register message handlers
        """
        self.kernel_client = kernel_client
        kernel_client.iopub_channel.message_received.connect(self._dispatch)

    @QtCore.Slot(object)
    def _dispatch(self, msg):
        """Dispatch messages"""
        msg_type = msg['header']['msg_type']
        handled_msg_types = ['comm_open', 'comm_msg', 'comm_close']
        if msg_type in handled_msg_types:
            getattr(self, msg_type)(msg)

    def new_comm(self, target_name, data=None, metadata=None,
                 comm_id=None, buffers=None):
        """
        Create a new Comm, register it, and open its Kernel-side counterpart
        Mimics the auto-registration in `Comm.__init__` in the Jupyter Comm.
        
        argument comm_id is optional
        """
        comm = Comm(target_name, self.kernel_client, comm_id)
        self.register_comm(comm)
        try:
            comm.open(data, metadata, buffers)
        except:
            self.unregister_comm(comm)
            raise
        return comm

    def register_target(self, target_name, f):
        """Register a callable f for a given target name
        
        f will be called with two arguments when a comm_open message is
        received with `target`:
            
        - the Comm instance
        - the `comm_open` message itself.
        
        f can be a Python callable or an import string for one.
        """
        if isinstance(f, string_types):
            f = import_item(f)

        self.targets[target_name] = f

    def unregister_target(self, target_name, f):
        """Unregister a callable registered with register_target"""
        return self.targets.pop(target_name)

    def register_comm(self, comm):
        """Register a new comm"""
        comm_id = comm.comm_id
        comm.kernel_client = self.kernel_client
        self.comms[comm_id] = comm
        comm.sig_is_closing.connect(self.unregister_comm)
        return comm_id

    @QtCore.Slot(object)
    def unregister_comm(self, comm):
        """Unregister a comm, and close its counterpart."""
        # unlike get_comm, this should raise a KeyError
        comm.sig_is_closing.disconnect(self.unregister_comm)
        self.comms.pop(comm.comm_id)

    def get_comm(self, comm_id):
        """Get a comm with a particular id
        
        Returns the comm if found, otherwise None.
        
        This will not raise an error,
        it will log messages if the comm cannot be found.
        """
        try:
            return self.comms[comm_id]
        except KeyError:
            self.log.warning("No such comm: %s", comm_id)
            if self.log.isEnabledFor(logging.DEBUG):
                # don't create the list of keys if debug messages aren't enabled
                self.log.debug("Current comms: %s", list(self.comms.keys()))

    # comm message handlers
    def comm_open(self, msg):
        """Handler for comm_open messages"""
        content = msg['content']
        comm_id = content['comm_id']
        target_name = content['target_name']
        f = self.targets.get(target_name, None)

        comm = Comm(target_name, self.kernel_client, comm_id)
        self.register_comm(comm)

        if f is None:
            self.log.error("No such comm target registered: %s", target_name)
        else:
            try:
                f(comm, msg)
                return
            except Exception:
                self.log.error("Exception opening comm with target: %s",
                               target_name, exc_info=True)

        # Failure.
        try:
            comm.close()
        except:
            self.log.error("""Could not close comm during `comm_open` failure
                clean-up.  The comm may not have been opened yet.""",
                exc_info=True)

    def comm_close(self, msg):
        """Handler for comm_close messages"""
        content = msg['content']
        comm_id = content['comm_id']
        comm = self.get_comm(comm_id)
        if comm is None:
            return

        self.unregister_comm(comm)

        try:
            comm.handle_close(msg)
        except Exception:
            self.log.error('Exception in comm_close for %s', comm_id,
                           exc_info=True)

    def comm_msg(self, msg):
        """Handler for comm_msg messages"""
        content = msg['content']
        comm_id = content['comm_id']
        comm = self.get_comm(comm_id)
        if comm is None:
            return
        try:
            comm.handle_msg(msg)
        except Exception:
            self.log.error('Exception in comm_msg for %s', comm_id,
                           exc_info=True)

class Comm(MetaQObjectHasTraits(
        'NewBase', (LoggingConfigurable, SuperQObject), {})):
    """
    Comm base class
    """
    sig_is_closing = QtCore.Signal(object)

    def __init__(self, target_name, kernel_client, comm_id=None,
                 msg_callback=None, close_callback=None):
        """
        Create a new comm. Must call open to use.
        """
        super(Comm, self).__init__(target_name=target_name)
        self.target_name = target_name
        self.kernel_client = kernel_client
        if comm_id is None:
            comm_id = uuid.uuid1().hex
        self.comm_id = comm_id
        self._msg_callback = msg_callback
        self._close_callback = close_callback
        self._send_channel = self.kernel_client.shell_channel

    def __del__(self):
        self.close()

    def _send_msg(self, msg_type, content, data, metadata, buffers):
        """
        Send a message on the shell channel.
        """
        if data is None:
            data = {}
        if content is None:
            content = {}
        content['comm_id'] = self.comm_id
        content['data'] = data

        msg = self.kernel_client.session.msg(
            msg_type, content, metadata=metadata)
        if buffers:
            msg['buffers'] = buffers
        return self._send_channel.send(msg)

    # methods for sending messages
    def open(self, data=None, metadata=None, buffers=None):
        """Open the kernel-side version of this comm"""
        return self._send_msg(
            'comm_open', {'target_name': self.target_name},
            data, metadata, buffers)

    def send(self, data=None, metadata=None, buffers=None):
        """Send a message to the kernel-side version of this comm"""
        return self._send_msg(
            'comm_msg', {}, data, metadata, buffers)

    def close(self, data=None, metadata=None, buffers=None):
        """Close the kernel-side version of this comm"""
        self.sig_is_closing.emit(self)
        return self._send_msg(
            'comm_close', {}, data, metadata, buffers)

    # methods for registering callbacks for incoming messages

    def on_msg(self, callback):
        """Register a callback for comm_msg
        
        Will be called with the `data` of any comm_msg messages.
        
        Call `on_msg(None)` to disable an existing callback.
        """
        self._msg_callback = callback

    def on_close(self, callback):
        """Register a callback for comm_close
        
        Will be called with the `data` of the close message.
        
        Call `on_close(None)` to disable an existing callback.
        """
        self._close_callback = callback
    
    # methods for handling incoming messages
    def handle_msg(self, msg):
        """Handle a comm_msg message"""
        self.log.debug("handle_msg[%s](%s)", self.comm_id, msg)
        if self._msg_callback:
            return self._msg_callback(msg)

    def handle_close(self, msg):
        """Handle a comm_close message"""
        self.log.debug("handle_close[%s](%s)", self.comm_id, msg)
        if self._close_callback:
            return self._close_callback(msg)

__all__ = ['CommManager']
