# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

import sys
import traceback as tb

from marimo import _loggers
from marimo._ast.cell import CellImpl
from marimo._ast.toplevel import TopLevelExtraction
from marimo._data._external_storage.get_storage import (
    get_storage_backends_from_variables,
    storage_backend_to_storage_namespace,
)
from marimo._data.get_datasets import (
    get_datasets_from_variables,
    has_updates_to_datasource,
)
from marimo._dependencies.dependencies import DependencyManager
from marimo._messaging.cell_output import CellChannel
from marimo._messaging.errors import (
    MarimoExceptionRaisedError,
    MarimoInterruptionError,
    MarimoSQLError,
    MarimoStrictExecutionError,
)
from marimo._messaging.notification import (
    DatasetsNotification,
    DataSourceConnectionsNotification,
    StorageNamespacesNotification,
    VariableValuesNotification,
)
from marimo._messaging.notification_utils import (
    CellNotificationUtils,
    broadcast_notification,
)
from marimo._messaging.tracebacks import (
    _highlight_traceback,
    format_exception_message,
    write_traceback,
)
from marimo._messaging.variables import create_variable_value
from marimo._output import formatting
from marimo._plugins.ui._core.ui_element import UIElement
from marimo._runtime.context.types import (
    get_context,
    get_global_context,
)
from marimo._runtime.control_flow import MarimoInterrupt, MarimoStopError
from marimo._runtime.runner import cell_runner
from marimo._runtime.runner.hook_context import PostExecutionHookContext
from marimo._runtime.runner.hooks import PostExecutionHook
from marimo._runtime.side_effect import SideEffect
from marimo._sql.engines.duckdb import (
    INTERNAL_DUCKDB_ENGINE,
    DuckDBEngine,
)
from marimo._sql.get_engines import (
    engine_to_data_source_connection,
    get_engines_from_variables,
)
from marimo._tracer import kernel_tracer
from marimo._types.ids import VariableName
from marimo._utils.flatten import contains_instance

LOGGER = _loggers.marimo_logger()


@kernel_tracer.start_as_current_span("set_imported_defs")
def _set_imported_defs(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    del run_result
    LOGGER.debug("Acquiring graph lock to update cell import workspace")
    with ctx.graph.lock:
        LOGGER.debug("Acquired graph lock to update import workspace.")
        if cell.import_workspace.is_import_block:
            cell.import_workspace.imported_defs = {
                name for name in cell.defs if name in ctx.glbls
            }


@kernel_tracer.start_as_current_span("set_status_idle")
def _set_status_idle(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    del run_result
    del ctx
    cell.set_runtime_state(status="idle")


@kernel_tracer.start_as_current_span("set_run_result_status")
def _set_run_result_status(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    if isinstance(run_result.exception, MarimoInterrupt):
        # `MarimoInterruptionError` is a broadcast payload (never raised);
        # the exception held here is the raised `MarimoInterrupt`.
        cell.set_run_result_status("interrupted")
    elif cell.cell_id in ctx.cancelled_cells:
        cell.set_run_result_status("cancelled")
    elif run_result.exception is not None:
        cell.set_run_result_status(
            "exception",
            (
                # TODO(akshayka): "run_result.exception" can unfortunately
                # hold things that are not exceptions; remove this check
                # if/when that is ever cleaned up.
                run_result.exception
                if isinstance(run_result.exception, Exception)
                else None
            ),
        )
    else:
        cell.set_run_result_status("success")


@kernel_tracer.start_as_current_span("broadcast_variables")
def _broadcast_variables(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    if not ctx.should_broadcast_data:
        return

    del run_result
    values = [
        create_variable_value(
            name=variable,
            value=(ctx.glbls.get(variable, None)),
        )
        for variable in cell.defs
    ]
    if values:
        broadcast_notification(VariableValuesNotification(variables=values))


@kernel_tracer.start_as_current_span("broadcast_datasets")
def _broadcast_datasets(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    if not ctx.should_broadcast_data:
        return

    del run_result
    tables = get_datasets_from_variables(
        [
            (VariableName(variable), ctx.glbls[variable])
            for variable in cell.defs
            if variable in ctx.glbls
        ]
    )
    if tables:
        LOGGER.debug("Broadcasting data tables")
        broadcast_notification(DatasetsNotification(tables=tables))


@kernel_tracer.start_as_current_span("broadcast_data_source_connection")
def _broadcast_data_source_connection(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    if not ctx.should_broadcast_data:
        return

    del run_result
    engines = get_engines_from_variables(
        [
            (VariableName(variable), ctx.glbls[variable])
            for variable in cell.defs
            if variable in ctx.glbls
        ]
    )

    if not engines:
        return

    LOGGER.debug("Broadcasting data source connections")
    broadcast_notification(
        DataSourceConnectionsNotification(
            connections=[
                engine_to_data_source_connection(variable, engine)
                for variable, engine in engines
            ]
        )
    )


@kernel_tracer.start_as_current_span("broadcast_duckdb_datasource")
def _broadcast_duckdb_datasource(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    if not ctx.should_broadcast_data:
        return

    del run_result
    if not DependencyManager.duckdb.has():
        return

    try:
        sqls = cell.sqls
        if not sqls:
            return
        modifies_datasources = any(
            has_updates_to_datasource(sql) for sql in sqls
        )
        if not modifies_datasources:
            return

        LOGGER.debug("Broadcasting internal duckdb datasource")
        broadcast_notification(
            DataSourceConnectionsNotification(
                connections=[
                    engine_to_data_source_connection(
                        INTERNAL_DUCKDB_ENGINE, DuckDBEngine()
                    )
                ]
            )
        )
    except Exception:
        return


@kernel_tracer.start_as_current_span("broadcast_storage_backends")
def broadcast_storage_backends(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    del run_result

    if not ctx.should_broadcast_data:
        return

    try:
        storage_backends = get_storage_backends_from_variables(
            [
                (VariableName(variable), ctx.glbls[variable])
                for variable in cell.defs
                if variable in ctx.glbls
            ]
        )
        if storage_backends:
            LOGGER.debug("Broadcasting storage namespaces")
            namespaces = [
                storage_backend_to_storage_namespace(storage_backend)
                for _, storage_backend in storage_backends
            ]
            broadcast_notification(
                StorageNamespacesNotification(namespaces=namespaces)
            )
    except Exception:
        LOGGER.debug("Error getting storage backends", exc_info=True)


@kernel_tracer.start_as_current_span("store_reference_to_output")
def _store_reference_to_output(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    del ctx

    # Stores a reference to the output if it contains a UIElement
    # (required for RPCs on unnamed UI elements) or if it has a
    # _repr_mimebundle_ method (required to keep descriptor-based
    # anywidgets alive — their comm is closed on GC).
    # The _repr_mimebundle_ check is intentionally broad; the cost
    # is just one extra reference that's cleared on re-run.
    if isinstance(run_result.output, UIElement):
        cell.set_output(run_result.output)
    elif run_result.output is not None:
        if contains_instance(run_result.output, UIElement) or callable(
            getattr(run_result.output, "_repr_mimebundle_", None)
        ):
            cell.set_output(run_result.output)


def _store_state_reference(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    del run_result
    # Associate state variables with variable names
    runtime_ctx = get_context()
    runtime_ctx.state_registry.register_scope(ctx.glbls, defs=cell.defs)
    # Use pre-computed all_temporaries
    runtime_ctx.state_registry.retain_active_states(
        set(ctx.graph.definitions.keys()) | ctx.all_temporaries
    )


@kernel_tracer.start_as_current_span("issue_exception_side_effect")
def _issue_exception_side_effect(
    _cell: CellImpl,
    _ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    runtime_ctx = get_context()
    if run_result.exception is not None:
        exception = run_result.exception
        key = type(exception).__name__
        traceback = getattr(exception, "__traceback__", None)
        if traceback:
            # Side effect hash is the exception name
            # AND the instruction pointer.
            # Side effect has to be relatively robust to code changes
            #  - Bars line number since comments should not effect
            #  - Bars stacktrace since file path should be agnostic
            # Code content is already utilized in hash, and tb_lasti is the
            # bytecode instruction pointer. So if the code is the same, then the
            # difference can be captured by where in the evaluation the exception
            # was raised.
            key += f":{traceback.tb_lasti}"
        # NB. This is on a cell level.
        runtime_ctx.cell_lifecycle_registry.add(SideEffect(key))


@kernel_tracer.start_as_current_span("broadcast_outputs")
def _broadcast_outputs(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    # TODO: clean this logic up ...
    #
    # Send the output to the frontend
    #
    # Don't rebroadcast an output that was already sent
    #
    # 1. if run_result.output is not None, need to send it
    # 2. otherwise if accumulated_output is empty, then need to send
    #    the (empty) output (to clear it)
    should_send_output = (
        run_result.output is not None or not run_result.accumulated_output
    )
    if (
        run_result.success()
        or isinstance(run_result.exception, MarimoStopError)
    ) and should_send_output:
        formatted_output = formatting.try_format(run_result.output)
        if formatted_output.exception is not None:
            # Try a plain formatter; maybe an opinionated one failed.
            formatted_output = formatting.try_format(
                run_result.output, include_opinionated=False
            )
        # For ImportError and ModuleNotFoundError, store the exception in the
        # context so it can be reported by the missing_packages_hook.
        if isinstance(
            formatted_output.exception, (ImportError, ModuleNotFoundError)
        ):
            ctx.exceptions[cell.cell_id] = formatted_output.exception
        if formatted_output.traceback is not None:
            write_traceback(formatted_output.traceback)

        CellNotificationUtils.broadcast_output(
            channel=CellChannel.OUTPUT,
            mimetype=formatted_output.mimetype,
            data=formatted_output.data,
            cell_id=cell.cell_id,
            status=None,
        )
    elif isinstance(run_result.exception, MarimoStrictExecutionError):
        LOGGER.debug("Cell %s raised a strict error", cell.cell_id)
        # Cell never runs, so clear console
        CellNotificationUtils.broadcast_error(
            data=[run_result.output],
            clear_console=True,
            cell_id=cell.cell_id,
        )
    elif isinstance(run_result.exception, MarimoInterrupt):
        LOGGER.debug("Cell %s was interrupted", cell.cell_id)
        # don't clear console because this cell was running and
        # its console outputs are not stale
        CellNotificationUtils.broadcast_error(
            data=[MarimoInterruptionError()],
            clear_console=False,
            cell_id=cell.cell_id,
        )
    elif isinstance(run_result.exception, MarimoExceptionRaisedError):
        CellNotificationUtils.broadcast_error(
            data=[run_result.exception],
            clear_console=False,
            cell_id=cell.cell_id,
        )
    elif isinstance(run_result.exception, MarimoSQLError):
        LOGGER.debug("Cell %s raised a SQL error", cell.cell_id)
        CellNotificationUtils.broadcast_error(
            data=[run_result.exception],
            clear_console=True,
            cell_id=cell.cell_id,
        )
    elif run_result.exception is not None:
        LOGGER.debug(
            "Cell %s raised %s",
            cell.cell_id,
            type(run_result.exception).__name__,
        )
        # don't clear console because this cell was running and
        # its console outputs are not stale
        exception_type = type(run_result.exception).__name__
        if isinstance(run_result.exception, BaseException):
            msg = format_exception_message(run_result.exception)
        else:
            msg = str(run_result.exception)
        if not msg:
            msg = f"This cell raised an exception: {exception_type}"

        # Include formatted traceback if enabled in config
        formatted_traceback = None
        show_tracebacks = False
        if ctx.user_config is not None:
            show_tracebacks = bool(
                ctx.user_config["runtime"].get("show_tracebacks", False)
            )

        if show_tracebacks and (
            isinstance(run_result.exception, BaseException)
            and run_result.exception.__traceback__
        ):
            tb_lines = tb.format_exception(run_result.exception)
            formatted_traceback = _highlight_traceback("".join(tb_lines))

        CellNotificationUtils.broadcast_error(
            data=[
                MarimoExceptionRaisedError(
                    msg=msg,
                    exception_type=exception_type,
                    raising_cell=None,
                    traceback=formatted_traceback,
                )
            ],
            clear_console=False,
            cell_id=cell.cell_id,
        )


@kernel_tracer.start_as_current_span("render_toplevel_defs")
def render_toplevel_defs(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    del run_result
    variable = cell.toplevel_variable
    served = ctx.graph.cells_serving_serialization_hint
    if variable is not None:
        extractor = TopLevelExtraction.from_graph(ctx.graph, cell=cell)
        serialization = list(iter(extractor))[-1]
        CellNotificationUtils.broadcast_serialization(
            serialization=serialization,
            cell_id=cell.cell_id,
        )
        served.add(cell.cell_id)
    elif cell.cell_id in served:
        # Cell stopped being a top-level definition: clear the prior hint.
        # Only broadcast on this transition so ordinary cells don't emit an
        # extra cell-op on every run.
        CellNotificationUtils.broadcast_serialization_cleared(
            cell_id=cell.cell_id,
        )
        served.discard(cell.cell_id)


@kernel_tracer.start_as_current_span("run_pytest")
def attempt_pytest(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    del run_result
    if cell._test:
        try:
            import marimo._runtime.pytest as marimo_pytest

            if ctx.execution_context is not None:
                with ctx.execution_context(cell.cell_id):
                    result = marimo_pytest.run_pytest(cell.defs, ctx.glbls)
                    if result.output:
                        sys.stdout.write(result.output)
        except ImportError:
            pass


@kernel_tracer.start_as_current_span("reset_matplotlib_context")
def _reset_matplotlib_context(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    del cell
    del run_result
    if get_global_context().mpl_installed:
        # ensures that every cell gets a fresh axis.
        exec("__marimo__._output.mpl.close_figures()", ctx.glbls)


@kernel_tracer.start_as_current_span("delete_local_variables")
def _delete_local_variables(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    # Remove temporary (local) variables from the kernel globals to relieve
    # memory pressure once the cell has finished running.
    #
    # To prevent breaking existing notebooks, we make an exception for
    # temporaries that are closed over by a function, lambda, or class defined
    # in this cell: deleting those would cause a NameError when the closure is
    # later invoked, since closures do late-binding. These are precomputed at
    # compile time as `closed_over_temporaries`.
    del run_result
    for name in cell.temporaries - cell.closed_over_temporaries:
        ctx.glbls.pop(name, None)


@kernel_tracer.start_as_current_span("flush_console")
def _flush_console(
    cell: CellImpl,
    ctx: PostExecutionHookContext,
    run_result: cell_runner.RunResult,
) -> None:
    """Flush buffered console output before marking the cell idle.

    Console messages (stdout/stderr) are batched by a background thread
    for performance.  Without an explicit flush, the messages may arrive
    at the frontend *after* the cell is marked idle and `completed-run`
    is sent.  A subsequent run would then clear the console (via
    `console=[]`) before the user sees the output.
    """
    del cell
    del run_result
    del ctx
    stream = get_context().stream
    stream.flush_console()


POST_EXECUTION_HOOKS: list[PostExecutionHook] = [
    _set_imported_defs,
    _set_run_result_status,
    _store_reference_to_output,
    _store_state_reference,
    _issue_exception_side_effect,
    _broadcast_variables,
    _broadcast_datasets,
    _broadcast_data_source_connection,
    _broadcast_duckdb_datasource,
    _broadcast_outputs,
    _reset_matplotlib_context,
    _delete_local_variables,
    # Flush buffered console output so that stderr/stdout arrives at the
    # frontend before the cell transitions to idle.
    _flush_console,
    # set status to idle after all post-processing is done, in case the
    # other hooks take a long time (broadcast outputs can take a long time
    # if a formatter is slow).
    _set_status_idle,
    # NB. Other hooks are added ad-hoc or manually due to priority.
    # Consider implementing priority sort to keep everything more centralized.
]
