"""IPython kernel for parallel computing"""
import asyncio
import inspect
import sys
from functools import partial

import ipykernel
from ipykernel.ipkernel import IPythonKernel
from traitlets import Integer
from traitlets import Type

from .datapub import ZMQDataPublisher
from ipyparallel.serialize import serialize_object
from ipyparallel.serialize import unpack_apply_message
from ipyparallel.util import utcnow


class IPythonParallelKernel(IPythonKernel):
    """Extend IPython kernel for parallel computing"""

    engine_id = Integer(-1)

    @property
    def int_id(self):
        return self.engine_id

    msg_types = getattr(IPythonKernel, 'msg_types', []) + ['apply_request']
    control_msg_types = getattr(IPythonKernel, 'control_msg_types', []) + [
        'abort_request',
        'clear_request',
    ]
    _execute_sleep = 0
    data_pub_class = Type(ZMQDataPublisher)

    def _topic(self, topic):
        """prefixed topic for IOPub messages"""
        base = "engine.%s" % self.engine_id

        return f"{base}.{topic}".encode()

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # add apply_request, in anticipation of upstream deprecation
        self.shell_handlers['apply_request'] = self.apply_request
        # set up data pub
        data_pub = self.shell.data_pub = self.data_pub_class(parent=self)
        self.shell.configurables.append(data_pub)
        data_pub.session = self.session
        data_pub.pub_socket = self.iopub_socket
        self.aborted = set()

    def _abort_queues(self):
        # forward-port ipython/ipykernel#853
        # may remove after requiring ipykernel 6.9.1

        # while this flag is true,
        # execute requests will be aborted
        self._aborting = True
        self.log.info("Aborting queue")

        # Callback to signal that we are done aborting
        def stop_aborting():
            self.log.info("Finishing abort")
            self._aborting = False
            # must be awaitable for ipykernel >= 3.6
            # must also be sync for ipykernel < 3.6
            f = asyncio.Future()
            f.set_result(None)
            return f

        # put stop_aborting on the message queue
        # so that it's handled after processing of already-pending messages
        if ipykernel.version_info < (6,):
            # 10 is SHELL priority in ipykernel 5.x
            streams = self.shell_streams
            schedule_stop_aborting = partial(self.schedule_dispatch, 10, stop_aborting)
        else:
            streams = [self.shell_stream]
            schedule_stop_aborting = partial(self.schedule_dispatch, stop_aborting)

        # flush streams, so all currently waiting messages
        # are added to the queue
        for stream in streams:
            stream.flush()

        # if we have a delay, give messages this long to arrive on the queue
        # before we start accepting requests
        asyncio.get_event_loop().call_later(
            self.stop_on_error_timeout, schedule_stop_aborting
        )

        # for compatibility, return a completed Future
        # so this is still awaitable
        f = asyncio.Future()
        f.set_result(None)
        return f

    def should_handle(self, stream, msg, idents):
        """Check whether a shell-channel message should be handled

        Allows subclasses to prevent handling of certain messages (e.g. aborted requests).
        """
        if not super().should_handle(stream, msg, idents):
            return False
        msg_id = msg['header']['msg_id']
        msg_type = msg['header']['msg_type']
        if msg_id in self.aborted:
            # is it safe to assume a msg_id will not be resubmitted?
            self.aborted.remove(msg_id)
            self._send_abort_reply(stream, msg, idents)
            return False
        self.log.info(f"Handling {msg_type}: {msg_id}")
        return True

    def init_metadata(self, parent):
        """init metadata dict, for execute/apply_reply"""
        parent_metadata = parent.get('metadata', {})
        return {
            'started': utcnow(),
            'dependencies_met': True,
            'engine': self.ident,
            'is_broadcast': parent_metadata.get('is_broadcast', False),
            'is_coalescing': parent_metadata.get('is_coalescing', False),
            'original_msg_id': parent_metadata.get('original_msg_id', ''),
        }

    def finish_metadata(self, parent, metadata, reply_content):
        """Finish populating metadata.

        Run after completing a request handler.
        """
        metadata['status'] = reply_content['status']
        if reply_content['status'] == 'error':
            if reply_content['ename'] == 'UnmetDependency':
                metadata['dependencies_met'] = False
            metadata['engine_info'] = self.get_engine_info()

        return metadata

    def get_engine_info(self, method=None):
        """Return engine_info dict"""
        engine_info = dict(
            engine_uuid=self.ident,
            engine_id=self.engine_id,
        )
        if method:
            engine_info["method"] = method
        return engine_info

    def apply_request(self, stream, ident, parent):
        try:
            content = parent['content']
            bufs = parent['buffers']
            msg_id = parent['header']['msg_id']
        except:
            self.log.error("Got bad msg: %s", parent, exc_info=True)
            return

        md = self.init_metadata(parent)
        reply_content, result_buf = self.do_apply(content, bufs, msg_id, md)

        # put 'ok'/'error' status in header, for scheduler introspection:
        md = self.finish_metadata(parent, md, reply_content)

        # flush i/o
        sys.stdout.flush()
        sys.stderr.flush()
        self.log.debug(f'Sending apply_reply: {msg_id}')
        self.session.send(
            stream,
            'apply_reply',
            reply_content,
            parent=parent,
            ident=ident,
            buffers=result_buf,
            metadata=md,
        )

    def do_apply(self, content, bufs, msg_id, reply_metadata):
        shell = self.shell
        try:
            working = shell.user_ns

            prefix = "_" + str(msg_id).replace("-", "") + "_"

            f, args, kwargs = unpack_apply_message(bufs, working, copy=False)

            fname = getattr(f, '__name__', 'f')

            fname = prefix + "f"
            argname = prefix + "args"
            kwargname = prefix + "kwargs"
            resultname = prefix + "result"

            ns = {fname: f, argname: args, kwargname: kwargs, resultname: None}
            # print ns
            working.update(ns)
            code = f"{resultname} = {fname}(*{argname},**{kwargname})"
            try:
                exec(code, shell.user_global_ns, shell.user_ns)
                result = working.get(resultname)
            finally:
                for key in ns:
                    working.pop(key)

            result_buf = serialize_object(
                result,
                buffer_threshold=self.session.buffer_threshold,
                item_threshold=self.session.item_threshold,
            )

        except BaseException as e:
            # invoke IPython traceback formatting
            # this sends the 'error' message
            shell.showtraceback()

            try:
                str_evalue = str(e)
            except Exception as str_error:
                str_evalue = f"Failed to cast exception to string: {str_error}"
            reply_content = {
                'traceback': [],
                'ename': str(type(e).__name__),
                'evalue': str_evalue,
            }
            # get formatted traceback, which ipykernel recorded
            if hasattr(shell, '_last_traceback'):
                # ipykernel 4.4
                reply_content['traceback'] = shell._last_traceback or []
            else:
                self.log.warning("Didn't find a traceback where I expected to")
            shell._last_traceback = None
            reply_content["engine_info"] = self.get_engine_info(method="apply")

            self.log.info(
                "Exception in apply request:\n%s", '\n'.join(reply_content['traceback'])
            )
            result_buf = []
            reply_content['status'] = 'error'
        else:
            reply_content = {'status': 'ok'}

        return reply_content, result_buf

    async def _do_execute_async(self, *args, **kwargs):
        super_execute = super().do_execute(*args, **kwargs)
        if inspect.isawaitable(super_execute):
            reply_content = await super_execute
        else:
            reply_content = super_execute
        # add engine info
        if reply_content['status'] == 'error':
            reply_content["engine_info"] = self.get_engine_info(method="execute")
        return reply_content

    def do_execute(self, *args, **kwargs):
        coro = self._do_execute_async(*args, **kwargs)
        if ipykernel.version_info < (6,):
            # ipykernel 5 uses gen.maybe_future which doesn't accept async def coroutines,
            # but it does accept asyncio.Futures
            return asyncio.ensure_future(coro)
        else:
            return coro

    # Control messages for msgspec extensions:

    def abort_request(self, stream, ident, parent):
        """abort a specific msg by id"""
        msg_ids = parent['content'].get('msg_ids', None)
        if isinstance(msg_ids, str):
            msg_ids = [msg_ids]
        if not msg_ids:
            f = self._abort_queues()
            if inspect.isawaitable(f):
                asyncio.ensure_future(f)
        for mid in msg_ids:
            self.aborted.add(str(mid))

        content = dict(status='ok')
        reply_msg = self.session.send(
            stream, 'abort_reply', content=content, parent=parent, ident=ident
        )
        self.log.debug("%s", reply_msg)

    def clear_request(self, stream, idents, parent):
        """Clear our namespace."""
        self.shell.reset(False)
        content = dict(status='ok')
        self.session.send(
            stream, 'clear_reply', ident=idents, parent=parent, content=content
        )

    def _send_abort_reply(self, stream, msg, idents):
        """Send a reply to an aborted request"""
        # FIXME: forward-port ipython/ipykernel#684
        self.log.info(
            f"Aborting {msg['header']['msg_id']}: {msg['header']['msg_type']}"
        )
        reply_type = msg["header"]["msg_type"].rsplit("_", 1)[0] + "_reply"
        status = {"status": "aborted"}
        md = self.init_metadata(msg)
        md = self.finish_metadata(msg, md, status)
        md.update(status)

        self.session.send(
            stream,
            reply_type,
            metadata=md,
            content=status,
            parent=msg,
            ident=idents,
        )
