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

import os
import weakref
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Literal, cast

from marimo._messaging.mimetypes import KnownMimeType
from marimo._output.mime import MIME
from marimo._output.rich_help import mddoc
from marimo._output.utils import flatten_string
from marimo._utils.methods import getcallable

if TYPE_CHECKING:
    from collections.abc import Iterator, Sequence

    from marimo._plugins.ui._core.ui_element import UIElement
    from marimo._plugins.ui._impl.batch import batch as batch_plugin


def _hypertext_cleanup(virtual_filenames: list[str]) -> None:
    """Cleanup side-effects related to initialization of Html."""
    from marimo._runtime.context import (
        ContextNotInitializedError,
        get_context,
    )

    try:
        ctx = get_context()
    except ContextNotInitializedError:
        return

    if ctx is not None and ctx.virtual_files_supported:
        for f in virtual_filenames:
            ctx.virtual_file_registry.dereference(f)


@mddoc
@dataclass
class Html(MIME):
    """A wrapper around HTML text that can be used as an output.

    Output an `Html` object as the last expression of a cell to render it in
    your app.

    Use f-strings to embed Html objects as text into other HTML or markdown
    strings. For example:

    ```python3
    hello_world = Html("<h2>Hello, World</h2>")
    Html(
        f'''
        <h1>Hello, Universe!</h1>
        {hello_world}
        '''
    )
    ```

    Attributes:
        text: a string of HTML

    Args:
        text: a string of HTML

    Methods:
        batch: convert this HTML element into a batched UI element
        callout: wrap this element in a callout
        center: center this element in the output area
        right: right-justify this element in the output area
    """

    # Some libraries (e.g. polars) will serialize dataclasses so we add this
    # field to serialize the mimetype. This is to support rich display in tables/dfs.
    _serialized_mime_bundle: dict[Literal["mimetype", "data"], str] = field(
        default_factory=dict,
        repr=False,
        init=False,
    )

    def __init__(self, text: str) -> None:
        """Initialize the HTML element.

        Subclasses of HTML MUST call this method.
        """
        self._text = text
        mimetype, data = self._mime_()

        self._serialized_mime_bundle = {
            "mimetype": mimetype,
            "data": data,
        }
        # Whenever _serialized_mime_bundle is set, ensure a public copy exists.
        # This avoids declaring a public attribute (does not show up in docs)
        # Pandas does not serialize private variables, so we need this.
        self.__setattr__(
            "serialized_mime_bundle", self._serialized_mime_bundle
        )

        # A list of the virtual file names referenced by this HTML element.
        self._virtual_filenames: list[str] = []

        from marimo._runtime.context import (
            ContextNotInitializedError,
            get_context,
        )

        try:
            ctx = get_context()
        except ContextNotInitializedError:
            return

        # Virtual File Refcounting
        #
        # HTML elements are responsible for maintaining the reference counts
        # of virtual files: virtual files cannot be disposed while HTML
        # elements reference them. For example, a user might cache HTML
        # referencing a virtual file if they create it using functools.cache.
        #
        # flatten the text to make sure searching isn't broken by newlines
        flat_text = flatten_string(self._text)
        for virtual_filename in ctx.virtual_file_registry.filenames():
            if virtual_filename in flat_text:
                ctx.virtual_file_registry.reference(virtual_filename)
                self._virtual_filenames.append(virtual_filename)

        # Dereference virtual files on object destruction
        finalizer = weakref.finalize(
            self, _hypertext_cleanup, self._virtual_filenames
        )
        finalizer.atexit = False

    @property
    def text(self) -> str:
        """A string of HTML representing this element."""
        return self._text

    def _mime_(self) -> tuple[KnownMimeType, str]:
        if not is_non_interactive():
            return ("text/html", self.text)

        # Try PNG representation first (for objects with _repr_png_)
        repr_png = getcallable(self, "_repr_png_")
        if repr_png is not None:
            png_bytes = cast(bytes, repr_png())
            return ("image/png", png_bytes.decode())

        # Try markdown representation (for objects with _repr_markdown_)
        repr_markdown = getcallable(self, "_repr_markdown_")
        if repr_markdown is not None:
            markdown_text = cast(str, repr_markdown())
            return ("text/markdown", markdown_text)

        # Default to HTML
        return ("text/html", self.text)

    def __format__(self, spec: str) -> str:
        """Format `self` as HTML text"""
        del spec
        return " ".join(
            [line.strip() for line in self.text.strip().split("\n")]
        )

    @mddoc
    def batch(self, **elements: UIElement[Any, Any]) -> batch_plugin:
        """Convert an HTML object with templated text into a UI element.

        This method lets you create custom UI elements that are represented
        by arbitrary HTML.

        Example:
            ```python3
            user_info = mo.md(
                '''
                - What's your name?: {name}
                - When were you born?: {birthday}
                '''
            ).batch(name=mo.ui.text(), birthday=mo.ui.date())
            ```

            In this example, `user_info` is a UI Element whose output is markdown
            and whose value is a dict with keys `'name'` and '`birthday`'
            (and values equal to the values of their corresponding elements).

        Args:
            elements: the UI elements to interpolate into the HTML template.
        """
        from marimo._plugins.ui._impl.batch import batch as batch_plugin

        return batch_plugin(html=self, elements=elements)

    @mddoc
    def center(self) -> Html:
        """Center an item.

        Example:
            ```python3
            mo.md("# Hello, world").center()
            ```

        Returns:
            An `Html` object.
        """
        from marimo._plugins.stateless import flex

        return flex.vstack([_BlockWrapped(self)], align="center", gap=0)

    @mddoc
    def right(self) -> Html:
        """Right-justify.

        Example:
            ```python3
            mo.md("# Hello, world").right()
            ```

        Returns:
            An `Html` object.
        """
        from marimo._plugins.stateless import flex

        return flex.vstack([_BlockWrapped(self)], align="end", gap=0)

    @mddoc
    def left(self) -> Html:
        """Left-justify.

        Example:
            ```python3
            mo.md("# Hello, world").left()
            ```

        Returns:
            An `Html` object.
        """
        from marimo._plugins.stateless import flex

        return flex.vstack([_BlockWrapped(self)], align="start", gap=0)

    @mddoc
    def callout(
        self,
        kind: Literal[
            "neutral", "danger", "warn", "success", "info"
        ] = "neutral",
        title: str | None = None,
    ) -> Html:
        """Create a callout containing this HTML element.

        A callout renders your HTML element in a flat, bordered box — the
        same style as markdown admonitions — emphasizing its importance. You
        can style the callout for different situations with the `kind`
        argument, and optionally add a bold `title` line.

        Examples:
            ```python3
            mo.md("Hooray, you did it!").callout(kind="success")
            ```

            ```python3
            mo.md("It's dangerous to go alone!").callout(
                kind="warn", title="Warning"
            )
            ```
        """

        from marimo._plugins.stateless.callout import callout as _callout

        return _callout(self, kind=kind, title=title)

    @mddoc
    def style(
        self, style: dict[str, Any] | None = None, **kwargs: Any
    ) -> Html:
        """Wrap an object in a styled container.

        Example:
            ```python
            mo.md("...").style({"max-height": "300px", "overflow": "auto"})
            mo.md("...").style(max_height="300px", overflow="auto")
            ```

        Args:
            style: an optional dict of CSS styles, keyed by property name
            **kwargs: CSS styles as keyword arguments
        """
        from marimo._plugins.stateless import style as _style

        return _style.style(self, style=style, **kwargs)

    def _repr_html_(self) -> str:
        return self.text


class ContainerHtml(Html):
    """Base class for Html wrappers that contain other renderables.

    A `ContainerHtml` holds *strong* references to its children and re-renders
    from their live `.text` on every access. This matters for two reasons:

    - UI elements are only weakly referenced by the UI element registry, so a
      container that merely froze `child.text` at construction time would let
      its children be garbage collected -- breaking interactivity for elements
      returned from helper functions. Holding strong references keeps the
      children alive, mirroring `mo.hstack`/`mo.vstack`.
    - Mutable children (e.g. `mo.status.spinner`) re-render live rather than
      being frozen at construction time.

    Subclasses store any extra rendering state on `self` *before* calling
    `super().__init__()` and implement `_build_text` in terms of
    `self._children`.
    """

    def __init__(self, children: Sequence[Html]) -> None:
        self._children: list[Html] = list(children)
        super().__init__(self._build_text())

    def _build_text(self) -> str:
        raise NotImplementedError

    @property
    def text(self) -> str:  # type: ignore[override]
        """Re-render children live on every access."""
        return self._build_text()


class _BlockWrapped(ContainerHtml):
    """Wraps another Html in a plain block `<div>`.

    Used by :meth:`Html.center`, :meth:`Html.left`, and :meth:`Html.right`
    so that the wrapped content renders inside a single block container
    (preserving normal block flow and margin collapsing between inner
    elements) rather than becoming multiple flex-column siblings when the
    inner wrapper uses `display: contents` (as `mo.md` does).

    The inner Html is re-rendered on every `text` access so mutable
    children (e.g. `mo.status.spinner`) keep updating live.
    """

    def __init__(self, inner: Html) -> None:
        super().__init__([inner])

    def _build_text(self) -> str:
        from marimo._output.builder import h

        return h.div(self._children[0].text)


MARIMO_NO_JS_KEY = "MARIMO_NO_JS"


@contextmanager
def patch_html_for_non_interactive_output() -> Iterator[None]:
    """
    Patch Html to return text/markdown for simpler non-interactive outputs,
    that can be rendered without JS/CSS (just as in the GitHub viewer).
    """
    # HACK: we must set MARIMO_NO_JS since the rendering may happen in another
    # thread
    # This won't work when we are running a marimo server and are auto-exporting
    # with this enabled.
    old_no_js = os.getenv(MARIMO_NO_JS_KEY, "false")
    try:
        os.environ[MARIMO_NO_JS_KEY] = "true"
        yield
    finally:
        os.environ[MARIMO_NO_JS_KEY] = old_no_js


def is_non_interactive() -> bool:
    """Whether to render HTML objects as best as possible assuming
    that this will be rendered without javascript.

    For example, prefer images or markdown over rich HTML output.
    """
    return os.getenv(MARIMO_NO_JS_KEY, "false").lower() == "true"
