"""Serve dot diagrams inline."""
import base64
import re

import markdown
import pydot


class InlinePydot(markdown.Extension):
    """Wrap the markdown prepcoressor."""

    def extendMarkdown(self, md, md_globals):
        """Add InlinePydotPreprocessor to the Markdown instance."""
        md.registerExtension(self)
        md.preprocessors.add('dot_block', InlinePydotPreprocessor(md), '_begin')


class InlinePydotPreprocessor(markdown.preprocessors.Preprocessor):
    """Identify dot codeblocks and run them through pydot."""

    BLOCK_RE = re.compile(r'~~~pydot:(?P<filename>[^\s]+)\n(?P<content>.*?)~~~', re.DOTALL)

    def run(self, lines):
        """Match and generate diagrams from dot code blocks."""
        text = '\n'.join(lines)
        for match in self.BLOCK_RE.finditer(text):
            filename = match.group(1)
            dot_string = match.group(2)

            # use pydot to turn the text into pydot
            graphs = pydot.graph_from_dot_data(dot_string)

            # encode the image and provide as an inline image in markdown
            encoded_image = base64.b64encode(graphs[0].create_png()).decode('ascii')
            data_path = f'data:image/png;base64,{encoded_image}'
            inline_image = f'![{filename}]({data_path})'

            # replace the image in the output markdown
            text = f'{text[:match.start()]}\n{inline_image}\n{text[match.end():]}'

        return text.split('\n')


def makeExtension(*args, **kwargs):
    """Provide the extension to the markdown extension loader."""
    return InlinePydot(*args, **kwargs)