import re

import mistune

mistune._pre_tags.append('latex')


class MathInlineGrammar(mistune.InlineGrammar):
    block_math = re.compile(r'^\$\$(.*?)\$\$|^\\\[(.*?)\\\]', re.DOTALL)
    math = re.compile(r'^~(.*?)~|^\\\((.*?)\\\)', re.DOTALL)
    text = re.compile(r'^[\s\S]+?(?=[\\<!\[_*`~$]|\\[\[(]|https?://| {2,}\n|$)')


class MathInlineLexer(mistune.InlineLexer):
    grammar_class = MathInlineGrammar

    def __init__(self, *args, **kwargs):
        self.default_rules = self.default_rules[:]
        self.inline_html_rules = self.default_rules
        self.default_rules.insert(self.default_rules.index('strikethrough') + 1, 'math')
        self.default_rules.insert(self.default_rules.index('strikethrough') + 1, 'block_math')
        super(MathInlineLexer, self).__init__(*args, **kwargs)

    def output_block_math(self, m):
        return self.renderer.block_math(m.group(1) or m.group(2))

    def output_math(self, m):
        return self.renderer.math(m.group(1) or m.group(2))

    def output_inline_html(self, m):
        tag = m.group(1)
        text = m.group(3)
        if self._parse_inline_html and text:
            if tag == 'a':
                self._in_link = True
                text = self.output(text)
                self._in_link = False
            else:
                text = self.output(text)
            extra = m.group(2) or ''
            html = '<%s%s>%s</%s>' % (tag, extra, text, tag)
        else:
            html = m.group(0)
        return self.renderer.inline_html(html)


class MathRenderer(mistune.Renderer):
    def __init__(self, *args, **kwargs):
        super(MathRenderer, self).__init__(*args, **kwargs)

    def block_math(self, math):
        return r'\[%s\]' % mistune.escape(math)

    def math(self, math):
        return r'\(%s\)' % mistune.escape(math)
