"""Adds a new directive for using upstream's code annotations feature.
See comments in visit_literal_block() override below for detailed explanation about
why this directive is needed."""
from typing import List, cast
from docutils import nodes
from sphinx.application import Sphinx
from sphinx.util.docutils import SphinxDirective
from sphinx.writers.html import HTMLTranslator
from sphinx.util.logging import getLogger
from . import html_translator_mixin

LOGGER = getLogger(__name__)


class annotations_list(nodes.container):
    pass


class CodeAnnotations(SphinxDirective):
    """Take a list of annotations and move them into a previous node's parent."""

    has_content = True

    def run(self) -> List[nodes.Node]:
        """Run the directive."""
        self.assert_has_content()

        div = annotations_list("", used=False)
        self.state.nested_parse(self.content, self.content_offset, div)
        self.set_source_info(div)
        if not isinstance(div.children[0], nodes.enumerated_list):
            LOGGER.error(
                "The code-annotations directive only accepts an enumerated list.",
                location=div,
            )
        return [div]


def visit_annotations_list(self: HTMLTranslator, node: annotations_list):
    # the content should have been moved into the literal block before calling this
    if not node["used"]:  # see visit_literal_block() override below
        node.children[0].walkabout(self)
    raise nodes.SkipNode()


@html_translator_mixin.override
def visit_literal_block(
    self: html_translator_mixin.HTMLTranslatorMixin,
    node: nodes.literal_block,
    super_func: html_translator_mixin.BaseVisitCallback[nodes.literal_block],
):
    #  we need to be sure that the annotated_list is the next node
    annotation_parent = cast(nodes.Element, node)
    if isinstance(node.parent.children[0], nodes.caption):
        # literal blocks with a caption have added parent
        annotation_parent = node.parent
    next_element = annotation_parent.next_node(
        nodes.Element, siblings=True, descend=False
    )
    if isinstance(next_element, annotations_list):
        try:
            super_func(self, node)
        except nodes.SkipNode:
            pass  # we're not done yet
        # the last line appended to self.body should now have the entire code-block:
        # <div class="highlight-xxxxx notranslate">
        #     <div class="highlight">
        #         <pre><span></span><code>...</code></pre>
        #     </div>
        #     <!-- we need to insert annotations' list here -->
        # </div>
        assert (  # to try detecting future changes in sphinx/pygments
            self.body[-1].endswith("</div>\n")
            and '<div class="highlight">' in self.body[-1]
        )
        # Add annotate class to the div.highlight tag (generated by pygments), and
        # remove the closing div tag (`[:-7]`) of the div.highlight's parent node.
        self.body[-1] = self.body[-1][:-7].replace(
            '<div class="highlight">', '<div class="annotate highlight">'
        )
        next_element.children[0].walkabout(self)  # add annotations' enumerated_list
        self.body.append("</div>\n")  # re-add the closing div for the parent
        next_element["used"] = True
        raise nodes.SkipNode()  # now we're done
    # else no annotations list found; proceed like normal
    super_func(self, node)


def setup(app: Sphinx):
    app.add_directive("code-annotations", CodeAnnotations)
    app.add_node(annotations_list, html=(visit_annotations_list, None))
    return {
        "parallel_read_safe": True,
        "parallel_write_safe": True,
    }
