"""A kernel manager for multiple kernels"""
# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.
import asyncio
import os
import socket
import typing as t
import uuid
from functools import wraps

import zmq
from traitlets import Any, Bool, Dict, DottedObjectName, Instance, Unicode, default, observe
from traitlets.config.configurable import LoggingConfigurable
from traitlets.utils.importstring import import_item

from .kernelspec import NATIVE_KERNEL_NAME, KernelSpecManager
from .manager import KernelManager
from .utils import ensure_async, run_sync


class DuplicateKernelError(Exception):
    pass


def kernel_method(f: t.Callable) -> t.Callable:
    """decorator for proxying MKM.method(kernel_id) to individual KMs by ID"""

    @wraps(f)
    def wrapped(
        self: t.Any, kernel_id: str, *args: t.Any, **kwargs: t.Any
    ) -> t.Union[t.Callable, t.Awaitable]:
        # get the kernel
        km = self.get_kernel(kernel_id)
        method = getattr(km, f.__name__)
        # call the kernel's method
        r = method(*args, **kwargs)
        # last thing, call anything defined in the actual class method
        # such as logging messages
        f(self, kernel_id, *args, **kwargs)
        # return the method result
        return r

    return wrapped


class MultiKernelManager(LoggingConfigurable):
    """A class for managing multiple kernels."""

    default_kernel_name = Unicode(
        NATIVE_KERNEL_NAME, help="The name of the default kernel to start"
    ).tag(config=True)

    kernel_spec_manager = Instance(KernelSpecManager, allow_none=True)

    kernel_manager_class = DottedObjectName(
        "jupyter_client.ioloop.IOLoopKernelManager",
        help="""The kernel manager class.  This is configurable to allow
        subclassing of the KernelManager for customized behavior.
        """,
    ).tag(config=True)

    @observe("kernel_manager_class")
    def _kernel_manager_class_changed(self, change):
        self.kernel_manager_factory = self._create_kernel_manager_factory()

    kernel_manager_factory = Any(help="this is kernel_manager_class after import")

    @default("kernel_manager_factory")
    def _kernel_manager_factory_default(self):
        return self._create_kernel_manager_factory()

    def _create_kernel_manager_factory(self) -> t.Callable:
        kernel_manager_ctor = import_item(self.kernel_manager_class)

        def create_kernel_manager(*args: t.Any, **kwargs: t.Any) -> KernelManager:
            if self.shared_context:
                if self.context.closed:
                    # recreate context if closed
                    self.context = self._context_default()
                kwargs.setdefault("context", self.context)
            km = kernel_manager_ctor(*args, **kwargs)
            return km

        return create_kernel_manager

    shared_context = Bool(
        True,
        help="Share a single zmq.Context to talk to all my kernels",
    ).tag(config=True)

    context = Instance("zmq.Context")

    _created_context = Bool(False)

    _pending_kernels = Dict()

    @property
    def _starting_kernels(self):
        """A shim for backwards compatibility."""
        return self._pending_kernels

    @default("context")  # type:ignore[misc]
    def _context_default(self) -> zmq.Context:
        self._created_context = True
        return zmq.Context()

    connection_dir = Unicode("")

    _kernels = Dict()

    def __del__(self):
        """Handle garbage collection.  Destroy context if applicable."""
        if self._created_context and self.context and not self.context.closed:
            if self.log:
                self.log.debug("Destroying zmq context for %s", self)
            self.context.destroy()
        try:
            super_del = super().__del__  # type:ignore[misc]
        except AttributeError:
            pass
        else:
            super_del()

    def list_kernel_ids(self) -> t.List[str]:
        """Return a list of the kernel ids of the active kernels."""
        # Create a copy so we can iterate over kernels in operations
        # that delete keys.
        return list(self._kernels.keys())

    def __len__(self) -> int:
        """Return the number of running kernels."""
        return len(self.list_kernel_ids())

    def __contains__(self, kernel_id: str) -> bool:
        return kernel_id in self._kernels

    def pre_start_kernel(
        self, kernel_name: t.Optional[str], kwargs: t.Any
    ) -> t.Tuple[KernelManager, str, str]:
        # kwargs should be mutable, passing it as a dict argument.
        kernel_id = kwargs.pop("kernel_id", self.new_kernel_id(**kwargs))
        if kernel_id in self:
            raise DuplicateKernelError("Kernel already exists: %s" % kernel_id)

        if kernel_name is None:
            kernel_name = self.default_kernel_name
        # kernel_manager_factory is the constructor for the KernelManager
        # subclass we are using. It can be configured as any Configurable,
        # including things like its transport and ip.
        constructor_kwargs = {}
        if self.kernel_spec_manager:
            constructor_kwargs["kernel_spec_manager"] = self.kernel_spec_manager
        km = self.kernel_manager_factory(
            connection_file=os.path.join(self.connection_dir, "kernel-%s.json" % kernel_id),
            parent=self,
            log=self.log,
            kernel_name=kernel_name,
            **constructor_kwargs,
        )
        return km, kernel_name, kernel_id

    async def _add_kernel_when_ready(
        self, kernel_id: str, km: KernelManager, kernel_awaitable: t.Awaitable
    ) -> None:
        try:
            await kernel_awaitable
            self._kernels[kernel_id] = km
            self._pending_kernels.pop(kernel_id, None)
        except Exception as e:
            self.log.exception(e)

    async def _remove_kernel_when_ready(
        self, kernel_id: str, kernel_awaitable: t.Awaitable
    ) -> None:
        try:
            await kernel_awaitable
            self.remove_kernel(kernel_id)
            self._pending_kernels.pop(kernel_id, None)
        except Exception as e:
            self.log.exception(e)

    def _using_pending_kernels(self):
        """Returns a boolean; a clearer method for determining if
        this multikernelmanager is using pending kernels or not
        """
        return getattr(self, 'use_pending_kernels', False)

    async def _async_start_kernel(
        self, *, kernel_name: t.Optional[str] = None, **kwargs: t.Any
    ) -> str:
        """Start a new kernel.

        The caller can pick a kernel_id by passing one in as a keyword arg,
        otherwise one will be generated using new_kernel_id().

        The kernel ID for the newly started kernel is returned.
        """
        km, kernel_name, kernel_id = self.pre_start_kernel(kernel_name, kwargs)
        if not isinstance(km, KernelManager):
            self.log.warning(
                "Kernel manager class ({km_class}) is not an instance of 'KernelManager'!".format(
                    km_class=self.kernel_manager_class.__class__
                )
            )
        kwargs['kernel_id'] = kernel_id  # Make kernel_id available to manager and provisioner

        starter = ensure_async(km.start_kernel(**kwargs))
        task = asyncio.create_task(self._add_kernel_when_ready(kernel_id, km, starter))
        self._pending_kernels[kernel_id] = task
        # Handling a Pending Kernel
        if self._using_pending_kernels():
            # If using pending kernels, do not block
            # on the kernel start.
            self._kernels[kernel_id] = km
        else:
            await task
            # raise an exception if one occurred during kernel startup.
            if km.ready.exception():
                raise km.ready.exception()  # type: ignore

        return kernel_id

    start_kernel = run_sync(_async_start_kernel)

    async def _async_shutdown_kernel(
        self,
        kernel_id: str,
        now: t.Optional[bool] = False,
        restart: t.Optional[bool] = False,
    ) -> None:
        """Shutdown a kernel by its kernel uuid.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel to shutdown.
        now : bool
            Should the kernel be shutdown forcibly using a signal.
        restart : bool
            Will the kernel be restarted?
        """
        self.log.info("Kernel shutdown: %s", kernel_id)
        # If the kernel is still starting, wait for it to be ready.
        if kernel_id in self._pending_kernels:
            task = self._pending_kernels[kernel_id]
            try:
                await task
                km = self.get_kernel(kernel_id)
                await t.cast(asyncio.Future, km.ready)
            except asyncio.CancelledError:
                pass
            except Exception:
                self.remove_kernel(kernel_id)
                return
        km = self.get_kernel(kernel_id)
        # If a pending kernel raised an exception, remove it.
        if not km.ready.cancelled() and km.ready.exception():
            self.remove_kernel(kernel_id)
            return
        stopper = ensure_async(km.shutdown_kernel(now, restart))
        fut = asyncio.ensure_future(self._remove_kernel_when_ready(kernel_id, stopper))
        self._pending_kernels[kernel_id] = fut
        # Await the kernel if not using pending kernels.
        if not self._using_pending_kernels():
            await fut
            # raise an exception if one occurred during kernel shutdown.
            if km.ready.exception():
                raise km.ready.exception()  # type: ignore

    shutdown_kernel = run_sync(_async_shutdown_kernel)

    @kernel_method
    def request_shutdown(self, kernel_id: str, restart: t.Optional[bool] = False) -> None:
        """Ask a kernel to shut down by its kernel uuid"""

    @kernel_method
    def finish_shutdown(
        self,
        kernel_id: str,
        waittime: t.Optional[float] = None,
        pollinterval: t.Optional[float] = 0.1,
    ) -> None:
        """Wait for a kernel to finish shutting down, and kill it if it doesn't"""
        self.log.info("Kernel shutdown: %s", kernel_id)

    @kernel_method
    def cleanup_resources(self, kernel_id: str, restart: bool = False) -> None:
        """Clean up a kernel's resources"""

    def remove_kernel(self, kernel_id: str) -> KernelManager:
        """remove a kernel from our mapping.

        Mainly so that a kernel can be removed if it is already dead,
        without having to call shutdown_kernel.

        The kernel object is returned, or `None` if not found.
        """
        return self._kernels.pop(kernel_id, None)

    async def _async_shutdown_all(self, now: bool = False) -> None:
        """Shutdown all kernels."""
        kids = self.list_kernel_ids()
        kids += list(self._pending_kernels)
        kms = list(self._kernels.values())
        futs = [self._async_shutdown_kernel(kid, now=now) for kid in set(kids)]
        await asyncio.gather(*futs)
        # If using pending kernels, the kernels will not have been fully shut down.
        if self._using_pending_kernels():
            for km in kms:
                try:
                    await km.ready
                except asyncio.CancelledError:
                    self._pending_kernels[km.kernel_id].cancel()
                except Exception:
                    # Will have been logged in _add_kernel_when_ready
                    pass

    shutdown_all = run_sync(_async_shutdown_all)

    def interrupt_kernel(self, kernel_id: str) -> None:
        """Interrupt (SIGINT) the kernel by its uuid.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel to interrupt.
        """
        kernel = self.get_kernel(kernel_id)
        if not kernel.ready.done():
            msg = "Kernel is in a pending state. Cannot interrupt."
            raise RuntimeError(msg)
        out = kernel.interrupt_kernel()
        self.log.info("Kernel interrupted: %s", kernel_id)
        return out

    @kernel_method
    def signal_kernel(self, kernel_id: str, signum: int) -> None:
        """Sends a signal to the kernel by its uuid.

        Note that since only SIGTERM is supported on Windows, this function
        is only useful on Unix systems.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel to signal.
        signum : int
            Signal number to send kernel.
        """
        self.log.info("Signaled Kernel %s with %s", kernel_id, signum)

    async def _async_restart_kernel(self, kernel_id: str, now: bool = False) -> None:
        """Restart a kernel by its uuid, keeping the same ports.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel to interrupt.
        now : bool, optional
            If True, the kernel is forcefully restarted *immediately*, without
            having a chance to do any cleanup action.  Otherwise the kernel is
            given 1s to clean up before a forceful restart is issued.

            In all cases the kernel is restarted, the only difference is whether
            it is given a chance to perform a clean shutdown or not.
        """
        kernel = self.get_kernel(kernel_id)
        if self._using_pending_kernels() and not kernel.ready.done():
            msg = "Kernel is in a pending state. Cannot restart."
            raise RuntimeError(msg)
        await ensure_async(kernel.restart_kernel(now=now))
        self.log.info("Kernel restarted: %s", kernel_id)

    restart_kernel = run_sync(_async_restart_kernel)

    @kernel_method
    def is_alive(self, kernel_id: str) -> bool:  # type:ignore[empty-body]
        """Is the kernel alive.

        This calls KernelManager.is_alive() which calls Popen.poll on the
        actual kernel subprocess.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel.
        """

    def _check_kernel_id(self, kernel_id: str) -> None:
        """check that a kernel id is valid"""
        if kernel_id not in self:
            raise KeyError("Kernel with id not found: %s" % kernel_id)

    def get_kernel(self, kernel_id: str) -> KernelManager:
        """Get the single KernelManager object for a kernel by its uuid.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel.
        """
        self._check_kernel_id(kernel_id)
        return self._kernels[kernel_id]

    @kernel_method
    def add_restart_callback(
        self, kernel_id: str, callback: t.Callable, event: str = "restart"
    ) -> None:
        """add a callback for the KernelRestarter"""

    @kernel_method
    def remove_restart_callback(
        self, kernel_id: str, callback: t.Callable, event: str = "restart"
    ) -> None:
        """remove a callback for the KernelRestarter"""

    @kernel_method
    def get_connection_info(self, kernel_id: str) -> t.Dict[str, t.Any]:  # type:ignore[empty-body]
        """Return a dictionary of connection data for a kernel.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel.

        Returns
        =======
        connection_dict : dict
            A dict of the information needed to connect to a kernel.
            This includes the ip address and the integer port
            numbers of the different channels (stdin_port, iopub_port,
            shell_port, hb_port).
        """

    @kernel_method
    def connect_iopub(  # type:ignore[empty-body]
        self, kernel_id: str, identity: t.Optional[bytes] = None
    ) -> socket.socket:
        """Return a zmq Socket connected to the iopub channel.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel
        identity : bytes (optional)
            The zmq identity of the socket

        Returns
        =======
        stream : zmq Socket or ZMQStream
        """

    @kernel_method
    def connect_shell(  # type:ignore[empty-body]
        self, kernel_id: str, identity: t.Optional[bytes] = None
    ) -> socket.socket:
        """Return a zmq Socket connected to the shell channel.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel
        identity : bytes (optional)
            The zmq identity of the socket

        Returns
        =======
        stream : zmq Socket or ZMQStream
        """

    @kernel_method
    def connect_control(  # type:ignore[empty-body]
        self, kernel_id: str, identity: t.Optional[bytes] = None
    ) -> socket.socket:
        """Return a zmq Socket connected to the control channel.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel
        identity : bytes (optional)
            The zmq identity of the socket

        Returns
        =======
        stream : zmq Socket or ZMQStream
        """

    @kernel_method
    def connect_stdin(  # type:ignore[empty-body]
        self, kernel_id: str, identity: t.Optional[bytes] = None
    ) -> socket.socket:
        """Return a zmq Socket connected to the stdin channel.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel
        identity : bytes (optional)
            The zmq identity of the socket

        Returns
        =======
        stream : zmq Socket or ZMQStream
        """

    @kernel_method
    def connect_hb(  # type:ignore[empty-body]
        self, kernel_id: str, identity: t.Optional[bytes] = None
    ) -> socket.socket:
        """Return a zmq Socket connected to the hb channel.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel
        identity : bytes (optional)
            The zmq identity of the socket

        Returns
        =======
        stream : zmq Socket or ZMQStream
        """

    def new_kernel_id(self, **kwargs: t.Any) -> str:
        """
        Returns the id to associate with the kernel for this request. Subclasses may override
        this method to substitute other sources of kernel ids.
        :param kwargs:
        :return: string-ized version 4 uuid
        """
        return str(uuid.uuid4())


class AsyncMultiKernelManager(MultiKernelManager):
    kernel_manager_class = DottedObjectName(
        "jupyter_client.ioloop.AsyncIOLoopKernelManager",
        config=True,
        help="""The kernel manager class.  This is configurable to allow
        subclassing of the AsyncKernelManager for customized behavior.
        """,
    )

    use_pending_kernels = Bool(
        False,
        help="""Whether to make kernels available before the process has started.  The
        kernel has a `.ready` future which can be awaited before connecting""",
    ).tag(config=True)

    context = Instance("zmq.asyncio.Context")

    @default("context")  # type:ignore[misc]
    def _context_default(self) -> zmq.asyncio.Context:
        self._created_context = True
        return zmq.asyncio.Context()

    start_kernel: t.Callable[
        ..., t.Awaitable
    ] = MultiKernelManager._async_start_kernel  # type:ignore[assignment]
    restart_kernel: t.Callable[
        ..., t.Awaitable
    ] = MultiKernelManager._async_restart_kernel  # type:ignore[assignment]
    shutdown_kernel: t.Callable[
        ..., t.Awaitable
    ] = MultiKernelManager._async_shutdown_kernel  # type:ignore[assignment]
    shutdown_all: t.Callable[
        ..., t.Awaitable
    ] = MultiKernelManager._async_shutdown_all  # type:ignore[assignment]
