# latex_extension.py
import marko
import marko.block
import marko.inline
from marko.md_renderer import MarkdownRenderer
import re
from log.Logger import Logger
logger = Logger()

class BlockFormula(marko.block.BlockElement):
    pattern = re.compile(r"\$\$ *\n([\s\S]+?)^\$\$ *$", re.MULTILINE)
    def __init__(self, match):
        logger.log_debug("Did shit at __init__ for blockformula")
        self.children = [marko.inline.RawText(match.group(1))]
    @classmethod
    def match(cls, source):
        return source.expect_re(cls.pattern)
    @classmethod
    def parse(cls, source):
        logger.log_debug("Did some shit with Latex")
        match = source.match
        source.consume()
        return match

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):
        # Render as HTML with MathJax-compatible format
        return '\n<div class="math-block">$$\n' + self.render_children(element) + '$$</div>\n'

class LaTeXExtension:
    logger.log_debug("Did shit at __init__ for latexextension")
    elements = [BlockFormula, Paragraph]
    parser_mixins = []
    renderer_mixins = [Renderer]