import marko
import marko.block
import marko.inline
import re
from log.Logger import Logger

logger = Logger()

class BlockFormula(marko.block.BlockElement):
    pattern = re.compile(r"^\$\$\s*\n?([\s\S]+?)\n?\$\$\s*$", re.MULTILINE)

    def __init__(self, match):
        self.children = [marko.inline.RawText(match.group(1).strip())]

    @classmethod
    def match(cls, source):
        return source.expect_re(cls.pattern)

    @classmethod
    def parse(cls, source):
        match = source.match
        source.consume()
        return match

class InlineFormula(marko.inline.InlineElement):
    pattern = re.compile(r"\$(?!\$)([^\$]+?)\$")
    parse_children = False  # Math content is raw text

    def __init__(self, match):
        self.children = [marko.inline.RawText(match.group(1).strip())]

    @classmethod
    def match(cls, source):
        return source.expect_re(cls.pattern)

class Paragraph(marko.block.Paragraph):
    override = True

    @classmethod
    def break_paragraph(cls, source, lazy=False):
        if BlockFormula.match(source):
            return True
        return super().break_paragraph(source, lazy=lazy)

class Renderer:
    def render_block_formula(self, element):
        # MathJax compatible block math
        # logger.log_debug(f"render_block_formula@LaTeXRenderer.py returned => {element}")
        return f'\n<div class="math-block">$$\n{self.render_children(element)}\n$$</div>\n'

    def render_inline_formula(self, element):
        # MathJax compatible inline math
        # logger.log_debug(f"render_inline_formula@LaTeXRenderer.py returned => {element}")
        return f'\\({self.render_children(element)}\\)'

class LaTeXExtension:
    elements = [BlockFormula, InlineFormula, Paragraph]
    parser_mixins = []
    renderer_mixins = [Renderer]