import asyncio
import sys
from contextlib import contextmanager

from pyxll import schedule_call
from spyder_kernels.console.kernel import SpyderKernel
from traitlets import MetaHasTraits
from zmq.eventloop.ioloop import IOLoop


@contextmanager
def push_stdio(stdout, stderr):
    """Context manage to temporarily replace stdout/stderr."""
    orig_stdout = sys.stdout
    orig_stderr = sys.stderr
    try:
        sys.stdout = stdout
        sys.stderr = stderr
        yield
    finally:
        sys.stdout = orig_stdout
        sys.stderr = orig_stderr


class PyXLLSpyderKernelMeta(MetaHasTraits):
    def __new__(cls, name, bases, dct):
        # Update our class dict with everything from SpyderKernel.
        # This is needed because in 'register_comm_handlers' (called from
        # the SpyderKernel constructor) only looks for methods in the top
        # class dict and so the SpyderKernel methods are excluded since
        # they're now members of the base class.
        if SpyderKernel in bases:
            dct = SpyderKernel.__dict__ | dct

        return super().__new__(cls, name, bases, dct)


class PyXLLSpyderKernel(SpyderKernel, metaclass=PyXLLSpyderKernelMeta):
    """
    Spyder kernel for running code inside Excel.

    This kernel runs dispatched commands via pyxll.schedule_call
    to ensure that user code is run safely in an Excel macro context.

    Doing this in the kernal instead of when polling the ioloop in
    the kernelapp ensures that schedule_call isn't being called when
    it's not needed. Calling schedule_call runs an Excel macro, which
    can have unwanted side effects like causing Excel's focus to change
    so we want to avoid using it when there's nothing to do.
    """

    def __init__(self, *args, **kwargs):
        # sys.stdout/stderrs get modified when the kernel is started
        self.__orig_stdout = self.__ipy_stdout = sys.stdout
        self.__orig_stderr = self.__ipy_stderr = sys.stderr

        # Event set when a message is dispatched to wake up polling the queue
        self.__disp_evt = asyncio.Event()

        # Run the base class constructor
        super().__init__(*args, **kwargs)

    def start(self):
        """Start the kernel"""
        super().start()

        # Capture the ipy stdout/stderr and restore sys.stdout/stderr after starting the kernel
        self.__ipy_stdout = sys.stdout
        self.__ipy_stderr = sys.stderr
        sys.stdout = self.__orig_stdout
        sys.stderr = self.__orig_stderr

    def schedule_dispatch(self, dispatch, *args):
        """Schedule a message for dispatch"""
        # Add to the dispatch queue
        super().schedule_dispatch(dispatch, *args)

        # Signal dispatch_queue to process the queue
        self.io_loop.add_callback(self.__disp_evt.set)

    async def dispatch_queue(self):
        """Coroutine to preserve order of message handling

        Ensures that only one message is processing at a time,
        even when the handler is async
        """
        # Wrap dispatch_queue with schedule_call.
        #
        # This ensures all code run in the kernel happens inside an Excel macro,
        # meaning that the Excel COM API can be used safely and the code can call
        # into Excel.
        #
        # This is run by the kernalapp's IOLoop, which is also run on the main
        # thread rather than PyXLL's asyncio event loop, so there are no issues with
        # different threads being used here.
        loop = IOLoop.current()

        def sync_dispatch():
            # Use the IPython stdout/stderr while processing the queue
            with push_stdio(self.__ipy_stdout, self.__ipy_stderr):

                async def process_queue():
                    while not self.msg_queue.empty():
                        await self.process_one()

                # Process the queue on the IOLoop and stop the loop when it's done
                fut = asyncio.ensure_future(process_queue())
                loop.add_future(fut, lambda _: loop.stop())

                # Run IOLoop until it's done
                loop.start()

        while True:
            self.__disp_evt.clear()

            # Process eveything in the queue in an Excel macro call
            await schedule_call(sync_dispatch)

            # Wait for more messages to be queued
            await self.__disp_evt.wait()
