"""
PyMOL IPython integration
"""

import pymol

START_SCRIPT = """
import pymol
from pymol import cmd
from IPython.core.magic import register_line_magic
for key in cmd.keyword:
    register_line_magic(key)(cmd.keyword[key][0])
del key
%clear
"""


def __init_plugin__(app=None):
    from pymol.plugins import addmenuitemqt
    addmenuitemqt('IPython Console', run_plugin_gui)


def run_plugin_gui():
    import os
    qt_api_env = os.getenv('QT_API', '')
    if qt_api_env:
        os.environ['QT_API'] = qt_api_env.lower()

    import warnings
    from pymol.Qt import QtCore, QtWidgets
    from qtconsole import inprocess
    from qtconsole import rich_jupyter_widget
    from ipykernel.inprocess import ipkernel

    class JupyterWidget(rich_jupyter_widget.RichJupyterWidget):
        def sizeHint(self):
            return QtCore.QSize(300, 400)

    class DockIPyWidget(QtWidgets.QDockWidget):
        def __init__(self, parent=None):
            super(DockIPyWidget, self).__init__(parent)

            warnings.filterwarnings("ignore",
                    category=DeprecationWarning,
                    module=r".*(jupyter)|(traitlets)")

            with warnings.catch_warnings(record=True):
                kernel_manager = KernelManagerWithUserModule()
                kernel_manager.start_kernel()
                kernel_manager.kernel.gui = 'qt'

                kernel_client = kernel_manager.client()
                kernel_client.start_channels()

                widget = JupyterWidget()
                widget.kernel_manager = kernel_manager
                widget.kernel_client = kernel_client
                widget.execute(START_SCRIPT)

            self.setWidget(widget)

    class KernelManagerWithUserModule(inprocess.QtInProcessKernelManager):
        def __init__(self, module=None):
            super(KernelManagerWithUserModule, self).__init__()

            if module is None:
                import types
                module = types.ModuleType("ipython_kernel_namespace")

            self._user_module = module

        def start_kernel(self, **kwargs):
            self.kernel = ipkernel.InProcessKernel(
                parent=self,
                session=self.session,
                user_module=self._user_module)

    window = pymol.gui.get_qtwindow()
    dockwidget = DockIPyWidget(window)
    window.addDockWidget(QtCore.Qt.LeftDockWidgetArea, dockwidget)
    dockwidget.show()