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

import functools
import html
import json
import os
import re
from pathlib import Path
from textwrap import dedent
from typing import TYPE_CHECKING, Any, Literal, cast

from marimo._ast.app_config import _AppConfig
from marimo._config.config import MarimoConfig, PartialMarimoConfig
from marimo._convert.common.filename import parse_title
from marimo._convert.converters import MarimoConvert
from marimo._messaging.notification import ModelLifecycleNotification
from marimo._output.utils import uri_encode_component
from marimo._schemas.notebook import NotebookV1
from marimo._schemas.session import NotebookSessionV1
from marimo._server.tokens import SkewProtectionToken
from marimo._session.model import SessionMode
from marimo._session.notebook import read_css_file, read_html_head_file
from marimo._utils.versions import is_editable
from marimo._version import __version__

if TYPE_CHECKING:
    from marimo._server.api.endpoints.assets import LspWorkspace


MOUNT_CONFIG_TEMPLATE = "'{{ mount_config }}'"


_json_script_escapes = {
    ord(">"): "\\u003E",
    ord("<"): "\\u003C",
    ord("&"): "\\u0026",
}


def _html_escape(text: str) -> str:
    """Escape HTML special characters."""
    return html.escape(text, quote=True)


def json_script(data: Any) -> str:
    # See https://github.com/django/django/blob/main/django/utils/html.py#L88C1-L92C2
    # Only escape values that can break out of a script tag
    return json.dumps(data, sort_keys=True).translate(_json_script_escapes)


def _export_context_block(*, notebook_code: str) -> str:
    """Emit the trusted export marker consumed by islands/export runtime code."""
    return dedent(
        f"""
    <script data-marimo="true">
        Object.defineProperty(window, "__MARIMO_EXPORT_CONTEXT__", {{
            value: Object.freeze({{
                trusted: true,
                notebookCode: {json_script(notebook_code)},
            }}),
            writable: false,
            configurable: false,
        }});
    </script>
    """
    )


def _static_state_block(
    *, files: dict[str, str], model_notifications: list[Any]
) -> str:
    """Emit the static-export virtual file table as a frozen, non-configurable
    global. Locking the shape prevents notebook-authored content from
    redirecting `@file/...` fetches by mutating the map after emission."""
    return dedent(
        f"""
    <script data-marimo="true">
        Object.defineProperty(window, "__MARIMO_STATIC__", {{
            value: Object.freeze({{
                files: Object.freeze({json_script(files)}),
                modelNotifications: Object.freeze({json_script(model_notifications)}),
            }}),
            writable: false,
            configurable: false,
        }});
    </script>
    """
    )


def _get_mount_config(
    *,
    filename: str | None,
    cwd: str | None = None,
    lsp_workspace: LspWorkspace | None = None,
    mode: Literal["edit", "home", "read", "gallery"],
    server_token: SkewProtectionToken,
    user_config: MarimoConfig,
    config_overrides: PartialMarimoConfig,
    app_config: _AppConfig | None,
    version: str | None = None,
    show_app_code: bool = True,
    session_snapshot: NotebookSessionV1 | None = None,
    notebook_snapshot: NotebookV1 | None = None,
    runtime_config: list[dict[str, Any]] | None = None,
) -> str:
    """
    Return a JSON string with custom indentation and sorting.
    """

    options: dict[str, Any] = {
        "filename": filename or "",
        "cwd": cwd or "",
        "lsp_workspace": lsp_workspace,
        "mode": mode,
        "version": version or get_version(),
        "server_token": str(server_token),
        "user_config": user_config,
        "config_overrides": config_overrides,
        "app_config": _del_none_or_empty(app_config.asdict())
        if app_config
        else {},
        "view": {
            "showAppCode": show_app_code,
        },
        "notebook": notebook_snapshot,
        "session": session_snapshot,
        "runtime_config": runtime_config,
    }

    return """{{
            "filename": {filename},
            "cwd": {cwd},
            "lspWorkspace": {lsp_workspace},
            "mode": {mode},
            "version": {version},
            "serverToken": {server_token},
            "config": {user_config},
            "configOverrides": {config_overrides},
            "appConfig": {app_config},
            "view": {view},
            "notebook": {notebook},
            "session": {session},
            "runtimeConfig": {runtime_config},
        }}
""".format(**{k: json_script(v) for k, v in options.items()}).strip()


def home_page_template(
    html: str,
    base_url: str,
    user_config: MarimoConfig,
    config_overrides: PartialMarimoConfig,
    server_token: SkewProtectionToken,
    mode: SessionMode,
    asset_url: str | None = None,
) -> str:
    html = html.replace("{{ base_url }}", base_url)
    html = html.replace("{{ title }}", "marimo")
    html = html.replace("{{ filename }}", "")

    # TODO(Trevor): Legacy, required by VS Code plugin. Remove when plugin is updated (see frontend/index.html)
    html = html.replace("{{ version }}", get_version())
    html = html.replace(
        "{{ user_config }}", _html_escape(json.dumps(user_config))
    )
    html = html.replace("{{ server_token }}", str(server_token))

    html = _replace_asset_urls(html, asset_url)

    app_mode: Literal["home", "gallery"] = (
        "home" if mode == SessionMode.EDIT else "gallery"
    )
    html = html.replace(
        MOUNT_CONFIG_TEMPLATE,
        _get_mount_config(
            filename=None,
            mode=app_mode,
            server_token=server_token,
            user_config=user_config,
            config_overrides=config_overrides,
            app_config=None,
        ),
    )

    # Add custom CSS from display config
    html = _inject_custom_css_for_config(html, user_config)
    html = _inject_custom_css_for_config(html, config_overrides)
    return html


def opengraph_metadata_template(
    *,
    base_url: str,
    mode: SessionMode,
    app_config: _AppConfig,
    filename: str | None,
    filepath: str | None,
    execute_opengraph_generators: bool = False,
) -> str:
    """Return OpenGraph `<meta>` tags for a notebook, or an empty string."""
    if not filepath:
        return ""

    try:
        from marimo._metadata.opengraph import (
            OpenGraphContext,
            is_https_url,
            resolve_opengraph_metadata,
        )

        file_key = (
            filename if filename and not Path(filename).is_absolute() else None
        )
        opengraph = resolve_opengraph_metadata(
            filepath,
            app_title=app_config.app_title,
            context=OpenGraphContext(
                filepath=filepath,
                file_key=file_key,
                base_url=base_url,
                mode=mode.value,
            ),
            execute_generator=execute_opengraph_generators,
        )
    except Exception:
        return ""

    if not opengraph.title and not opengraph.description:
        return ""

    if opengraph.image and is_https_url(opengraph.image):
        thumbnail_url = opengraph.image
    else:
        # Server-resolvable thumbnail URL, falling back to a placeholder when
        # no screenshot exists.
        thumbnail_url = f"{base_url}/og/thumbnail"
        if filename and not Path(filename).is_absolute():
            thumbnail_url = (
                f"{thumbnail_url}?file={uri_encode_component(filename)}"
            )

    meta_tags: list[str] = []
    if opengraph.title:
        meta_tags.append(
            f'<meta property="og:title" content="{_html_escape(opengraph.title)}" />'
        )
    if opengraph.description:
        meta_tags.append(
            f'<meta property="og:description" content="{_html_escape(opengraph.description)}" />'
        )
        meta_tags.append(
            f'<meta name="description" content="{_html_escape(opengraph.description)}" />'
        )
    meta_tags.append(
        f'<meta property="og:image" content="{_html_escape(thumbnail_url)}" />'
    )
    meta_tags.append(
        '<meta name="twitter:card" content="summary_large_image" />'
    )
    return "\n".join(meta_tags)


def notebook_page_template(
    *,
    html: str,
    base_url: str,
    user_config: MarimoConfig,
    config_overrides: PartialMarimoConfig,
    server_token: SkewProtectionToken,
    app_config: _AppConfig,
    filename: str | None,
    filepath: str | None = None,
    lsp_workspace: LspWorkspace | None = None,
    mode: SessionMode,
    session_snapshot: NotebookSessionV1 | None = None,
    notebook_snapshot: NotebookV1 | None = None,
    runtime_config: list[dict[str, Any]] | None = None,
    asset_url: str | None = None,
    html_head: str | None = None,
    execute_opengraph_generators: bool = False,
) -> str:
    html = html.replace("{{ base_url }}", base_url)

    # When we have a remote URL, let's pre-populate the index.html page
    # with a view of the notebook.
    if runtime_config and notebook_snapshot is None:
        # Prefer the absolute path for IO, since `filename` can be a display
        # path (workspace-relative) in gallery mode.
        path = filepath or filename
        if path:
            path_obj = Path(path)
            if path_obj.exists():
                notebook_snapshot = MarimoConvert.from_py(
                    path_obj.read_text(encoding="utf-8")
                ).to_notebook_v1()

    html = html.replace("{{ filename }}", _html_escape(filename or ""))

    # TODO(Trevor): Legacy, required by VS Code plugin. Remove when plugin is updated (see frontend/index.html)
    html = html.replace("{{ version }}", get_version())
    html = html.replace(
        "{{ user_config }}", _html_escape(json.dumps(user_config))
    )
    html = html.replace("{{ server_token }}", str(server_token))

    html = _replace_asset_urls(html, asset_url)

    cwd = os.path.dirname(filepath) if filepath else None
    html = html.replace(
        MOUNT_CONFIG_TEMPLATE,
        _get_mount_config(
            filename=filename,
            cwd=cwd,
            lsp_workspace=lsp_workspace,
            mode="read" if mode == SessionMode.RUN else "edit",
            server_token=server_token,
            user_config=user_config,
            config_overrides=config_overrides,
            app_config=app_config,
            runtime_config=runtime_config,
            notebook_snapshot=notebook_snapshot,
            session_snapshot=session_snapshot,
        ),
    )

    html = html.replace(
        "{{ title }}",
        _html_escape(
            parse_title(filename)
            if app_config.app_title is None
            else app_config.app_title
        ),
    )

    opengraph_tags = opengraph_metadata_template(
        base_url=base_url,
        mode=mode,
        app_config=app_config,
        filename=filename,
        filepath=filepath,
        execute_opengraph_generators=execute_opengraph_generators,
    )
    if opengraph_tags:
        html = html.replace("</head>", f"{opengraph_tags}\n</head>")

    # If has custom css, inline the css and add to the head
    # Prefer the absolute path for IO, since `filename` can be a display
    # path (workspace-relative) in gallery/workspace mode.
    if app_config.css_file:
        css_contents = read_css_file(
            app_config.css_file, filename=filepath or filename
        )
        if css_contents:
            css_contents = _custom_css_block(css_contents)
            # Append to head
            html = html.replace("</head>", f"{css_contents}</head>")

    # Add custom CSS from display config
    html = _inject_custom_css_for_config(
        html, user_config, filepath or filename
    )
    html = _inject_custom_css_for_config(
        html, config_overrides, filepath or filename
    )

    # Add global HTML head contents if specified (from create_asgi_app)
    if html_head:
        html = html.replace("</head>", f"{html_head}</head>")

    # html_head_file is blocked in edit mode: it can contain arbitrary scripts
    # and markup that could exfiltrate data or redress the UI. CSS-only styling
    # (css_file) is permitted. Run mode is unaffected.
    if mode == SessionMode.RUN and app_config.html_head_file:
        head_contents = read_html_head_file(
            app_config.html_head_file, filename=filepath or filename
        )
        if head_contents:
            # Append to head
            html = html.replace("</head>", f"{head_contents}</head>")

    return html


def static_notebook_template(
    html: str,
    user_config: MarimoConfig,
    config_overrides: PartialMarimoConfig,
    server_token: SkewProtectionToken,
    app_config: _AppConfig,
    filepath: str | None,
    code: str,
    code_hash: str,
    session_snapshot: NotebookSessionV1,
    notebook_snapshot: NotebookV1,
    files: dict[str, str],
    model_notifications: list[ModelLifecycleNotification] | None = None,
    asset_url: str | None = None,
) -> str:
    if asset_url is None:
        version = str(__version__).replace(".dev", "-dev")
        asset_url = f"https://cdn.jsdelivr.net/npm/@marimo-team/frontend@{version}/dist"

    html = html.replace("{{ base_url }}", "")
    filename = os.path.basename(filepath or "")
    html = html.replace("{{ filename }}", _html_escape(filename))

    # We don't need all this user config when we export the notebook,
    # but we do need some:
    # - display.theme
    # - display.cell_output
    html = html.replace(
        MOUNT_CONFIG_TEMPLATE,
        _get_mount_config(
            filename=filename,
            cwd=None,  # We don't need the cwd for static notebooks
            mode="read",
            server_token=server_token,
            user_config=user_config,
            config_overrides=config_overrides,
            app_config=app_config,
            session_snapshot=session_snapshot,
            notebook_snapshot=notebook_snapshot,
            runtime_config=None,
        ),
    )

    html = html.replace(
        "{{ title }}",
        _html_escape(
            parse_title(filepath)
            if app_config.app_title is None
            else app_config.app_title
        ),
    )

    static_block = _static_state_block(
        files=files,
        model_notifications=[
            n.to_json_serializable() for n in model_notifications or []
        ],
    )
    static_block += _export_context_block(notebook_code=code)

    # Add HTML head file contents if specified
    if app_config.html_head_file:
        head_contents = read_html_head_file(
            app_config.html_head_file, filename=filepath
        )
        if head_contents:
            static_block += dedent(
                f"""
            {head_contents}
            """
            )

    # If has custom css, inline the css and add to the head
    if app_config.css_file:
        css_contents = read_css_file(app_config.css_file, filename=filepath)
        if css_contents:
            static_block += _custom_css_block(css_contents)

    # Add custom CSS from display config
    static_block = _inject_custom_css_for_config(
        static_block, user_config, filepath
    )
    static_block = _inject_custom_css_for_config(
        static_block, config_overrides, filepath
    )

    code_block = dedent(
        f"""
    <marimo-code hidden="">
        {uri_encode_component(code)}
    </marimo-code>
    """
    )
    if not code:
        code_block = '<marimo-code hidden=""></marimo-code>'

    # Add a 256-bit hash of the code, for cache busting or CI checks
    code_block += (
        f'\n<marimo-code-hash hidden="">{code_hash}</marimo-code-hash>\n'
    )

    # Replace all relative href and src with absolute URL
    html = _replace_asset_urls(html, asset_url)

    # Append to head
    html = html.replace("</head>", f"{static_block}</head>")
    # Append to body
    html = html.replace("</body>", f"{code_block}</body>")

    html = _inject_custom_css_for_config(html, user_config, filepath)
    html = _inject_custom_css_for_config(html, config_overrides, filepath)
    return html


def wasm_notebook_template(
    *,
    html: str,
    version: str,
    filename: str,
    user_config: MarimoConfig,
    config_overrides: PartialMarimoConfig,
    app_config: _AppConfig,
    mode: Literal["edit", "run"],
    code: str,
    show_code: bool,
    asset_url: str | None = None,
    session_snapshot: NotebookSessionV1 | None = None,
    notebook_snapshot: NotebookV1 | None = None,
) -> str:
    """Template for WASM notebooks."""
    import re

    body = html

    if asset_url is not None:
        body = re.sub(r'="./assets/', f'="{asset_url}/assets/', body)

    body = body.replace("{{ base_url }}", "")
    body = body.replace(
        "{{ title }}",
        _html_escape(
            parse_title(filename)
            if app_config.app_title is None
            else app_config.app_title
        ),
    )

    body = body.replace("{{ filename }}", _html_escape("notebook.py"))
    body = body.replace(
        MOUNT_CONFIG_TEMPLATE,
        _get_mount_config(
            # WASM runtime currently expect this to be notebook.py instead of the actual filename
            filename="notebook.py",
            mode="edit" if mode == "edit" else "read",
            server_token=SkewProtectionToken("unused"),
            user_config=user_config,
            config_overrides=config_overrides,
            app_config=app_config,
            version=version,
            show_app_code=show_code,
            runtime_config=None,
            session_snapshot=session_snapshot,
            notebook_snapshot=notebook_snapshot,
        ),
    )

    body = body.replace(
        "</head>", '<marimo-wasm hidden=""></marimo-wasm></head>'
    )

    warning_script = """
    <script>
        if (window.location.protocol === 'file:') {
            alert('Warning: This file must be served by an HTTP server to function correctly.');
        }
    </script>
    """
    body = body.replace("</head>", f"{warning_script}</head>")

    # Hide save button in WASM mode
    wasm_styles = """
    <style>
        #save-button {
            display: none !important;
        }
        #filename-input {
            display: none !important;
        }
    </style>
    """
    body = body.replace("</head>", f"{wasm_styles}</head>")

    # If has custom css, inline the css and add to the head
    if app_config.css_file:
        css_contents = read_css_file(app_config.css_file, filename=filename)
        if css_contents:
            css_contents = _custom_css_block(css_contents)
            # Append to head
            body = body.replace("</head>", f"{css_contents}</head>")

    # Add custom CSS from display config
    body = _inject_custom_css_for_config(body, user_config, filename)
    body = _inject_custom_css_for_config(body, config_overrides, filename)

    # Add HTML head file contents if specified
    if app_config.html_head_file:
        head_contents = read_html_head_file(
            app_config.html_head_file, filename=filename
        )
        if head_contents:
            # Append to head
            body = body.replace("</head>", f"{head_contents}</head>")

    body = body.replace(
        "</head>",
        (
            f"{_export_context_block(notebook_code=code)}"
            f'<marimo-code hidden="">{uri_encode_component(code)}</marimo-code></head>'
        ),
    )

    return body


def inject_script(html: str, script: str) -> str:
    """Inject a script into the HTML before the closing body tag."""
    script_tag = f"<script>{script}</script>"
    return html.replace("</body>", f"{script_tag}</body>")


def _del_none_or_empty(d: Any) -> Any:
    return {
        key: (
            _del_none_or_empty(cast(Any, value))
            if isinstance(value, dict)
            else value
        )
        for key, value in d.items()
        if value is not None and value != []
    }


@functools.lru_cache(maxsize=1)
def get_version() -> str:
    # Invariant for the lifetime of the process; cache so the three
    # render-path call sites (mount_config + two .replace passes) don't
    # each pay the importlib.metadata lookup.
    return (
        f"{__version__} (editable)" if is_editable("marimo") else __version__
    )


# HTML5 terminates a <style> raw-text element at the first "</style"
# (case-insensitive, optional whitespace). Neutralize that sequence so
# notebook-controlled CSS cannot break out of the block and inject HTML.
_STYLE_END_RE = re.compile(r"</(?=\s*style)", re.IGNORECASE)


def _sanitize_css_for_style_block(css_contents: str) -> str:
    # "\/" is a valid CSS escape for "/", so the CSS tokenizer still sees
    # a (harmlessly invalid) token, while the HTML parser no longer
    # matches the end-tag.
    return _STYLE_END_RE.sub(r"<\\/", css_contents)


def _custom_css_block(css_contents: str) -> str:
    # marimo-custom is used by the frontend to identify this stylesheet
    # comes from marimo
    safe = _sanitize_css_for_style_block(css_contents)
    return f"<style title='marimo-custom'>{safe}</style>"


def _inject_custom_css_for_config(
    html: str,
    config: MarimoConfig | PartialMarimoConfig,
    filename: str | None = None,
) -> str:
    """Inject custom CSS from display config into HTML."""
    custom_css = config.get("display", {}).get("custom_css", [])
    if not custom_css:
        return html

    css_contents: list[str] = []
    for css_path in custom_css:
        css_content = read_css_file(css_path, filename=filename)
        if css_content:
            css_contents.append(_custom_css_block(css_content))

    if not css_contents:
        return html

    css_block = "\n".join(css_contents)
    return html.replace("</head>", f"{css_block}</head>")


def _replace_asset_urls(html: str, asset_url: str | None) -> str:
    """Replace asset URLs with the given asset URL.

    These are naturally relative URLs. This can be used to load assets
    from a CDN instead of from the marimo server.

    The asset URL can be parameterized with {version}
    """
    if asset_url is None:
        return html

    if "{version}" in asset_url:
        asset_url = asset_url.replace("{version}", __version__)

    return (
        html.replace("href='./", f"crossorigin='anonymous' href='{asset_url}/")
        .replace("src='./", f"crossorigin='anonymous' src='{asset_url}/")
        .replace('href="./', f'crossorigin="anonymous" href="{asset_url}/')
        .replace('src="./', f'crossorigin="anonymous" src="{asset_url}/')
    )
