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

import json
import re
from typing import Any
from urllib.request import urlopen

from marimo._config.config import Theme
from marimo._dependencies.dependencies import DependencyManager
from marimo._loggers import marimo_logger
from marimo._messaging.mimetypes import METADATA_KEY, KnownMimeType, MimeBundle
from marimo._output.formatters.formatter_factory import FormatterFactory
from marimo._plugins.core.media import io_to_data_url
from marimo._plugins.ui._impl.altair_chart import (
    AltairChartType,
    chart_to_json,
    get_chart_mimetype,
    maybe_fix_vegafusion_background,
)

LOGGER = marimo_logger()


class AltairFormatter(FormatterFactory):
    @staticmethod
    def package_name() -> str:
        return "altair"

    def register(self) -> None:
        import altair

        from marimo._output import formatting
        from marimo._plugins.ui._impl.charts.altair_transformer import (
            register_transformers,
        )

        # add marimo transformers
        register_transformers()

        @formatting.formatter(altair.TopLevelMixin)
        def _show_chart(chart: AltairChartType) -> tuple[KnownMimeType, str]:
            import altair as alt

            # Try to get the _repr_mimebundle_ method from the chart
            # If its HTML, we want to handle this ourselves
            # if its svg, vega, or png, then we want to pass that instead
            # because that means the user has configured the that renderer
            mimebundle: MimeBundle | tuple[MimeBundle, MimeBundle] = {}
            try:
                mimebundle = chart._repr_mimebundle_() or {}  # type: ignore
            except Exception as e:
                LOGGER.warning("Failed to get mimebundle from chart: %s", e)

            # When the mimebundle is a tuple, it follows the format
            # (data_dict, metadata_dict).
            if isinstance(mimebundle, tuple) and "image/png" in mimebundle[0]:
                return _format_png_mimebundle(mimebundle)

            # Handle where there are multiple mime types
            # return as a mimebundle
            if len(mimebundle) > 1:
                return (
                    "application/vnd.marimo+mimebundle",
                    json.dumps(mimebundle),
                )

            # Handle non-HTML mime types
            non_html_mime_types: list[KnownMimeType] = [
                "image/svg+xml",
                "image/png",
                "application/vnd.vega.v5+json",
                "application/vnd.vegalite.v5+json",
                "application/vnd.vega.v6+json",
                "application/vnd.vegalite.v6+json",
            ]
            for mime_type in non_html_mime_types:
                if mime_type in mimebundle:
                    mime_response = mimebundle[mime_type]
                    if isinstance(mime_response, bytes):
                        data_url = io_to_data_url(mime_response, mime_type)
                        return (mime_type, data_url or "")
                    if isinstance(mime_response, str):
                        if (
                            mime_type == "image/svg+xml"
                            and not altair.renderers.options.get("raw_svg")
                        ):
                            _maybe_warn_external_resources(mime_response)
                            svg_bytes = mime_response.encode()
                            data_url = io_to_data_url(svg_bytes, mime_type)
                            return (mime_type, data_url or "")
                        return mime_type, mime_response
                    return mime_type, json.dumps(mime_response)

            chart = _apply_embed_options(chart)
            chart = maybe_fix_vegafusion_background(chart)

            # If vegafusion is enabled, just wrap in altair_chart
            if alt.data_transformers.active.startswith("vegafusion"):
                return (
                    get_chart_mimetype(spec_format="vega"),
                    chart_to_json(chart=chart, spec_format="vega"),
                )

            # If the user has not set the max_rows option, we set it to 20_000
            # since we are able to handle the larger sizes (default is 5000)
            if "max_rows" not in alt.data_transformers.options:
                alt.data_transformers.options["max_rows"] = 20_000

            # Return the chart as a vega-lite chart with embed options
            return (
                get_chart_mimetype(spec_format="vega-lite"),
                chart_to_json(chart=chart, validate=False),
            )

    def apply_theme(self, theme: Theme) -> None:
        del theme
        # We don't need to apply this here because the theme is set in the
        # vega-lite component


def _format_png_mimebundle(
    png_mimebundle: tuple[MimeBundle, MimeBundle],
) -> tuple[KnownMimeType, str]:
    data_url = io_to_data_url(png_mimebundle[0]["image/png"], "image/png")
    metadata = png_mimebundle[1]["image/png"]

    mimebundle = {
        "image/png": data_url or "",
        METADATA_KEY: {
            "image/png": {
                "width": round(metadata["width"]),
                "height": round(metadata["height"]),
            }
        },
    }
    return "application/vnd.marimo+mimebundle", json.dumps(mimebundle)


# Check if the SVG contains external resources that may not render
# correctly when encoded as a Data URL.
# https://github.com/marimo-team/marimo/pull/9104
def _maybe_warn_external_resources(svg: str) -> None:
    # Strictly detecting external resource usage in SVG is difficult;
    # as a heuristic, we check for 'href' or 'xlink:href' attributes
    # that point to external resources (ignoring internal '#' or 'data:' URLs).
    if re.search(
        r'<[^>]*\b(?:xlink:)?href\s*=\s*["\'](?!\s*(?:#|data:))[^"\']+', svg
    ):
        msg = (
            "This SVG contains external resources (href/xlink:href) "
            "that may not render correctly when encoded as a Data URL. "
            "If images are missing, try enabling raw SVG rendering with: "
            "altair.renderers.enable('svg', raw_svg=True)."
        )
        LOGGER.warning(msg)


# This is only needed since it seems that altair does not
# handle this internally.
# https://github.com/marimo-team/marimo/issues/2302
def _apply_embed_options(chart: AltairChartType) -> AltairChartType:
    import altair as alt

    # Respect user-set embed options
    # Note:
    # The python key is `embed_options`
    # The javascript key is `embedOptions`
    embed_options = alt.renderers.options.get("embed_options", {})
    prev_usermeta = {} if alt.Undefined is chart.usermeta else chart.usermeta

    embed_options = _apply_format_locales(embed_options)

    chart["usermeta"] = {
        **prev_usermeta,
        "embedOptions": {
            **embed_options,
            **prev_usermeta.get("embedOptions", {}),
        },
    }
    return chart


FETCH_TIMEOUT = 3
TIME_FORMAT_LOCALE_URL = (
    "https://unpkg.com/d3-time-format@latest/locale/{locale}.json"
)
FORMAT_LOCALE_URL = "https://unpkg.com/d3-format@latest/locale/{locale}.json"


def _apply_format_locales(embed_options: dict[str, Any]) -> dict[str, Any]:
    """Apply format localizations to embed options using either vl_convert or d3 format files."""

    def get_time_format_locale(locale: str) -> dict[str, Any]:
        try:
            if DependencyManager.vl_convert_python.has():
                import vl_convert as vlc  # type: ignore

                return dict(vlc.get_time_format_locale(locale))  # type: ignore[arg-type]
            else:
                with urlopen(
                    TIME_FORMAT_LOCALE_URL.format(locale=locale),
                    timeout=FETCH_TIMEOUT,
                ) as response:
                    return dict(json.loads(response.read()))
        except Exception as e:
            LOGGER.warning(f"Error getting time format locale: {e}")
            return {}

    def get_format_locale(locale: str) -> dict[str, Any]:
        try:
            if DependencyManager.vl_convert_python.has():
                import vl_convert as vlc  # type: ignore

                return dict(vlc.get_format_locale(locale))  # type: ignore[arg-type]
            else:
                with urlopen(
                    FORMAT_LOCALE_URL.format(locale=locale),
                    timeout=FETCH_TIMEOUT,
                ) as response:
                    return dict(json.loads(response.read()))
        except Exception as e:
            LOGGER.warning(f"Error getting format locale: {e}")
            return {}

    if "timeFormatLocale" in embed_options:
        time_format_locale = embed_options["timeFormatLocale"]
        if isinstance(time_format_locale, str):
            embed_options["timeFormatLocale"] = get_time_format_locale(
                time_format_locale
            )

    if "formatLocale" in embed_options:
        format_locale = embed_options["formatLocale"]
        if isinstance(format_locale, str):
            embed_options["formatLocale"] = get_format_locale(format_locale)

    return embed_options
