"""Lektor Markdown figures with arbitrary CSS classes and Markdown captions.

Syntax::

    :::figure right w50
    ![Alternative text](image.jpg)

    Caption with *Markdown* and [links](https://example.com/).
    :::

The first non-empty content line must be a Markdown image. Every remaining
line is treated as the optional caption.
"""

from __future__ import annotations

import html
import re
from typing import Any

import mistune
from lektor.pluginsystem import Plugin

FIGURE_PATTERN = re.compile(
    r" {0,3}:::figure(?:[ \t]+(?P<classes>[^\n]*?))?[ \t]*\n"
    r"(?P<body>[\s\S]*?)"
    r"\n {0,3}:::[ \t]*(?:\n+|$)"
)

# Mistune 0.8 renders this marker as inline HTML. The renderer mixin removes
# it and, at the same time, avoids wrapping the figure image in a paragraph.
_FIGURE_IMAGE_MARKER = "<!--lektor-markdown-figure:image-->"


def _normalise_classes(value: str | None) -> str:
    """Return whitespace-normalised, otherwise unrestricted CSS classes."""
    return " ".join((value or "").split())


def _class_attribute(classes: str) -> str:
    if not classes:
        return ""
    return f' class="{html.escape(classes, quote=True)}"'


def _split_figure(match: re.Match[str]) -> tuple[str, str, str]:
    """Extract classes, the image line, and the optional caption source."""
    classes = _normalise_classes(match.group("classes"))
    lines = match.group("body").splitlines()

    while lines and not lines[0].strip():
        lines.pop(0)
    while lines and not lines[-1].strip():
        lines.pop()

    if not lines:
        raise ValueError("A figure directive must contain a Markdown image.")

    image_source = lines.pop(0).strip()
    if not image_source.startswith("!["):
        raise ValueError(
            "The first non-empty line in a figure directive must be a Markdown image."
        )

    caption_source = "\n".join(lines).strip()
    return classes, image_source, caption_source


# ---------------------------------------------------------------------------
# Mistune 0.8 support (used by Lektor 3.3.x)
# ---------------------------------------------------------------------------

if hasattr(mistune, "BlockGrammar"):

    class FigureBlockGrammar(mistune.BlockGrammar):  # type: ignore[misc]
        figure = FIGURE_PATTERN

    class FigureBlockLexer(mistune.BlockLexer):  # type: ignore[misc]
        grammar_class = FigureBlockGrammar
        default_rules = list(mistune.BlockLexer.default_rules)
        default_rules.insert(default_rules.index("fences"), "figure")
        caption_rules = [rule for rule in default_rules if rule != "figure"]

        def parse_figure(self, match: re.Match[str]) -> None:
            classes, image_source, caption_source = _split_figure(match)
            class_attribute = _class_attribute(classes)

            # ``close_html`` is Mistune 0.8's raw block-HTML output token.
            self.tokens.append(
                {
                    "type": "close_html",
                    "text": f"<figure{class_attribute}>\n",
                }
            )

            # The marker lets FigureRendererMixin remove the otherwise
            # automatic <p> wrapper while retaining Lektor's normal inline
            # image renderer and attachment-relative URL resolution.
            self.tokens.append(
                {
                    "type": "paragraph",
                    "text": _FIGURE_IMAGE_MARKER + image_source,
                }
            )

            if caption_source:
                self.tokens.append({"type": "close_html", "text": "<figcaption>\n"})
                self.parse(caption_source, self.caption_rules)
                self.tokens.append({"type": "close_html", "text": "</figcaption>\n"})

            self.tokens.append({"type": "close_html", "text": "</figure>\n"})

    class FigureRendererMixin:
        """Remove the paragraph wrapper from the directive's image only."""

        def paragraph(self, text: str) -> str:
            if text.startswith(_FIGURE_IMAGE_MARKER):
                return text[len(_FIGURE_IMAGE_MARKER) :] + "\n"
            return super().paragraph(text)  # type: ignore[misc]

else:
    FigureBlockLexer = None  # type: ignore[assignment,misc]
    FigureRendererMixin = None  # type: ignore[assignment,misc]


# ---------------------------------------------------------------------------
# Mistune 2 support (used by Lektor's newer Markdown controller)
# ---------------------------------------------------------------------------


def _render_figure(text: str, classes: str) -> str:
    return f"<figure{_class_attribute(classes)}>\n{text}</figure>\n"


def _render_figure_image(text: str) -> str:
    return text + "\n"


def _render_figure_caption(text: str) -> str:
    return f"<figcaption>\n{text}</figcaption>\n"


def _render_ast_figure(children: list[dict[str, Any]], classes: str) -> dict[str, Any]:
    return {"type": "figure", "classes": classes.split(), "children": children}


def _render_ast_figure_image(children: list[dict[str, Any]]) -> dict[str, Any]:
    return {"type": "figure_image", "children": children}


def _render_ast_figure_caption(children: list[dict[str, Any]]) -> dict[str, Any]:
    return {"type": "figure_caption", "children": children}


def mistune_figure_plugin(md: Any) -> None:
    """Register the figure block with Mistune 2."""

    def parse_figure(
        block: Any, match: re.Match[str], state: dict[str, Any]
    ) -> dict[str, Any]:
        classes, image_source, caption_source = _split_figure(match)

        children: list[dict[str, Any]] = [
            {"type": "figure_image", "text": image_source}
        ]

        if caption_source:
            caption_rules = [rule for rule in block.rules if rule != "figure"]
            caption_children = block.parse(caption_source, state, caption_rules)
            children.append({"type": "figure_caption", "children": caption_children})

        return {
            "type": "figure",
            "children": children,
            "params": (classes,),
        }

    md.block.register_rule("figure", FIGURE_PATTERN, parse_figure)
    if "figure" not in md.block.rules:
        md.block.rules.insert(0, "figure")

    if md.renderer.NAME == "html":
        md.renderer.register("figure", _render_figure)
        md.renderer.register("figure_image", _render_figure_image)
        md.renderer.register("figure_caption", _render_figure_caption)
    elif md.renderer.NAME == "ast":
        md.renderer.register("figure", _render_ast_figure)
        md.renderer.register("figure_image", _render_ast_figure_image)
        md.renderer.register("figure_caption", _render_ast_figure_caption)


class MarkdownFigurePlugin(Plugin):
    name = "Markdown Figure"
    description = (
        "Adds :::figure blocks with arbitrary CSS classes and optional "
        "Markdown captions."
    )

    def on_markdown_config(self, config: Any, **extra: Any) -> None:
        # Newer Lektor controller using Mistune 2.
        if hasattr(config, "parser_options"):
            print("Using mistune 2")
            plugins = config.parser_options.setdefault("plugins", [])
            if mistune_figure_plugin not in plugins:
                plugins.append(mistune_figure_plugin)
            return

        # Stable Lektor 3.3 controller using Mistune 0.8.
        if FigureBlockLexer is None or FigureRendererMixin is None:
            raise RuntimeError("Unsupported Mistune version for Markdown Figure.")

        existing_block = config.options.get("block")
        if existing_block not in (None, FigureBlockLexer):
            raise RuntimeError(
                "Markdown Figure cannot be combined with another plugin "
                "that replaces Mistune's block lexer."
            )

        config.options["block"] = FigureBlockLexer
        if FigureRendererMixin not in config.renderer_mixins:
            config.renderer_mixins.insert(0, FigureRendererMixin)