# Based on https://gist.github.com/mikefromit/5a6fdfecc9310712f15a872df9f41f03

from jinja2 import pass_environment, nodes
from jinja2.ext import Extension
import mistletoe as md
from mistletoe.contrib.mathjax import MathJaxRenderer

@pass_environment
def markdown(env, value):
    """
    Markdown filter with support for extensions.
    """
    output = value

    return md.markdown(output, MathJaxRenderer)


class Markdown(Extension):
    """
    A wrapper around the markdown filter for syntactic sugar.
    """

    tags = set(["markdown"])

    def parse(self, parser):  # pragma: no cover
        """
        Parses the statements and defers to the callback
        for markdown processing.
        """
        lineno = next(parser.stream).lineno
        body = parser.parse_statements(["name:endmarkdown"], drop_needle=True)

        return nodes.CallBlock(
            self.call_method("_render_markdown"), [], [], body
        ).set_lineno(lineno)

    def _render_markdown(self, caller=None):  # pragma: no cover
        """
        Calls the markdown filter to transform the output.
        """
        if not caller:
            return ""
        output = caller().strip()
        return markdown(self.environment, output)