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

import asyncio
import os
import sys
from dataclasses import dataclass, replace
from functools import cached_property
from typing import TYPE_CHECKING, Literal, cast

from marimo import _loggers
from marimo._ast.app import InternalApp
from marimo._ast.errors import CycleError, MultipleDefinitionError
from marimo._ast.load import load_app
from marimo._cli.print import echo
from marimo._config.config import RuntimeConfig
from marimo._config.manager import (
    get_default_config_manager,
)
from marimo._convert.common.filename import get_download_filename
from marimo._convert.converters import MarimoConvert
from marimo._messaging.cell_output import CellChannel, CellOutput
from marimo._messaging.errors import Error, is_unexpected_error
from marimo._messaging.notification import (
    CellNotification,
    CompletedRunNotification,
)
from marimo._messaging.serde import deserialize_kernel_message
from marimo._messaging.types import KernelMessage
from marimo._output.hypertext import patch_html_for_non_interactive_output
from marimo._runtime.commands import AppMetadata, SerializedCLIArgs
from marimo._runtime.patches import extract_docstring_from_header
from marimo._schemas.serialization import NotebookSerialization
from marimo._server.export._status import emit_pdf_export_status
from marimo._server.export.exporter import Exporter
from marimo._server.models.export import (
    ExportAsHTMLRequest,
    ExportPDFPreset,
)
from marimo._server.models.models import InstantiateNotebookRequest
from marimo._session.model import ConnectionState, SessionMode
from marimo._session.notebook import AppFileManager, load_notebook
from marimo._types.ids import ConsumerId
from marimo._utils.inline_script_metadata import (
    pin_pep723_dependencies_for_wasm,
)
from marimo._utils.marimo_path import MarimoPath

LOGGER = _loggers.marimo_logger()

if TYPE_CHECKING:
    from collections.abc import Mapping

    from marimo._server.export._pdf_raster import PDFRasterizationOptions
    from marimo._server.export._status import PDFExportStatusCallback
    from marimo._session.state.session_view import SessionView
    from marimo._session.types import Session
    from marimo._types.ids import CellId_t


@dataclass
class ExportResult:
    contents: bytes | str
    download_filename: str
    did_error: bool

    @cached_property
    def bytez(self) -> bytes:
        """Return UTF-8 encoded bytes (cached)."""
        if isinstance(self.contents, bytes):
            return self.contents
        return self.contents.encode("utf-8")

    @cached_property
    def text(self) -> str:
        """Return UTF-8 decoded text (cached)."""
        if isinstance(self.contents, str):
            return self.contents
        return self.contents.decode("utf-8")


def _as_ir(path: MarimoPath) -> NotebookSerialization:
    if path.is_python():
        py_contents = path.read_text(encoding="utf-8")
        converter = MarimoConvert.from_py(py_contents)
        return _with_filename(converter.ir, path.short_name)
    elif path.is_markdown():
        md_contents = path.read_text(encoding="utf-8")
        converter = MarimoConvert.from_md(md_contents)
        return _with_filename(converter.ir, path.short_name)
    raise ValueError(f"Unsupported file type: {path.path.suffix}")


def export_as_script(path: MarimoPath) -> ExportResult:
    from marimo._convert.script import convert_from_ir_to_script

    ir = _as_ir(path)
    return ExportResult(
        contents=convert_from_ir_to_script(ir),
        download_filename=get_download_filename(path.short_name, "script.py"),
        did_error=False,
    )


def export_as_md(path: MarimoPath) -> ExportResult:
    ir = _as_ir(path)
    return ExportResult(
        contents=MarimoConvert.from_ir(ir).to_markdown(),
        download_filename=get_download_filename(path.short_name, "md"),
        did_error=False,
    )


def export_as_ipynb(
    path: MarimoPath, sort_mode: Literal["top-down", "topological"]
) -> ExportResult:
    app = load_app(path.absolute_name)
    if app is None:
        return ExportResult(
            contents=b"",
            download_filename=get_download_filename(path.short_name, "ipynb"),
            did_error=True,
        )

    # Try the requested sort mode, fall back to top-down if cycles exist
    internal_app = InternalApp(app)
    actual_sort_mode = sort_mode
    if sort_mode == "topological":
        try:
            # Check if graph can be accessed (raises CycleError/MultipleDefinitionError)
            _ = internal_app.graph
        except (CycleError, MultipleDefinitionError):
            echo(
                "Warning: Notebook has errors, "
                "using top-down order instead of topological.",
                err=True,
            )
            actual_sort_mode = "top-down"

    result = Exporter().export_as_ipynb(
        app=internal_app,
        sort_mode=actual_sort_mode,
    )
    return ExportResult(
        contents=result,
        download_filename=get_download_filename(path.short_name, "ipynb"),
        did_error=False,
    )


def export_as_wasm(
    path: MarimoPath,
    mode: Literal["edit", "run"],
    show_code: bool,
    asset_url: str | None = None,
) -> ExportResult:
    _app = load_app(path.absolute_name)
    if _app is None:
        return ExportResult(
            contents=b"",
            download_filename=get_download_filename(
                path.short_name, "wasm.html"
            ),
            did_error=True,
        )
    app = InternalApp(_app)
    # Inline the layout file, if it exists
    app.inline_layout_file()
    config = get_default_config_manager(current_path=path.absolute_name)
    resolved = config.get_config()

    result = Exporter().export_as_wasm(
        filename=path.short_name,
        app=app,
        display_config=resolved["display"],
        mode=mode,
        code=app.to_py(),
        asset_url=asset_url,
        show_code=show_code,
        sharing_config=resolved.get("sharing"),
    )
    return ExportResult(
        contents=result[0],
        download_filename=result[1],
        did_error=False,
    )


def notebook_uses_slides_layout(filepath: MarimoPath) -> bool:
    """Return whether a notebook declares the slides layout."""
    try:
        file_manager = load_notebook(filepath.absolute_name)
        layout_config = file_manager.read_layout_config()
        return layout_config is not None and layout_config.type == "slides"
    except Exception as e:
        LOGGER.debug(
            "Unable to infer notebook layout for %s: %s",
            filepath.absolute_name,
            e,
        )
        return False


async def run_app_then_export_as_ipynb(
    filepath: MarimoPath,
    sort_mode: Literal["top-down", "topological"],
    cli_args: SerializedCLIArgs,
    argv: list[str] | None,
) -> ExportResult:
    file_manager = load_notebook(filepath.absolute_name)

    with patch_html_for_non_interactive_output():
        # Use quiet=True to suppress runtime stdout/stderr since outputs
        # are captured in the session_view and will be included in the ipynb
        (session_view, did_error) = await run_app_until_completion(
            file_manager,
            cli_args,
            argv,
            quiet=True,
        )

    result = Exporter().export_as_ipynb(
        app=file_manager.app,
        sort_mode=sort_mode,
        session_view=session_view,
    )
    return ExportResult(
        contents=result,
        download_filename=get_download_filename(filepath.short_name, "ipynb"),
        did_error=did_error,
    )


async def run_app_then_export_as_pdf(
    filepath: MarimoPath,
    *,
    include_outputs: bool,
    webpdf: bool,
    cli_args: SerializedCLIArgs,
    argv: list[str] | None,
    export_as: ExportPDFPreset | None,
    include_inputs: bool = True,
    rasterization_options: PDFRasterizationOptions | None = None,
    status_callback: PDFExportStatusCallback | None = None,
) -> tuple[bytes | None, bool]:
    file_manager = load_notebook(filepath.absolute_name)

    session_view: SessionView | None = None
    png_fallbacks: Mapping[CellId_t, str] | None = None
    did_error = False

    if include_outputs:
        emit_pdf_export_status(
            status_callback,
            phase="execute",
            message="executing notebook...",
        )
        with patch_html_for_non_interactive_output():
            # Using quiet=True to suppress runtime stdout/stderr since outputs
            # are captured in the session_view and will be included in the PDF
            (session_view, did_error) = await run_app_until_completion(
                file_manager,
                cli_args,
                argv,
                quiet=True,
            )
        emit_pdf_export_status(
            status_callback,
            phase="execute_complete",
            message="notebook execution finished.",
        )

        if (
            session_view is not None
            and rasterization_options is not None
            and rasterization_options.enabled
        ):
            from marimo._server.export._pdf_raster import (
                collect_pdf_png_fallbacks,
            )

            png_fallbacks = await collect_pdf_png_fallbacks(
                app=file_manager.app,
                session_view=session_view,
                filename=filepath.short_name,
                filepath=filepath.absolute_name,
                argv=argv,
                options=rasterization_options,
                status_callback=status_callback,
            )
    emit_pdf_export_status(
        status_callback,
        phase="prepare",
        message="serializing notebook for PDF rendering...",
    )
    exporter = Exporter()
    if export_as == "slides":
        pdf_data = await exporter.export_as_slides_pdf(
            app=file_manager.app,
            session_view=session_view,
            png_fallbacks=png_fallbacks,
            include_inputs=include_inputs,
            status_callback=status_callback,
        )
    else:
        pdf_data = exporter.export_as_pdf(
            app=file_manager.app,
            session_view=session_view,
            png_fallbacks=png_fallbacks,
            include_inputs=include_inputs,
            webpdf=webpdf,
            status_callback=status_callback,
        )
    if pdf_data is not None:
        emit_pdf_export_status(
            status_callback,
            phase="complete",
            message="done.",
        )
    return pdf_data, did_error


async def run_app_then_export_as_html(
    path: MarimoPath,
    include_code: bool,
    cli_args: SerializedCLIArgs,
    argv: list[str],
    *,
    asset_url: str | None = None,
) -> ExportResult:
    file_manager = load_notebook(path.absolute_name)

    # Inline the layout file, if it exists
    file_manager.app.inline_layout_file()

    config = get_default_config_manager(current_path=file_manager.path)
    resolved = config.get_config()
    display_config = resolved["display"]
    session_view, did_error = await run_app_until_completion(
        file_manager,
        cli_args,
        argv=argv,
    )
    # Export the session as HTML
    html, filename = Exporter().export_as_html(
        filename=file_manager.filename,
        app=file_manager.app,
        session_view=session_view,
        display_config=display_config,
        sharing_config=resolved.get("sharing"),
        request=ExportAsHTMLRequest(
            include_code=include_code,
            download=False,
            files=[],
            asset_url=asset_url,
        ),
    )
    return ExportResult(
        contents=html,
        download_filename=filename,
        did_error=did_error,
    )


async def run_app_then_export_as_wasm(
    path: MarimoPath,
    mode: Literal["edit", "run"],
    show_code: bool,
    cli_args: SerializedCLIArgs,
    argv: list[str],
    *,
    asset_url: str | None = None,
) -> ExportResult:
    """Execute notebook and export as WASM HTML with embedded session."""
    from marimo._session.state.serialize import (
        serialize_notebook,
        serialize_session_view,
    )

    file_manager = load_notebook(path.absolute_name)
    file_manager.app.inline_layout_file()

    config = get_default_config_manager(current_path=file_manager.path)
    resolved = config.get_config()
    display_config = resolved["display"]

    session_view, did_error = await run_app_until_completion(
        file_manager,
        cli_args,
        argv=argv,
    )

    session_snapshot = serialize_session_view(
        session_view,
        cell_ids=file_manager.app.cell_manager.cell_ids(),
        drop_virtual_file_outputs=True,
    )
    notebook_snapshot = serialize_notebook(
        session_view, file_manager.app.cell_manager
    )

    code = pin_pep723_dependencies_for_wasm(file_manager.app.to_py(), path)

    html, filename = Exporter().export_as_wasm(
        filename=file_manager.filename,
        app=file_manager.app,
        display_config=display_config,
        code=code,
        mode=mode,
        show_code=show_code,
        asset_url=asset_url,
        session_snapshot=session_snapshot,
        notebook_snapshot=notebook_snapshot,
        sharing_config=resolved.get("sharing"),
    )
    return ExportResult(
        contents=html,
        download_filename=filename,
        did_error=did_error,
    )


async def export_as_html_without_execution(
    path: MarimoPath,
    include_code: bool,
    *,
    asset_url: str | None = None,
) -> ExportResult:
    """Export a notebook to HTML without executing its cells."""
    from marimo._session.state.session_view import SessionView

    file_manager = load_notebook(path.absolute_name)

    # Inline the layout file, if it exists.
    file_manager.app.inline_layout_file()

    view = SessionView()
    for cell_data in file_manager.app.cell_manager.cell_data():
        view.last_executed_code[cell_data.cell_id] = cell_data.code

    config = get_default_config_manager(current_path=file_manager.path)
    display_config = config.get_config()["display"]

    html, filename = Exporter().export_as_html(
        filename=file_manager.filename,
        app=file_manager.app,
        session_view=view,
        display_config=display_config,
        request=ExportAsHTMLRequest(
            include_code=include_code,
            download=False,
            files=[],
            asset_url=asset_url,
        ),
    )
    return ExportResult(
        contents=html,
        download_filename=filename,
        did_error=False,
    )


async def run_app_then_export_as_reactive_html(
    path: MarimoPath,
    include_code: bool,
) -> ExportResult:
    from marimo._islands._island_generator import MarimoIslandGenerator

    generator = MarimoIslandGenerator.from_file(
        path.absolute_name, display_code=include_code
    )
    await generator.build()
    html = generator.render_html()
    basename = os.path.basename(path.absolute_name)
    filename = f"{os.path.splitext(basename)[0]}.html"
    return ExportResult(
        contents=html,
        download_filename=filename,
        did_error=False,
    )


async def run_app_until_completion(
    file_manager: AppFileManager,
    cli_args: SerializedCLIArgs,
    argv: list[str] | None,
    quiet: bool = False,
    persist_session: bool = True,
) -> tuple[SessionView, bool]:
    from marimo._session.consumer import SessionConsumer
    from marimo._session.events import SessionEventBus
    from marimo._session.session import SessionImpl

    instantiated_event = asyncio.Event()

    class RunUntilCompletionSessionConsumer(SessionConsumer):
        def __init__(self) -> None:
            self.did_error = False

        @property
        def consumer_id(self) -> ConsumerId:
            return ConsumerId("default")

        def notify(self, notification: KernelMessage) -> None:
            data = deserialize_kernel_message(notification)
            # Print errors to stderr (unless quiet mode)
            if isinstance(data, CellNotification):
                output = data.output
                console_output = data.console
                if output and output.channel == CellChannel.MARIMO_ERROR:
                    errors = cast(list[Error], output.data)
                    for err in errors:
                        # Not all errors are fatal
                        if is_unexpected_error(err):
                            if not quiet:
                                echo(
                                    f"{err.__class__.__name__}: {err.describe()}",
                                    file=sys.stderr,
                                )
                            self.did_error = True

                if console_output and not quiet:
                    console_as_list: list[CellOutput] = (
                        console_output
                        if isinstance(console_output, list)
                        else [console_output]
                    )
                    try:
                        for line in console_as_list:
                            # We print to stderr to not interfere with the
                            # piped output
                            mimetype = line.mimetype
                            if mimetype == "text/plain":
                                echo(line.data, file=sys.stderr, nl=False)
                    except Exception:
                        LOGGER.warning("Error printing console output")

            if isinstance(data, CompletedRunNotification):
                instantiated_event.set()

        def on_attach(
            self, session: Session, event_bus: SessionEventBus
        ) -> None:
            del session
            del event_bus

        def on_detach(self) -> None:
            return None

        def connection_state(self) -> ConnectionState:
            return ConnectionState.OPEN

    config_manager = get_default_config_manager(
        current_path=file_manager.path
    ).with_overrides(
        {
            "runtime": cast(
                RuntimeConfig,
                {
                    "on_cell_change": "autorun",
                    "auto_instantiate": True,
                    "auto_reload": "off",
                    "watcher_on_save": "lazy",
                    # We cast because we don't want to override the other
                    # config values
                },
            ),
        }
    )

    # Create a session
    session_consumer = RunUntilCompletionSessionConsumer()
    session = SessionImpl.create(
        # Any initialization ID will do
        initialization_id="_any_",
        session_consumer=session_consumer,
        # Run in EDIT mode so that console outputs are captured
        mode=SessionMode.EDIT,
        app_metadata=AppMetadata(
            query_params={},
            filename=file_manager.path,
            cli_args=cli_args,
            argv=argv,
            app_config=file_manager.app.config,
            docstring=extract_docstring_from_header(
                file_manager.app._app._header
            ),
        ),
        app_file_manager=file_manager,
        config_manager=config_manager,
        virtual_file_storage=None,
        redirect_console_to_browser=False,
        ttl_seconds=None,
        auto_instantiate=True,
    )

    # Run the notebook to completion once
    session.instantiate(
        InstantiateNotebookRequest(object_ids=[], values=[]),
        http_request=None,
    )
    await instantiated_event.wait()
    # Hack: yield to give the session view a chance to process the incoming
    # console operations.
    await asyncio.sleep(0.1)

    if persist_session:
        from marimo._server.export._session_cache import (
            persist_session_view_to_cache,
        )

        try:
            persist_session_view_to_cache(
                view=session.session_view,
                notebook_path=file_manager.path,
                cell_ids=file_manager.app.cell_manager.cell_ids(),
            )
        except Exception as e:
            LOGGER.warning(
                "Failed to persist session snapshot for %s: %s",
                file_manager.path,
                e,
            )

    # Stop distributor, terminate kernel process, etc -- all information is
    # captured by the session view.
    session.close()

    return session.session_view, session_consumer.did_error


def _with_filename(
    ir: NotebookSerialization, filename: str
) -> NotebookSerialization:
    return replace(ir, filename=filename)
