from __future__ import annotations

import abc
import datetime
import typing as t
import unittest
import uuid

from hyperscript import h
from rich.console import Console as RichConsole
from rich.progress import BarColumn, Progress, TaskID, TextColumn, TimeElapsedColumn
from rich.prompt import Confirm, Prompt
from rich.status import Status
from rich.syntax import Syntax
from rich.tree import Tree

from sqlmesh.core.snapshot import Snapshot, SnapshotChangeCategory
from sqlmesh.core.test import ModelTest
from sqlmesh.utils import rich as srich
from sqlmesh.utils.date import to_date

if t.TYPE_CHECKING:
    import ipywidgets as widgets

    from sqlmesh.core.context_diff import ContextDiff
    from sqlmesh.core.plan import Plan

    LayoutWidget = t.TypeVar("LayoutWidget", bound=t.Union[widgets.VBox, widgets.HBox])


SNAPSHOT_CHANGE_CATEGORY_STR = {
    SnapshotChangeCategory.BREAKING: "Breaking",
    SnapshotChangeCategory.NON_BREAKING: "Non-breaking",
    SnapshotChangeCategory.FORWARD_ONLY: "Forward-only",
}


class Console(abc.ABC):
    """Abstract base class for defining classes used for displaying information to the user and also interact
    with them when their input is needed"""

    @abc.abstractmethod
    def start_snapshot_progress(self, snapshot_name: str, total_batches: int) -> None:
        """Indicates that a new load progress has begun."""

    @abc.abstractmethod
    def update_snapshot_progress(self, snapshot_name: str, num_batches: int) -> None:
        """Update snapshot progress."""

    @abc.abstractmethod
    def stop_snapshot_progress(self, success: bool = True) -> None:
        """Stop the load progress"""

    @abc.abstractmethod
    def start_promotion_progress(self, environment: str, total_tasks: int) -> None:
        """Indicates that a new promotion progress has begun."""

    @abc.abstractmethod
    def update_promotion_progress(self, num_tasks: int) -> None:
        """Update promotion progress."""

    @abc.abstractmethod
    def stop_promotion_progress(self, success: bool = True) -> None:
        """Stop the promotion progress"""

    @abc.abstractmethod
    def show_model_difference_summary(
        self, context_diff: ContextDiff, detailed: bool = False
    ) -> None:
        """Displays a summary of differences for the given models"""

    @abc.abstractmethod
    def plan(self, plan: Plan, auto_apply: bool) -> None:
        """The main plan flow.

        The console should present the user with choices on how to backfill and version the snapshots
        of a plan.

        Args:
            plan: The plan to make choices for.
            auto_apply: Whether to automatically apply the plan after all choices have been made.
        """

    @abc.abstractmethod
    def log_test_results(
        self, result: unittest.result.TestResult, output: str, target_dialect: str
    ) -> None:
        """Display the test result and output

        Args:
            result: The unittest test result that contains metrics like num success, fails, ect.
            output: The generated output from the unittest
            target_dialect: The dialect that tests were run against. Assumes all tests run against the same dialect.
        """

    @abc.abstractmethod
    def show_sql(self, sql: str) -> None:
        """Display to the user SQL"""

    @abc.abstractmethod
    def log_status_update(self, message: str) -> None:
        """Display general status update to the user"""

    @abc.abstractmethod
    def log_error(self, message: str) -> None:
        """Display error info to the user"""

    @abc.abstractmethod
    def log_success(self, message: str) -> None:
        """Display a general successful message to the user"""

    @abc.abstractmethod
    def loading_start(self, message: t.Optional[str] = None) -> uuid.UUID:
        """Starts loading and returns a unique ID that can be used to stop the loading. Optionally can display a message"""

    @abc.abstractmethod
    def loading_stop(self, id: uuid.UUID) -> None:
        """Stop loading for the given id"""


class TerminalConsole(Console):
    """A rich based implementation of the console"""

    def __init__(self, console: t.Optional[RichConsole] = None, **kwargs: t.Any) -> None:
        self.console: RichConsole = console or srich.console
        self.evaluation_progress: t.Optional[Progress] = None
        self.evaluation_tasks: t.Dict[str, t.Tuple[TaskID, int]] = {}
        self.promotion_progress: t.Optional[Progress] = None
        self.promotion_task: t.Optional[TaskID] = None
        self.loading_status: t.Dict[uuid.UUID, Status] = {}

    def _print(self, value: t.Any, **kwargs: t.Any) -> None:
        self.console.print(value)

    def _prompt(self, message: str, **kwargs: t.Any) -> t.Any:
        return Prompt.ask(message, console=self.console, **kwargs)

    def _confirm(self, message: str, **kwargs: t.Any) -> bool:
        return Confirm.ask(message, console=self.console, **kwargs)

    def start_snapshot_progress(self, snapshot_name: str, total_batches: int) -> None:
        """Indicates that a new load progress has begun."""
        if not self.evaluation_progress:
            self.evaluation_progress = Progress(
                TextColumn("[bold blue]{task.fields[snapshot_name]}", justify="right"),
                BarColumn(bar_width=40),
                "[progress.percentage]{task.percentage:>3.1f}%",
                "•",
                srich.SchedulerBatchColumn(),
                "•",
                TimeElapsedColumn(),
                console=self.console,
            )
            self.evaluation_progress.start()
            self.evaluation_tasks = {}
        self.evaluation_tasks[snapshot_name] = (
            self.evaluation_progress.add_task(
                f"Running {snapshot_name}...",
                snapshot_name=snapshot_name,
                total=total_batches,
            ),
            total_batches,
        )

    def update_snapshot_progress(self, snapshot_name: str, num_batches: int) -> None:
        """Update snapshot progress."""
        if self.evaluation_progress and self.evaluation_tasks:
            task_id = self.evaluation_tasks[snapshot_name][0]
            self.evaluation_progress.update(task_id, refresh=True, advance=num_batches)

    def stop_snapshot_progress(self, success: bool = True) -> None:
        """Stop the load progress"""
        self.evaluation_tasks = {}
        if self.evaluation_progress:
            self.evaluation_progress.stop()
            self.evaluation_progress = None
            if success:
                self.log_success("All model batches have been executed successfully")

    def start_promotion_progress(self, environment: str, total_tasks: int) -> None:
        """Indicates that a new promotion progress has begun."""
        if self.promotion_progress is None:
            self.promotion_progress = Progress(
                TextColumn(f"[bold blue]Virtually Updating '{environment}'", justify="right"),
                BarColumn(bar_width=40),
                "[progress.percentage]{task.percentage:>3.1f}%",
                "•",
                TimeElapsedColumn(),
                console=self.console,
            )
            self.promotion_progress.start()
            self.promotion_task = self.promotion_progress.add_task(
                f"Virtually Updating {environment}...",
                total=total_tasks,
            )

    def update_promotion_progress(self, num_tasks: int) -> None:
        """Update promotion progress."""
        if self.promotion_progress is not None and self.promotion_task is not None:
            self.promotion_progress.update(self.promotion_task, refresh=True, advance=num_tasks)

    def stop_promotion_progress(self, success: bool = True) -> None:
        """Stop the promotion progress"""
        self.promotion_task = None
        if self.promotion_progress is not None:
            self.promotion_progress.stop()
            self.promotion_progress = None
            if success:
                self.log_success("The target environment has been updated successfully")

    def show_model_difference_summary(
        self, context_diff: ContextDiff, detailed: bool = False
    ) -> None:
        """Shows a summary of the differences.

        Args:
            context_diff: The context diff to use to print the summary
            detailed: Show the actual SQL differences if True.
        """
        if context_diff.is_new_environment:
            self._print(
                Tree(
                    f"[bold]New environment `{context_diff.environment}` will be created from `{context_diff.create_from}`"
                )
            )
            if not context_diff.has_snapshot_changes:
                return

        if not context_diff.has_changes:
            self._print(Tree(f"[bold]No differences when compared to `{context_diff.environment}`"))
            return

        tree = Tree(f"[bold]Summary of differences against `{context_diff.environment}`:")

        if context_diff.added:
            added_tree = Tree(f"[bold][added]Added Models:")
            for model in context_diff.added:
                added_tree.add(f"[added]{model}")
            tree.add(added_tree)

        if context_diff.removed:
            removed_tree = Tree(f"[bold][removed]Removed Models:")
            for model in context_diff.removed:
                removed_tree.add(f"[removed]{model}")
            tree.add(removed_tree)

        if context_diff.modified_snapshots:
            direct = Tree(f"[bold][direct]Directly Modified:")
            indirect = Tree(f"[bold][indirect]Indirectly Modified:")
            metadata = Tree(f"[bold][metadata]Metadata Updated:")
            for model in context_diff.modified_snapshots:
                if context_diff.directly_modified(model):
                    direct.add(
                        Syntax(f"{model}\n{context_diff.text_diff(model)}", "sql")
                        if detailed
                        else f"[direct]{model}"
                    )
                elif context_diff.indirectly_modified(model):
                    indirect.add(f"[indirect]{model}")
                elif context_diff.metadata_updated(model):
                    metadata.add(f"[metadata]{model}")
            if direct.children:
                tree.add(direct)
            if indirect.children:
                tree.add(indirect)
            if metadata.children:
                tree.add(metadata)
        self._print(tree)

    def plan(self, plan: Plan, auto_apply: bool) -> None:
        """The main plan flow.

        The console should present the user with choices on how to backfill and version the snapshots
        of a plan.

        Args:
            plan: The plan to make choices for.
            auto_apply: Whether to automatically apply the plan after all choices have been made.
        """
        self._prompt_categorize(plan, auto_apply)
        self._show_options_after_categorization(plan, auto_apply)

        if auto_apply:
            plan.apply()

    def _show_options_after_categorization(self, plan: Plan, auto_apply: bool) -> None:
        if plan.requires_backfill:
            self._show_missing_dates(plan)
            self._prompt_backfill(plan, auto_apply)
        elif plan.context_diff.has_changes and not auto_apply:
            self._prompt_promote(plan)

    def _prompt_categorize(self, plan: Plan, auto_apply: bool) -> None:
        """Get the user's change category for the directly modified models"""
        self.show_model_difference_summary(plan.context_diff)

        self._show_categorized_snapshots(plan)

        for snapshot in plan.uncategorized:
            self._print(Syntax(plan.context_diff.text_diff(snapshot.name), "sql"))
            tree = Tree(f"[bold][direct]Directly Modified: {snapshot.name}")
            indirect_tree = None

            for child in plan.indirectly_modified[snapshot.name]:
                if not indirect_tree:
                    indirect_tree = Tree(f"[indirect]Indirectly Modified Children:")
                    tree.add(indirect_tree)
                indirect_tree.add(f"[indirect]{child}")
            self._print(tree)
            self._get_snapshot_change_category(snapshot, plan, auto_apply)

    def _show_categorized_snapshots(self, plan: Plan) -> None:
        context_diff = plan.context_diff
        for snapshot in plan.categorized:
            if not context_diff.directly_modified(snapshot.name):
                continue

            category_str = SNAPSHOT_CHANGE_CATEGORY_STR[plan.snapshot_change_category(snapshot)]
            tree = Tree(f"[bold][direct]Directly Modified: {snapshot.name} ({category_str})")
            syntax_dff = Syntax(context_diff.text_diff(snapshot.name), "sql")
            indirect_tree = None
            for child in plan.indirectly_modified[snapshot.name]:
                if not indirect_tree:
                    indirect_tree = Tree(f"[indirect]Indirectly Modified Children:")
                    tree.add(indirect_tree)
                indirect_tree.add(f"[indirect]{child}")
            self._print(syntax_dff)
            self._print(tree)

    def _show_missing_dates(self, plan: Plan) -> None:
        """Displays the models with missing dates"""
        if not plan.missing_intervals:
            return
        backfill = Tree("[bold]Models needing backfill (missing dates):")
        for missing in plan.missing_intervals:
            backfill.add(f"{missing.snapshot_name}: {missing.format_missing_range()}")
        self._print(backfill)

    def _prompt_backfill(self, plan: Plan, auto_apply: bool) -> None:
        is_forward_only_dev = plan.is_dev and plan.forward_only
        backfill_or_preview = "preview" if is_forward_only_dev else "backfill"

        if plan.is_start_and_end_allowed:
            if not plan.override_start:
                blank_meaning = (
                    "to preview starting from yesterday"
                    if is_forward_only_dev
                    else "for the beginning of history"
                )
                start = self._prompt(
                    f"Enter the {backfill_or_preview} start date (eg. '1 year', '2020-01-01') or blank {blank_meaning}",
                )
                if start:
                    plan.start = start

            if not plan.override_end:
                end = self._prompt(
                    f"Enter the {backfill_or_preview} end date (eg. '1 month ago', '2020-01-01') or blank to {backfill_or_preview} up until now",
                )
                if end:
                    plan.end = end

        if not auto_apply and self._confirm(f"Apply - {backfill_or_preview.capitalize()} Tables"):
            plan.apply()

    def _prompt_promote(self, plan: Plan) -> None:
        if self._confirm(
            f"Apply - Virtual Update",
        ):
            plan.apply()

    def log_test_results(
        self, result: unittest.result.TestResult, output: str, target_dialect: str
    ) -> None:
        divider_length = 70
        if result.wasSuccessful():
            self._print("=" * divider_length)
            self._print(
                f"Successfully Ran {str(result.testsRun)} tests against {target_dialect}",
                style="green",
            )
            self._print("-" * divider_length)
        else:
            self._print("-" * divider_length)
            self._print("Test Failure Summary")
            self._print("=" * divider_length)
            self._print(
                f"Num Successful Tests: {result.testsRun - len(result.failures) - len(result.errors)}"
            )
            for test, _ in result.failures + result.errors:
                if isinstance(test, ModelTest):
                    self._print(f"Failure Test: {test.model_name} {test.test_name}")
            self._print("=" * divider_length)
            self._print(output)

    def show_sql(self, sql: str) -> None:
        self._print(Syntax(sql, "sql"))

    def log_status_update(self, message: str) -> None:
        self._print(message)

    def log_error(self, message: str) -> None:
        self._print(f"[red]{message}[/red]")

    def log_success(self, message: str) -> None:
        self._print(f"\n[green]{message}[/green]\n")

    def loading_start(self, message: t.Optional[str] = None) -> uuid.UUID:
        id = uuid.uuid4()
        self.loading_status[id] = Status(message or "", console=self.console, spinner="line")
        self.loading_status[id].start()
        return id

    def loading_stop(self, id: uuid.UUID) -> None:
        self.loading_status[id].stop()
        del self.loading_status[id]

    def _get_snapshot_change_category(
        self, snapshot: Snapshot, plan: Plan, auto_apply: bool
    ) -> None:
        choices = self._snapshot_change_choices(snapshot)
        response = self._prompt(
            "\n".join([f"[{i+1}] {choice}" for i, choice in enumerate(choices.values())]),
            show_choices=False,
            choices=[f"{i+1}" for i in range(len(choices))],
        )
        choice = list(choices)[int(response) - 1]
        plan.set_choice(snapshot, choice)

    def _snapshot_change_choices(
        self, snapshot: Snapshot, use_rich_formatting: bool = True
    ) -> t.Dict[SnapshotChangeCategory, str]:
        direct = snapshot.name
        if use_rich_formatting:
            direct = f"[direct]{direct}[/direct]"
        indirect = "indirectly modified children"
        if use_rich_formatting:
            indirect = f"[indirect]{indirect}[/indirect]"
        if snapshot.is_view_kind:
            choices = {
                SnapshotChangeCategory.BREAKING: f"Update {direct} and backfill {indirect}",
                SnapshotChangeCategory.NON_BREAKING: f"Update {direct} but don't backfill {indirect}",
            }
        elif snapshot.is_embedded_kind:
            choices = {
                SnapshotChangeCategory.BREAKING: f"Backfill {indirect}",
                SnapshotChangeCategory.NON_BREAKING: f"Don't backfill {indirect}",
            }
        else:
            choices = {
                SnapshotChangeCategory.BREAKING: f"Backfill {direct} and {indirect}",
                SnapshotChangeCategory.NON_BREAKING: f"Backfill {direct} but not {indirect}",
            }
        labeled_choices = {
            k: f"[{SNAPSHOT_CHANGE_CATEGORY_STR[k]}] {v}" for k, v in choices.items()
        }
        return labeled_choices


def add_to_layout_widget(target_widget: LayoutWidget, *widgets: widgets.Widget) -> LayoutWidget:
    """Helper function to add a widget to a layout widget
    Args:
        target_widget: The layout widget to add the other widget(s) to
        *widgets: The widgets to add to the layout widget

    Returns:
        The layout widget with the children added
    """
    target_widget.children += tuple(widgets)
    return target_widget


class NotebookMagicConsole(TerminalConsole):
    """
    Console to be used when using the magic notebook interface (`%<command>`).
    Generally reuses the Terminal console when possible by either directly outputing what it provides
    or capturing it and converting it into a widget.
    """

    def __init__(
        self,
        display: t.Optional[t.Callable] = None,
        console: t.Optional[RichConsole] = None,
        **kwargs: t.Any,
    ) -> None:
        import ipywidgets as widgets
        from IPython.display import display as ipython_display

        super().__init__(console, **kwargs)
        self.display = display or get_ipython().user_ns.get("display", ipython_display)  # type: ignore
        self.missing_dates_output = widgets.Output()
        self.dynamic_options_after_categorization_output = widgets.VBox()

    def _show_missing_dates(self, plan: Plan) -> None:
        self._add_to_dynamic_options(self.missing_dates_output)
        self.missing_dates_output.outputs = ()
        with self.missing_dates_output:
            super()._show_missing_dates(plan)

    def _apply(self, button: widgets.Button) -> None:
        button.disabled = True
        with button.output:
            button.plan.apply()

    def _prompt_promote(self, plan: Plan) -> None:
        import ipywidgets as widgets

        button = widgets.Button(
            description="Apply - Virtual Update",
            disabled=False,
            button_style="success",
            # Auto will make the button really large.
            # Likely changing this soon anyways to be just `Apply` with description above
            layout={"width": "10rem"},
        )
        self._add_to_dynamic_options(button)
        output = widgets.Output()
        self._add_to_dynamic_options(output)

        button.plan = plan
        button.on_click(self._apply)
        button.output = output

    def _prompt_backfill(self, plan: Plan, auto_apply: bool) -> None:
        import ipywidgets as widgets

        prompt = widgets.VBox()

        backfill_or_preview = "Preview" if plan.is_dev and plan.forward_only else "Backfill"

        def _date_picker(
            plan: Plan, value: t.Any, on_change: t.Callable, disabled: bool = False
        ) -> widgets.DatePicker:
            picker = widgets.DatePicker(
                disabled=disabled,
                value=value,
                layout={"width": "auto"},
            )

            picker.observe(on_change, "value")
            return picker

        def _checkbox(description: str, value: bool, on_change: t.Callable) -> widgets.Checkbox:
            checkbox = widgets.Checkbox(
                value=value,
                description=description,
                disabled=False,
                indent=False,
            )

            checkbox.observe(on_change, "value")
            return checkbox

        def start_change_callback(change: t.Dict[str, datetime.datetime]) -> None:
            plan.start = change["new"]
            self._show_options_after_categorization(plan, auto_apply)

        def end_change_callback(change: t.Dict[str, datetime.datetime]) -> None:
            plan.end = change["new"]
            self._show_options_after_categorization(plan, auto_apply)

        if plan.is_start_and_end_allowed:
            add_to_layout_widget(
                prompt,
                widgets.HBox(
                    [
                        widgets.Label(
                            f"Start {backfill_or_preview} Date:", layout={"width": "8rem"}
                        ),
                        _date_picker(plan, to_date(plan.start), start_change_callback),
                    ]
                ),
            )

            add_to_layout_widget(
                prompt,
                widgets.HBox(
                    [
                        widgets.Label(f"End {backfill_or_preview} Date:", layout={"width": "8rem"}),
                        _date_picker(
                            plan,
                            to_date(plan.end),
                            end_change_callback,
                        ),
                    ]
                ),
            )

        self._add_to_dynamic_options(prompt)

        if not auto_apply:
            button = widgets.Button(
                description=f"Apply - {backfill_or_preview} Tables",
                disabled=False,
                button_style="success",
            )
            self._add_to_dynamic_options(button)
            output = widgets.Output()
            self._add_to_dynamic_options(output)

            button.plan = plan
            button.on_click(self._apply)
            button.output = output

    def _show_options_after_categorization(self, plan: Plan, auto_apply: bool) -> None:
        self.dynamic_options_after_categorization_output.children = ()
        self.display(self.dynamic_options_after_categorization_output)
        super()._show_options_after_categorization(plan, auto_apply)

    def _add_to_dynamic_options(self, *widgets: widgets.Widget) -> None:
        add_to_layout_widget(self.dynamic_options_after_categorization_output, *widgets)

    def _get_snapshot_change_category(
        self, snapshot: Snapshot, plan: Plan, auto_apply: bool
    ) -> None:
        import ipywidgets as widgets

        def radio_button_selected(change: t.Dict[str, t.Any]) -> None:
            plan.set_choice(snapshot, choices[change["owner"].index])
            self._show_options_after_categorization(plan, auto_apply)

        choice_mapping = self._snapshot_change_choices(snapshot, use_rich_formatting=False)
        choices = list(choice_mapping)
        plan.set_choice(snapshot, choices[0])

        radio = widgets.RadioButtons(
            options=choice_mapping.values(),
            layout={"width": "max-content"},
            disabled=False,
        )
        radio.observe(
            radio_button_selected,
            "value",
        )
        self.display(radio)

    def log_test_results(
        self, result: unittest.result.TestResult, output: str, target_dialect: str
    ) -> None:
        import ipywidgets as widgets

        divider_length = 70
        shared_style = {
            "font-size": "11px",
            "font-weight": "bold",
            "font-family": "Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace",
        }
        if result.wasSuccessful():
            success_color = {"color": "#008000"}
            header = str(h("span", {"style": shared_style}, "-" * divider_length))
            message = str(
                h(
                    "span",
                    {"style": {**shared_style, **success_color}},
                    f"Successfully Ran {str(result.testsRun)} Tests Against {target_dialect}",
                )
            )
            footer = str(h("span", {"style": shared_style}, "=" * divider_length))
            self.display(widgets.HTML("<br>".join([header, message, footer])))
        else:
            fail_color = {"color": "#db3737"}
            fail_shared_style = {**shared_style, **fail_color}
            header = str(h("span", {"style": fail_shared_style}, "-" * divider_length))
            message = str(h("span", {"style": fail_shared_style}, "Test Failure Summary"))
            num_success = str(
                h(
                    "span",
                    {"style": fail_shared_style},
                    f"Num Successful Tests: {result.testsRun - len(result.failures) - len(result.errors)}",
                )
            )
            failure_tests = []
            for test, _ in result.failures + result.errors:
                if isinstance(test, ModelTest):
                    failure_tests.append(
                        str(
                            h(
                                "span",
                                {"style": fail_shared_style},
                                f"Failure Test: {test.model_name} {test.test_name}",
                            )
                        )
                    )
            failures = "<br>".join(failure_tests)
            footer = str(h("span", {"style": fail_shared_style}, "=" * divider_length))
            error_output = widgets.Textarea(output, layout={"height": "300px", "width": "100%"})
            test_info = widgets.HTML(
                "<br>".join([header, message, footer, num_success, failures, footer])
            )
            self.display(widgets.VBox(children=[test_info, error_output], layout={"width": "100%"}))


class DatabricksMagicConsole(TerminalConsole):
    """
    Note: Databricks Magic Console currently does not support progress bars while a plan is being applied. The
    NotebookMagicConsole does support progress bars, but they will time out after 5 minutes of execution
    and it makes it difficult to see the progress of the plan.
    """

    def _print(self, value: t.Any, **kwargs: t.Any) -> None:
        with self.console.capture() as capture:
            self.console.print(value, **kwargs)
        output = capture.get()
        print(output)

    def _prompt(self, message: str, **kwargs: t.Any) -> t.Any:
        self._print(message)
        return super()._prompt("", **kwargs)

    def _confirm(self, message: str, **kwargs: t.Any) -> bool:
        message = f"{message} \[y/n]"
        self._print(message)
        return super()._confirm("", **kwargs)


def get_console(**kwargs: t.Any) -> TerminalConsole | DatabricksMagicConsole | NotebookMagicConsole:
    """
    Returns the console that is appropriate for the current runtime environment.

    Note: Google Colab environment is untested and currently assumes is compatible with the base
    NotebookMagicConsole.
    """
    from sqlmesh import RuntimeEnv, runtime_env

    runtime_env_mapping = {
        RuntimeEnv.DATABRICKS: DatabricksMagicConsole,
        RuntimeEnv.JUPYTER: NotebookMagicConsole,
        RuntimeEnv.TERMINAL: TerminalConsole,
        RuntimeEnv.GOOGLE_COLAB: NotebookMagicConsole,
    }
    return runtime_env_mapping[runtime_env](**kwargs)
