"""Render MinerU's markup back into readable prose for `Chunk.text`.

Why this exists — measured, not assumed. Running the same extraction pipeline,
same document, same gold set, changing only the parse:

    PyMuPDF (plain text)                recall 0.8537   35/41
    MinerU, raw markup in text          recall 0.7561   31/41
    MinerU, formula + table rendered     recall 0.8293   34/41

The term filter is an NER model reading prose. MinerU writes formulas with every
character spaced out — `{ \\mathrm { P u r c h a s i n g ~ c o s t s } }` — and
tables as HTML. Neither yields a single mention, even for terms the same model
finds easily in plain text.

The markup is not discarded: `Chunk.latex` and `Chunk.table_html` keep it
verbatim, because the formula branch needs exactly that form.
"""

from __future__ import annotations

import html
import re

# --- LaTeX ---------------------------------------------------------------

_DROP_COMMANDS = re.compile(
    r"\\(?:qquad|quad|left|right|displaystyle|limits|nolimits)\b|\\[!,;:]"
)
_WRAPPERS = re.compile(r"\\(?:mathrm|mathbf|mathit|mathsf|text|textrm|operatorname)\s*")
# Operators map to non-letter glyphs on purpose. A letter here (e.g. "x" for
# \times) would itself look like a spelled-out single character and get merged
# into the token beside it — "C_p x T" becoming "C_p xT".
_SYMBOLS = {
    r"\times": "×", r"\cdot": "·", r"\div": "/", r"\pm": "±",
    r"\leq": "≤", r"\geq": "≥", r"\neq": "≠", r"\approx": "≈",
    r"\%": "%", r"\$": "$", r"\&": "&",
}

# "P u r c h a s i n g" -> "Purchasing", "1 0 0" -> "100". Runs of >= 2 single
# characters; letters and digits are joined separately so "C 5" is left alone.
# Subscripts are folded first (below), so a real variable like `C_p` is already
# one token and never gets swallowed into a neighbouring word.
_SPACED_RUN = re.compile(
    r"(?<!\S)(?:[A-Za-z]\s+){1,}[A-Za-z](?!\S)"
    r"|(?<!\S)(?:\d\s+){1,}\d(?!\S)"
)

# `~` is a non-breaking SPACE in LaTeX, i.e. a real word boundary. It has to
# survive the whitespace pass, otherwise "c o s t s" after it merges into the
# preceding word and "Purchasing costs" becomes "Purchasingcosts".
_WORD_GAP = "\x00"


def render_latex(latex: str) -> str:
    """LaTeX -> a readable line. Best effort: it feeds a term filter, not a parser."""
    s = latex.replace("$$", " ").replace("$", " ")
    s = s.replace("~", _WORD_GAP)

    # array/matrix scaffolding carries no meaning for a term filter
    s = re.sub(r"\\(?:begin|end)\s*\{[^{}]*\}", " ", s)
    s = re.sub(r"\{\s*(?:[rlc|]\s*){1,8}\}", " ", s)   # column spec, e.g. { r l }
    s = s.replace("\\\\", " ; ").replace("&", " ")

    # fold sub/superscripts first, so `C _ { p }` becomes one token `C_p`
    s = re.sub(r"\s*_\s*\{\s*([^{}]*?)\s*\}", r"_\1", s)
    s = re.sub(r"\s*\^\s*\{\s*([^{}]*?)\s*\}", r"^\1", s)
    s = re.sub(r"\s*_\s*([A-Za-z0-9])", r"_\1", s)
    s = re.sub(r"\s*\^\s*([A-Za-z0-9])", r"^\1", s)

    # \frac{a}{b} -> (a) / (b); repeat for nesting
    for _ in range(4):
        baru = re.sub(r"\\d?frac\s*\{([^{}]*)\}\s*\{([^{}]*)\}", r"(\1) / (\2)", s)
        if baru == s:
            break
        s = baru

    s = re.sub(r"\\sqrt\s*\{([^{}]*)\}", r"sqrt(\1)", s)
    for k, v in _SYMBOLS.items():
        s = s.replace(k, f" {v} ")
    s = _WRAPPERS.sub(" ", s)
    s = _DROP_COMMANDS.sub(" ", s)
    s = re.sub(r"\\[A-Za-z]+", " ", s)     # any command left over
    s = s.replace("{", " ").replace("}", " ").replace("\\", " ")
    s = re.sub(r"[ \t\r\n]+", " ", s).strip()

    # now that spacing is uniform, rejoin the spelled-out words and numbers
    s = _SPACED_RUN.sub(lambda m: m.group(0).replace(" ", ""), s)

    s = s.replace(_WORD_GAP, " ")
    return re.sub(r"\s+", " ", s).strip()


# --- HTML tables ---------------------------------------------------------

_CELL_END = re.compile(r"</\s*(?:td|th)\s*>", re.I)
_ROW_END = re.compile(r"</\s*tr\s*>", re.I)
_TAG = re.compile(r"<[^>]+>")

# Table cells routinely carry inline LaTeX ("$\frac{kVA}{...}$"). Stripping the
# HTML alone leaves that markup sitting in the prose, which is the same problem
# the formula rendering exists to solve — found in a real table on the first run.
_INLINE_MATH = re.compile(r"\$([^$]{1,400})\$")


def render_table(table_html: str) -> str:
    """HTML table -> plain rows, cells separated by ' | '.

    Keeps the reading order a person would use, so a term appearing only in a
    table header is still a mention.
    """
    s = _ROW_END.sub("\n", table_html)
    s = _CELL_END.sub(" | ", s)
    s = _TAG.sub(" ", s)
    s = html.unescape(s)
    s = _INLINE_MATH.sub(lambda m: " " + render_latex(m.group(1)) + " ", s)

    baris = []
    for raw in s.split("\n"):
        bersih = re.sub(r"[ \t]+", " ", raw).strip().strip("|").strip()
        if bersih:
            baris.append(re.sub(r"\s*\|\s*", " | ", bersih))
    return "\n".join(baris)