from docutils import nodes
from sphinx.application import Sphinx
from sphinx.util.docutils import SphinxDirective
from sphinx.writers.html import HTMLTranslator


class annotations_list(nodes.container):
    pass


class CodeAnnotations(SphinxDirective):
    has_content = True

    def run(self) -> list[nodes.Node]:
        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):
            raise self.error(
                "The code-annotations directive only accepts an enumerated list."
            )
        return [div]


def visit_annotations_list(self: HTMLTranslator, node: annotations_list):
    code_block = node.previous_sibling()
    in_container = False
    if isinstance(code_block, nodes.container) and (
        i := code_block.first_child_matching_class(nodes.literal_block)
    ):
        code_block = code_block.children[i]
        in_container = True
    if not isinstance(code_block, nodes.literal_block):
        raise nodes.SkipNode()

    if in_container:
        assert self.body[-1].endswith("</div>\n")
        self.body.pop()
        while self.body and (not self.body[-1] or self.body[-1].isspace()):
            self.body.pop()

    # the last line appended to self.body should now have the entire code-block:
    # <div class="highlight-xxxxx notranslate">
    #     <div class="annotate highlight">
    #         <pre><span>... (1) ... (2) ...</span></pre>
    #     </div>
    #     <ol>
    #         <li>Annotation 1</li>
    #         <li>Annotation 2</li>
    #     </ol>
    # </div>

    try:
        # to try detecting future changes in sphinx/pygments
        assert self.body[-1].endswith("</div>\n")
        assert '<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">'
        )
        node.children[0].walkabout(self)  # add annotations' enumerated_list
        self.body.append("</div>\n")
    finally:
        if in_container:
            self.body.append("</div>\n")
    raise nodes.SkipNode()


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,
    }
