"""
This script is executed on application startup if OVITO Pro is launched as
an IPython kernel (as specified in the kernel.json file).

It also backs the GUI "coding agent" feature, which runs an embedded IPython
kernel inside OVITO's main GUI thread. Unlike a normal kernel, the embedded
kernel is NOT driven by a blocking IOLoop: OVITO's Qt event loop steps the
kernel incrementally instead (see _EmbeddedKernelSession.step() and the C++
KernelSessionController). That way the main-thread call stack always unwinds
back to the Qt event loop between kernel messages, honoring OVITO's invariant
that shutdown/teardown only run once control has returned to loopLevel 0 -- no
nested QEventLoop is ever entered on the kernel's behalf.
"""

import os
import sys
from pathlib import Path


# The kernel session (if any) launched from the GUI. None when no kernel is running.
_active_session = None


# Cached IPKernelApp subclass used for the embedded kernel (built lazily so the
# heavy ipykernel import only happens when the feature is actually used).
_embedded_app_class = None


def _get_embedded_app_class():
    """
    Return an IPKernelApp subclass that disables kernel subshells.

    ipykernel 7 always creates a ShellChannelThread in init_control(); with it
    the shell_stream binds to that thread's IOLoop and shell messages are routed
    to subshells via inproc sockets. That topology cannot be driven cleanly from
    the Qt main loop with a single socket notifier. Dropping the
    shell_channel_thread makes ipykernel take the no-subshell path in
    KernelBase.start(): the shell_stream is bound directly to the main IOLoop and
    cells execute on the main thread, so watching the single shell socket is
    enough. The control channel keeps running on its own control_thread, so
    shutdown_request/interrupt_request are still handled independently.
    """
    global _embedded_app_class
    if _embedded_app_class is None:
        from ipykernel.kernelapp import IPKernelApp

        class _EmbeddedKernelApp(IPKernelApp):
            def init_control(self, context):
                super().init_control(context)
                # Drop the (just-constructed, never-started) shell channel
                # thread so init_kernel() binds the shell stream to the main
                # IOLoop instead of this thread's loop.
                self.shell_channel_thread = None

            def _init_asyncio_patch(self):
                # Deliberately a no-op, overriding IPKernelApp._init_asyncio_patch().
                # On Windows the base version, seeing the default Proactor policy,
                # calls asyncio.set_event_loop_policy(WindowsSelectorEventLoopPolicy())
                # at the very start of initialize(). That installs a *fresh* policy
                # whose per-thread current-loop slot is empty, discarding the
                # set_event_loop(self._loop) that _EmbeddedKernelSession just did --
                # so the shell ZMQStream built later in init_kernel() (with no explicit
                # loop) captures IOLoop.current() -> a brand-new orphan loop instead of
                # the loop we step, and the kernel never answers (see the detailed note
                # in _EmbeddedKernelSession.__init__).
                #
                # We don't need the base patch's effect anyway: the session installs a
                # SelectorEventLoop explicitly. Neutralizing the base method here -- rather
                # than pre-empting it by switching the policy ourselves -- also keeps this
                # module free of asyncio.get_event_loop_policy()/set_event_loop_policy(),
                # which are deprecated as of Python 3.14 and slated for removal in 3.16.
                pass

        _embedded_app_class = _EmbeddedKernelApp
    return _embedded_app_class


class _EmbeddedKernelSession:
    """
    One run of the GUI-embedded IPython kernel: from creation to client
    shutdown (whether the client disconnects, asks to restart, or code in the
    kernel calls exit()). OVITO is the kernel's host process, so there is no
    supervising process that could relaunch it on the same ports the way a
    subprocess-managed kernel would be restarted -- both cases just end this
    session, after which a new button click is needed to start a fresh one.

    The kernel's IOLoop is never started (io_loop.start() would block). Instead
    start() wires up the kernel and hands back the shell socket file descriptor,
    and step() runs the loop one non-blocking iteration at a time, called from
    OVITO's Qt event loop whenever that socket becomes readable (plus a
    low-frequency fallback timer for any scheduled asyncio callbacks).
    """

    def __init__(self) -> None:
        import asyncio

        # True once start() has wired up the kernel; cleared in _teardown().
        self._started = False
        # Set from _on_exit_now(): every path that ends the session (client
        # shutdown_request via the control thread, or in-kernel exit()/quit())
        # trips shell.exit_now, so this reliably signals step_kernel() to stop.
        self._should_stop = False

        # Give this session its own asyncio event loop and make it the current
        # loop *before* initialize()/start() run, so every IOLoop.current()
        # lookup (IPKernelApp, KernelBase.start(), the ZMQStreams in
        # init_kernel()) binds to a pristine loop. The first session leaves the
        # process's previous main-thread loop stopped-but-not-reset; reusing it
        # leaves the second kernel's shell messages undispatched (sockets bound,
        # but "kernel_info_request" never answered). Kept alive for the session's
        # lifetime and closed in _teardown(), which also clears the thread's
        # current loop so the next session starts clean.
        #
        # Force a SelectorEventLoop. On Windows the default asyncio loop is the
        # IOCP-based ProactorEventLoop, which is fundamentally incompatible with
        # the incremental-stepping design here for two reasons: (1) it does not
        # implement add_reader(), which ipykernel/pyzmq/tornado use to watch the
        # shell socket FD; and (2) its run_forever() always re-arms an internal
        # self-pipe overlapped read, so a run_forever() with a pre-scheduled
        # stop() never leaves loop._ready empty -- step()'s "drain until nothing
        # is ready" loop would spin forever, and once the queue does empty a
        # _run_once() computes an infinite select timeout and blocks the GUI
        # thread inside GetQueuedCompletionStatus. The SelectorEventLoop (the
        # default on macOS/Linux) provides both the add_reader() support and the
        # non-blocking select(0) semantics the stepping relies on. This mirrors
        # what Jupyter/tornado themselves force on Windows via the selector loop.
        #
        # Creating the loop explicitly here is what pins it to a selector loop, but
        # it is not sufficient on Windows on its own: IPKernelApp.initialize()
        # (called below) runs _init_asyncio_patch() first thing, and the stock
        # version -- seeing the default WindowsProactorEventLoopPolicy -- swaps in a
        # fresh event-loop policy whose current-loop slot is empty, discarding the
        # set_event_loop(self._loop) we do here. init_kernel() would then build the
        # shell ZMQStream (with no explicit loop) against IOLoop.current() -> a
        # brand-new orphan loop instead of self._loop, so the loop we step never
        # sees the shell socket and every shell message (starting with
        # kernel_info_request) goes unanswered -- the client times out. We prevent
        # that by overriding _init_asyncio_patch() to a no-op on _EmbeddedKernelApp
        # (see _get_embedded_app_class), so nothing discards our loop and the shell
        # stream binds to the loop we step. That override, rather than switching the
        # policy ourselves, is also what keeps this module off the deprecated
        # asyncio event-loop-policy API (removed in Python 3.16). macOS/Linux never
        # hit the swap in the first place -- _init_asyncio_patch() is a no-op there.
        self._loop = asyncio.SelectorEventLoop()
        asyncio.set_event_loop(self._loop)

        saved_argv = sys.argv
        sys.argv = ["ipykernel_launcher"]
        try:
            self._app = _get_embedded_app_class().instance()
            # Do not let ipykernel dup2() the process's stdout/stderr file
            # descriptors onto capture pipes. An embedded kernel shares OVITO's
            # terminal, so capture is unnecessary; worse, each initialize()
            # redirects fd 1/2 but teardown restores only the Python-level
            # sys.stdout/stderr, not the fds -- so a second session would
            # re-capture on top of a broken fd state and corrupt the new
            # kernel's IO wiring.
            self._app.capture_fd_output = False
            # Assign a per-session unique connection file name before initialize() runs.
            # ipykernel's default naming (kernel-<pid>.json) is constant across all
            # sessions in the same OVITO process.  A stale file with that deterministic
            # name -- from a previous session or a previous OVITO run with the same PID
            # -- causes init_connection_file() to load it via filefind() and then causes
            # write_connection_file() to skip the write (seeing identical content). When
            # that stale file is subsequently deleted by the previous session's cleanup,
            # this session has no connection file on disk.  A unique name sidesteps all
            # of this: filefind() never finds it, write_connection_file() always writes
            # it, and cleanup_connection_file() deletes the right file.
            import uuid
            self._app.connection_file = f"kernel-ovito-{uuid.uuid4().hex[:8]}.json"
            self._app.initialize()
        finally:
            sys.argv = saved_argv

        # Belt-and-suspenders: ensure the connection file actually made it to disk.
        if not Path(self._app.abs_connection_file).exists():
            self._app.write_connection_file()

        # exit_now can become True through several paths: a control-channel
        # shutdown_request handled on the control thread, or code running in the
        # kernel calling exit()/quit() (which sets it via IPython's own Quitter,
        # with no control message involved). Observe it directly so every such
        # path flips our stop flag and step_kernel() ends the session.
        self._app.kernel.shell.observe(self._on_exit_now, names="exit_now")

        cf = self._app.connection_file
        if not os.path.isabs(cf):
            cf = os.path.join(self._app.connection_dir, cf)
        self.connection_file = cf

    def _on_exit_now(self, change) -> None:
        if change["new"]:
            self._should_stop = True

    def is_running(self) -> bool:
        # True once the kernel is wired up and not yet torn down.
        return self._app is not None and self._started

    def should_stop(self) -> bool:
        return self._should_stop

    def start(self) -> int:
        """
        Wire up the kernel without entering a blocking IOLoop, and return the
        shell socket's file descriptor for OVITO to watch with a QSocketNotifier.

        Mirrors IPKernelApp.start() up to -- but excluding -- io_loop.start().
        """
        import asyncio
        import zmq
        from tornado import ioloop

        # Make sure the session's own loop is current before kernel.start()
        # registers its stream dispatchers on IOLoop.current().
        asyncio.set_event_loop(self._loop)

        app = self._app
        if app.poller is not None:
            app.poller.start()
        app.kernel.start()
        app.io_loop = ioloop.IOLoop.current()
        self._started = True

        # With subshells disabled the shell stream is bound to this main loop;
        # its socket FD is readable whenever a shell message is pending.
        return app.kernel.shell_stream.socket.getsockopt(zmq.FD)

    def step(self) -> None:
        """
        Advance the kernel by draining all work that is ready right now, without
        blocking and without entering a Qt event loop.

        A single asyncio iteration only polls I/O once and runs the callbacks
        that were already queued; receiving a shell message merely schedules the
        coroutine that executes it. Because the ZMQ socket FD stops being
        readable as soon as the message is read off it, the QSocketNotifier
        would not fire again for those follow-up iterations -- so we keep
        iterating here until no callbacks remain ready, which lets a full
        execute_request (recv -> run cell -> reply) complete in one activation.
        Work that is only scheduled for the future (e.g. an async cell awaiting a
        timer) leaves the ready queue empty and is picked up by the next
        notifier wakeup or the fallback timer.
        """
        loop = self._loop
        if not self._started or loop.is_running():
            # Guard against re-entrancy: user code that spins a Qt event loop
            # (modal dialog, processEvents) could deliver the notifier signal
            # while we are already inside run_forever().
            return
        while True:
            # Scheduling stop first makes each run_forever() perform exactly one
            # non-blocking iteration (timeout 0 because the ready queue is
            # non-empty) and then return.
            loop.call_soon(loop.stop)
            loop.run_forever()
            # loop._ready holds callbacks that are ready to run immediately.
            # Keep going while there is more such work; stop once the loop would
            # have to wait. (CPython internal, stable; OVITO ships CPython.)
            if not loop._ready:
                break

    def discard(self) -> None:
        # Tear the session down if that has not happened yet. Safe to call on an
        # already-finished session (whose _app is None) -- it becomes a no-op.
        if self._app is not None:
            self._teardown()

    def _teardown(self) -> None:
        # Best-effort: close sockets/threads, drop the connection file, and
        # clear the traitlets singletons (IPKernelApp/Kernel/ZMQInteractiveShell
        # are all SingletonConfigurable) so the next session's .instance() builds
        # a genuinely new app/kernel/shell trio instead of reusing this dead one.
        app = self._app
        assert app is not None
        self._app = None
        self._started = False
        kernel = getattr(app, "kernel", None)
        shell = getattr(kernel, "shell", None)
        try:
            app.close()
        except Exception:
            pass
        try:
            app.cleanup_connection_file()
        except Exception:
            pass
        for singleton in (shell, kernel, app):
            if singleton is None:
                continue
            try:
                type(singleton).clear_instance()
            except Exception:
                pass

        # Drop this session's dedicated asyncio loop and clear it from the
        # thread, so the next session installs a genuinely fresh loop rather
        # than reusing this now-stopped one.
        import asyncio

        try:
            asyncio.set_event_loop(None)
        except Exception:
            pass
        try:
            self._loop.close()
        except Exception:
            pass


def initialize_embedded_kernel() -> str:
    """
    Phase 1: Create the embedded kernel on the main thread (non-blocking), if
    none is currently active, and return the absolute path to its connection
    file. Call start_kernel_stepping() afterwards to begin processing messages.

    Calling this while a kernel session is already active is a no-op that
    returns the existing connection file.
    """
    global _active_session
    if _active_session is not None:
        if _active_session.is_running():
            # A kernel is already wired up and running: hand back its existing
            # connection file rather than tearing it down under OVITO's feet.
            return _active_session.connection_file
        # A previous session lingered (its stepping phase never started, or never
        # reset the global). Never silently hand back a stale kernel -- drop it
        # so every launch starts from a clean slate.
        _active_session.discard()
        _active_session = None
    _active_session = _EmbeddedKernelSession()
    return _active_session.connection_file


def start_kernel_stepping() -> int:
    """
    Phase 2: Wire up the active kernel session without blocking, and return the
    shell socket file descriptor that OVITO should watch. Returns -1 if there is
    no active session. The kernel is subsequently advanced by repeated
    step_kernel() calls driven from OVITO's Qt event loop.
    """
    session = _active_session
    if session is None:
        return -1
    return session.start()


def step_kernel() -> bool:
    """
    Advance the active kernel session by one batch of ready work. Returns True if
    the session should keep running, or False if it has ended (client
    disconnect/restart, or in-kernel exit()/quit()) -- in which case the caller
    must invoke shutdown_kernel_session() to tear it down.
    """
    session = _active_session
    if session is None:
        return False
    if session.should_stop():
        return False
    session.step()
    # An executed cell may have called exit()/quit(), or a shutdown_request may
    # have tripped exit_now on the control thread while we were stepping.
    return not session.should_stop()


def shutdown_kernel_session() -> None:
    """
    Tear down the active kernel session (if any), the same way a connected
    client's shutdown_request would. Used when step_kernel() reports the session
    has ended, and when OVITO itself is closed while a kernel is still running.
    Idempotent and safe to call more than once.
    """
    global _active_session
    session = _active_session
    _active_session = None
    if session is not None:
        session.discard()


# Start the kernel from an external jupyter or the --kernel command line option.
# This standalone path runs a dedicated OVITO-as-kernel process, where a blocking
# app.start() with a Qt event loop is exactly what we want; it is unrelated to
# the GUI-embedded, incrementally-stepped kernel above.
if __name__ == "__main__":
    # Remove the CWD from sys.path while we load stuff.
    # This is added back by InteractiveShellApp.init_path()
    if sys.path[0] == "" or Path(sys.path[0]) == Path.cwd():
        del sys.path[0]

    from ipykernel.kernelapp import IPKernelApp
    from ipykernel.eventloops import enable_gui

    app = IPKernelApp.instance()
    app.initialize()
    enable_gui("qt", app.kernel)
    app.start()
