from markdown.extensions import Extension
from markdown.treeprocessors import Treeprocessor
import xml.etree.ElementTree as etree

class GalleryTreeprocessor(Treeprocessor):
    def run(self, doc):
        for parent in doc.findall(".//ul/.."):
            for idx, ul in enumerate(parent):
                if ul.tag != "ul" or len(ul) == 0:
                    continue
                has_gallery = False

                # Option 1: In raw text in the <li>
                if ul[-1].text and ul[-1].text.endswith("{gallery}"):
                    has_gallery = True
                    ul[-1].text = ul[-1].text[:-9]
                # Option 2: After the last child (in its tail) \
                if len(ul[-1]) and ul[-1][-1].tail and \
                        ul[-1][-1].tail.endswith("{gallery}"):
                    has_gallery = True
                    ul[-1][-1].tail = ul[-1][-1].tail[:-9]

                if has_gallery:
                    # TODO: Manipulate the output tree
                    el = etree.Element("div")
                    p = etree.Element("p")
                    p.text = "<There is a gallery here:>"
                    parent.remove(ul)
                    el.append(p)
                    el.append(ul)
                    parent.insert(idx, el)

class GalleryExtension(Extension):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def extendMarkdown(self, md):
        md.treeprocessors.register(GalleryTreeprocessor(md), 'gallery', 8)
        md.registerExtension(self)