# *************************************************************************
#
#  Copyright (c) 2026 - Datatailr Inc.
#  All Rights Reserved.
#
#  This file is part of Datatailr and subject to the terms and conditions
#  defined in 'LICENSE.txt'. Unauthorized copying and/or distribution
#  of this file, in parts or full, via any medium is strictly prohibited.
# *************************************************************************

"""
IPython magic for exploring workflow runs.

Usage inside a Jupyter notebook::

    %load_ext datatailr.scheduler.jupyter_magic

    # Open the widget with a workflow dropdown
    %dt_runs

    # Preselect a workflow
    %dt_runs simple-data-pipeline

    # Specify environment
    %dt_runs simple-data-pipeline --env prod
"""

from datatailr import Environment

_IPYTHON_AVAILABLE = True
try:
    from IPython.core.magic import Magics, magics_class, line_magic  # type: ignore[import-not-found]
except ImportError:
    _IPYTHON_AVAILABLE = False

_WIDGET_AVAILABLE = True
try:
    from datatailr.scheduler.jupyter_widget import ResultViewerWidget
except ImportError:
    _WIDGET_AVAILABLE = False


_ENV_MAP = {
    "dev": Environment.DEV,
    "pre": Environment.PRE,
    "prod": Environment.PROD,
}


if _IPYTHON_AVAILABLE:

    @magics_class
    class DatatailrMagics(Magics):
        @line_magic
        def dt_runs(self, line: str):
            """Open the workflow runs explorer widget.

            Usage::

                %dt_runs                              # dropdown, dev environment
                %dt_runs my-workflow                   # preselect workflow
                %dt_runs my-workflow --env prod        # preselect workflow in prod
                %dt_runs --env pre                     # dropdown, pre environment
            """
            if not _WIDGET_AVAILABLE:
                print(
                    "\033[33mThe %dt_runs magic requires datatailr.scheduler.jupyter_widget "
                    "and its dependencies (ipywidgets, networkx, ipycytoscape).\033[0m"
                )
                return

            parts = line.strip().split()
            env = Environment.DEV
            workflow_name = None

            i = 0
            while i < len(parts):
                if parts[i] == "--env" and i + 1 < len(parts):
                    env_str = parts[i + 1].lower()
                    env = _ENV_MAP.get(env_str, Environment.DEV)
                    i += 2
                else:
                    workflow_name = parts[i]
                    i += 1

            widget = ResultViewerWidget(environment=env)
            widget.display()
            if workflow_name:
                widget.select_workflow(workflow_name)


def load_ipython_extension(ipython):
    if not _IPYTHON_AVAILABLE:
        raise ImportError(
            "IPython is required for the %dt_runs magic. "
            "Install it with: pip install ipython"
        )
    ipython.register_magics(DatatailrMagics)
