from functools import partial

from markdown_it import MarkdownIt
from markdown_it.common.utils import arrayReplaceAt
from markdown_it.rules_core import StateCore
from markdown_it.token import Token


def img_figures_plugin(md: MarkdownIt, lazyload: bool) -> None:
    """Plugin which wraps images with plain text descriptions in figures.

    :param MarkdownIt md: Parser
    :param bool lazyload: Whether to set loading to lazy on images
    """

    def _img_figures(state: StateCore, lazyload: bool) -> None:
        for inline_token in state.tokens:
            if inline_token.type != "inline":
                continue

            tokens = inline_token.children

            assert tokens is not None
            i = len(tokens)
            while i >= 1:
                i -= 1
                assert isinstance(tokens, list)
                current_token = tokens[i]

                if current_token.type == "image":
                    if lazyload:
                        current_token.attrs["loading"] = "lazy"
                    if current_token.children:
                        nodes = []
                        level = current_token.level

                        token = Token("figure_open", "figure", 1)
                        token.level = level
                        level += 1
                        nodes.append(token)

                        nodes.append(current_token)

                        token = Token("figcaption_open", "figcaption", 0)
                        token.level = level
                        level += 1
                        token.attrs = {"aria-hidden": "true"}
                        token.content = current_token.content
                        nodes.append(token)

                        token = Token("text", "", 0)
                        token.content = current_token.content
                        token.level = level
                        nodes.append(token)

                        token = Token("figcaption_close", "figcaption", 0)
                        token.level = level
                        level -= 1
                        nodes.append(token)

                        token = Token("figure_close", "figure", -1)
                        level -= 1
                        token.level = level

                        nodes.append(token)

                        inline_token.children = tokens = arrayReplaceAt(tokens, i, nodes)

    md.core.ruler.after(
        "inline",
        "img-figures",
        partial(_img_figures, lazyload=lazyload),
    )
