"""
LaTeX formula conversion module.
Converts LaTeX formulas to Word native formulas (OMML format).
Conversion flow: LaTeX → MathML → OMML
"""

from __future__ import annotations

import html.entities
import re
import uuid
from typing import NamedTuple

import latex2mathml.converter
import mathml2omml
from lxml import etree  # type: ignore[import-untyped]


class FormulaInfo(NamedTuple):
    """Formula information."""

    placeholder: str
    latex: str
    is_block: bool  # True for block formula, False for inline formula


def extract_latex_formulas(text: str) -> tuple[str, list[FormulaInfo]]:
    """
    Extract LaTeX formulas from text and replace with placeholders.

    Args:
        text: Text containing LaTeX formulas

    Returns:
        (processed text, list of formula info)
    """
    formulas: list[FormulaInfo] = []

    # Process block formulas $$...$$ first (to avoid matching by inline pattern)
    block_pattern = r"[$][$](.+?)[$][$]"

    def replace_block(match):
        latex = match.group(1).strip()
        placeholder = f"FORMULABLOCK{uuid.uuid4().hex}"
        formulas.append(FormulaInfo(placeholder, latex, is_block=True))
        return placeholder

    text = re.sub(block_pattern, replace_block, text, flags=re.DOTALL)

    # Process inline formulas $...$
    inline_pattern = r"(?<![$])[$](?![$])(.+?)(?<![$])[$](?![$])"

    def replace_inline(match):
        latex = match.group(1).strip()
        placeholder = f"FORMULAINLINE{uuid.uuid4().hex}"
        formulas.append(FormulaInfo(placeholder, latex, is_block=False))
        return placeholder

    text = re.sub(inline_pattern, replace_inline, text, flags=re.DOTALL)

    return text, formulas


def latex_to_omml(latex: str) -> str | None:
    """
    Convert LaTeX formula to OMML format.

    Args:
        latex: LaTeX formula string

    Returns:
        OMML XML string, or None if conversion fails
    """
    try:
        # LaTeX → MathML
        mathml = latex2mathml.converter.convert(latex)

        # Fix unescaped special characters in MathML
        mathml = _fix_mathml_special_chars(mathml)

        # MathML → OMML
        omml = mathml2omml.convert(mathml, html.entities.name2codepoint)

        return omml
    except Exception as e:
        print(f"[WARN] LaTeX formula conversion failed: {latex[:50]}... - {e}")
        return None


def _fix_mathml_special_chars(mathml: str) -> str:
    """Fix unescaped special characters in MathML."""

    def fix_element_content(match):
        tag_name = match.group(1)
        attrs = match.group(2) or ""
        content = match.group(3)

        # Protect already escaped entities
        content = content.replace("&lt;", "\x00LT\x00")
        content = content.replace("&gt;", "\x00GT\x00")

        # Escape bare < and >
        content = content.replace("<", "&lt;")
        content = content.replace(">", "&gt;")

        # Restore escaped entities
        content = content.replace("\x00LT\x00", "&lt;")
        content = content.replace("\x00GT\x00", "&gt;")

        return f"<{tag_name}{attrs}>{content}</{tag_name}>"

    mathml = re.sub(
        r"<(mtext)(\s[^>]*)?>(.+?)</mtext>",
        fix_element_content,
        mathml,
        flags=re.DOTALL,
    )

    return mathml


def _fix_omml_groupchr(omml_str: str) -> str:
    """Fix unclosed groupChrPr tag in OMML generated by mathml2omml."""
    pattern = r"(<m:groupChrPr>(?:<[^>]+/>)*)(</m:groupChr>)"
    omml_str = re.sub(pattern, r"\1</m:groupChrPr>", omml_str)
    return omml_str


def _create_omml_element(omml_str: str):
    """
    Parse OMML string to lxml Element.

    Args:
        omml_str: OMML XML string

    Returns:
        lxml Element object, or None if parsing fails
    """
    try:
        omml_str = _fix_omml_groupchr(omml_str)

        omml_ns = "http://schemas.openxmlformats.org/officeDocument/2006/math"

        if omml_str.startswith("<m:"):
            first_tag_end = omml_str.find(">")
            if first_tag_end > 0:
                omml_str = omml_str[:first_tag_end] + f' xmlns:m="{omml_ns}"' + omml_str[first_tag_end:]

        omml_elem = etree.fromstring(omml_str.encode("utf-8"))

        return omml_elem
    except Exception as e:
        print(f"[WARN] OMML parsing failed: {e}")
        return None


def replace_formula_placeholders(document, formulas: list[FormulaInfo]) -> None:
    """
    Replace formula placeholders with OMML formulas in Word document.

    Args:
        document: python-docx Document object
        formulas: List of formula info
    """
    if not formulas:
        return

    formula_map = {f.placeholder: f for f in formulas}

    for paragraph in document.paragraphs:
        _replace_in_paragraph(paragraph, formula_map)

    for table in document.tables:
        for row in table.rows:
            for cell in row.cells:
                for paragraph in cell.paragraphs:
                    _replace_in_paragraph(paragraph, formula_map)


def _replace_in_paragraph(paragraph, formula_map: dict[str, FormulaInfo]) -> None:
    """Replace formula placeholders in a paragraph."""
    full_text = paragraph.text

    if "FORMULA" not in full_text:
        return

    placeholder_pattern = re.compile(r"FORMULA(?:INLINE|BLOCK)[0-9a-f]{32}")
    matches = [m for m in placeholder_pattern.finditer(full_text) if m.group(0) in formula_map]
    if not matches:
        return

    from docx.oxml import OxmlElement as OE
    from docx.oxml.ns import qn as qn_func

    def append_text_run(p, text: str) -> None:
        if not text:
            return
        r = OE("w:r")
        t = OE("w:t")
        t.text = text
        t.set(qn_func("xml:space"), "preserve")
        r.append(t)
        p.append(r)

    p = paragraph._p
    for child in list(p):
        if child.tag != qn_func("w:pPr"):
            p.remove(child)

    last_idx = 0
    for match in matches:
        placeholder = match.group(0)

        before_text = full_text[last_idx : match.start()]
        append_text_run(p, before_text)

        formula_info = formula_map[placeholder]

        omml_str = latex_to_omml(formula_info.latex)
        if omml_str is None:
            append_text_run(p, f"[{formula_info.latex}]")
            last_idx = match.end()
            continue

        omml_elem = _create_omml_element(omml_str)
        if omml_elem is None:
            append_text_run(p, f"[{formula_info.latex}]")
            last_idx = match.end()
            continue

        p.append(omml_elem)
        last_idx = match.end()

        print(f"[INFO] Replaced formula: {formula_info.latex[:30]}...")

    tail_text = full_text[last_idx:]
    append_text_run(p, tail_text)
