import re
import os


def preprocess(input_file, output_file):
    with open(input_file, "r", encoding="utf-8") as f:
        content = f.read()

    # 1. Citations
    # <d-cite key="KEY"></d-cite> -> [@KEY]
    def replace_citation(match):
        key = match.group(1)
        return f"[@{key}]"

    content = re.sub(r'<d-cite key="([^"]+)"></d-cite>', replace_citation, content)

    # 2. Liquid Include Figures
    # {% include figure.html path="PATH" ... %} -> ![](PATH)
    # Also handle gif -> png replacement
    def replace_liquid_figure(match):
        path = match.group(1)
        # Check if gif
        if path.endswith(".gif"):
            png_path = path.replace(".gif", ".png")
            # Check if PNG exists in img folder instead of gif folder
            # Assuming path is like assets/gif/submission/trajectory.gif
            # We want assets/img/submission/trajectory.png
            possible_png_path = png_path.replace("/gif/", "/img/")
            if os.path.exists(possible_png_path):
                path = possible_png_path
            elif os.path.exists(png_path):
                path = png_path

        return f"![]({path})"

    # Updated regex to handle % inside style attributes
    content = re.sub(
        r'{% include figure.html path="([^"]+)"(.*?)%}', replace_liquid_figure, content
    )

    # 3. HTML Figure Blocks
    # Matches <figure>...</figure> and extracts src and caption.
    # We use DOTALL to match across lines.
    # We handle both iframe and img tags inside figure.

    def replace_html_block(match):
        block = match.group(0)

        # Extract src
        # Look for src="{{ 'PATH' | relative_url }}"
        src_match = re.search(r'src="\{\{ \'([^\']+)\' \| relative_url \}\}"', block)
        if not src_match:
            # Fallback: look for simple src="PATH"
            src_match = re.search(r'src="([^"]+)"', block)

        if not src_match:
            return block  # Could not find src, return original

        path = src_match.group(1)

        # Handle SVG -> PNG replacement
        if path.endswith(".svg"):
            # Specific override for mermaid diagram
            if "mermaid-diagram.svg" in path:
                specific_png = path.replace(
                    "mermaid-diagram.svg",
                    "mermaid-diagram.png",
                )
                if os.path.exists(specific_png):
                    path = specific_png

            # Generic fallback if still svg
            if path.endswith(".svg"):
                png_path = path.replace(".svg", ".png")
                if os.path.exists(png_path):
                    path = png_path

        # Extract caption
        caption = ""
        cap_match = re.search(r"<figcaption>(.*?)</figcaption>", block, re.DOTALL)
        if cap_match:
            caption = cap_match.group(1).strip()
        else:
            # Try alt attribute
            alt_match = re.search(r'alt="([^"]+)"', block)
            if alt_match:
                caption = alt_match.group(1)

        # Construct markdown image
        return f"![{caption}]({path})"

    content = re.sub(
        r"<figure[^>]*>.*?</figure>", replace_html_block, content, flags=re.DOTALL
    )

    # 4. Fix Math Nesting
    # Remove $$ around \begin{align*} or \begin{aligned}
    # Pandoc converts $$ to \[ \] which is invalid around align environments

    content = re.sub(r"\$\$\s*(\\begin\{align\*?\})", r"\1", content)
    content = re.sub(r"(\\end\{align\*?\})\s*\$\$", r"\1", content)

    # We DO NOT remove wrappers for 'aligned' as it requires math mode ($...$)
    # content = re.sub(r'\$\$\s*(\\begin\{aligned\})', r'\1', content)
    # content = re.sub(r'(\\end\{aligned\})\s*\$\$', r'\1', content)

    # 5. Remove Front Matter?
    # Pandoc usually handles YAML front matter, but "layout: distill" etc might cause issues or be ignored.
    # We keep it as is, Pandoc usually ignores unknown fields.

    with open(output_file, "w", encoding="utf-8") as f:
        f.write(content)


if __name__ == "__main__":
    preprocess("submission.md", "submission_processed.md")
