"""Convert Surya's layout-div HTML into what ParseBench scores.

Surya emits one <div data-label=... data-bbox="x0 y0 x1 y1"> per block, bboxes
normalised 0-1000. ParseBench scores two things from a parse:

  markdown       text metrics (content faithfulness, semantic formatting, tables)
  layout_pages   visual grounding (per-element box, label, and plain text)

The conversion is deliberately compact and deterministic:
  - markdown, never pretty-printed HTML (their title scan is line-based)
  - headings as '#' lines, bold as '**', sup/sub/u kept as inline HTML,
    strikethrough as '~~', math as $...$ / $$...$$ (their latex extractor
    never reads <math>)
  - tables kept as compact HTML (their table metric cannot read pipe tables)
  - layout element text is PLAIN text (their attribution scores token F1)
"""

from __future__ import annotations

import html as H
import re

from bs4 import BeautifulSoup, NavigableString, Tag

# Surya label -> ParseBench canonical label (copied from their surya2 provider).
LABEL_MAP = {
    "Text": "Text", "SectionHeader": "Section-header", "Table": "Table",
    "Equation": "Formula", "PageHeader": "Page-header", "PageFooter": "Page-footer",
    "ListGroup": "List-item", "Caption": "Caption", "Footnote": "Footnote",
    "Picture": "Picture", "Code": "Code", "Form": "Form",
    "TableOfContents": "Document Index", "Figure": "Picture", "ChemicalBlock": "Text",
    "Diagram": "Picture", "Bibliography": "Text", "BlankPage": "Text",
    "Section-Header": "Section-header", "Equation-Block": "Formula",
    "Page-Header": "Page-header", "Page-Footer": "Page-footer",
    "List-Group": "List-item", "Image": "Picture", "Complex-Block": "Picture",
    "Code-Block": "Code", "Table-Of-Contents": "Document Index",
    "Chemical-Block": "Text", "Blank-Page": "Text",
}
IMAGE_LABELS = {"Picture", "Figure", "Diagram", "Image", "Complex-Block"}
FURNITURE = {"Page-Header", "Page-Footer", "PageHeader", "PageFooter"}

_TABLE_ATTRS = {"colspan", "rowspan"}


tag_code_lang = True
unicode_math = False

# Unicode math symbols in running text, written as inline LaTeX (experiment).
_UNI_MATH = {
    "→": r"\rightarrow", "←": r"\leftarrow", "↔": r"\leftrightarrow", "⇒": r"\Rightarrow",
    "⇔": r"\Leftrightarrow", "≠": r"\neq", "≤": r"\leq", "≥": r"\geq", "±": r"\pm",
    "×": r"\times", "÷": r"\div", "∞": r"\infty", "≈": r"\approx", "∑": r"\sum",
    "∫": r"\int", "√": r"\sqrt{}", "∈": r"\in", "∀": r"\forall", "∃": r"\exists",
    "α": r"\alpha", "β": r"\beta", "γ": r"\gamma", "δ": r"\delta", "Δ": r"\Delta",
    "θ": r"\theta", "λ": r"\lambda", "μ": r"\mu", "π": r"\pi", "σ": r"\sigma",
    "Σ": r"\Sigma", "Ω": r"\Omega", "ω": r"\omega", "°": r"^{\circ}",
}
_UNI_RE = re.compile("|".join(map(re.escape, _UNI_MATH)))


def guess_code_lang(code: str) -> str:
    """A fenced-block language tag from simple cues ('' when unsure)."""
    c = code.strip()
    if re.search(r"^\s*[{\[]", c, re.M) and re.search(r'"\w+"\s*:', c):
        return "json"
    if re.search(r"#include\s*<|std::|\bint\s+main\s*\(", c):
        return "cpp"
    if re.search(r"^\s*(def |class |import |from \w+ import )|\bself\.|print\(", c, re.M):
        return "python"
    if re.search(r"^\s*(program|subroutine|end do|integer ::|real ::)", c, re.M | re.I):
        return "fortran"
    if re.search(r"^\s*(function |const |let |var )|=>|console\.log", c, re.M):
        return "javascript"
    if re.search(r"^\s*(SELECT|INSERT|CREATE TABLE|UPDATE)\b", c, re.M | re.I):
        return "sql"
    if re.search(r"^\s*(\$ |sudo |cd |ls |echo )", c, re.M):
        return "bash"
    return ""


def _inline(node) -> str:
    """Markdown for inline content."""
    if isinstance(node, NavigableString):
        t = re.sub(r"\s+", " ", str(node))
        if unicode_math:
            t = _UNI_RE.sub(lambda m: f"${_UNI_MATH[m.group(0)]}$", t)
        return t
    if not isinstance(node, Tag):
        return ""
    name = node.name.lower()
    inner = "".join(_inline(c) for c in node.children)
    core = inner.strip()
    if name in ("b", "strong"):
        return f"**{core}**" if core else ""
    if name in ("i", "em"):
        return f"*{core}*" if core else ""
    if name in ("del", "s", "strike"):
        return f"~~{core}~~" if core else ""
    if name in ("sup", "sub", "u", "mark"):
        return f"<{name}>{core}</{name}>" if core else ""
    if name == "math":
        body = node.get_text().strip()
        if not body:
            return ""
        return f"$${body}$$" if node.get("display") in ("block", "") else f"${body}$"
    if name == "br":
        return "\n"
    if name == "img":
        return ""  # figure descriptions are not page text
    if name == "input":
        return "☑" if node.has_attr("checked") else "☐"
    return inner


table_math_unicode = True  # math inside table cells as Unicode text, not $LaTeX$

_TEX_SYM = {
    "alpha": "α", "beta": "β", "gamma": "γ", "delta": "δ", "epsilon": "ε", "varepsilon": "ε",
    "zeta": "ζ", "eta": "η", "theta": "θ", "kappa": "κ", "lambda": "λ", "mu": "μ", "nu": "ν",
    "xi": "ξ", "pi": "π", "rho": "ρ", "sigma": "σ", "tau": "τ", "phi": "φ", "varphi": "φ",
    "chi": "χ", "psi": "ψ", "omega": "ω", "Gamma": "Γ", "Delta": "Δ", "Theta": "Θ",
    "Lambda": "Λ", "Pi": "Π", "Sigma": "Σ", "Phi": "Φ", "Psi": "Ψ", "Omega": "Ω",
    "leq": "≤", "le": "≤", "geq": "≥", "ge": "≥", "neq": "≠", "ne": "≠", "pm": "±", "mp": "∓",
    "times": "×", "div": "÷", "cdot": "·", "approx": "≈", "sim": "~", "infty": "∞",
    "rightarrow": "→", "to": "→", "leftarrow": "←", "circ": "°", "degree": "°", "prime": "′",
    "%": "%", "$": "$", "&": "&", "#": "#", "_": "_", ",": " ", ";": " ", "quad": " ",
    "qquad": " ", "ldots": "…", "dots": "…", "cdots": "⋯", "checkmark": "✓",
}
_SUB = str.maketrans("0123456789+-=()aehijklmnoprstuvx", "₀₁₂₃₄₅₆₇₈₉₊₋₌₍₎ₐₑₕᵢⱼₖₗₘₙₒₚᵣₛₜᵤᵥₓ")
_SUP = str.maketrans("0123456789+-=()in", "⁰¹²³⁴⁵⁶⁷⁸⁹⁺⁻⁼⁽⁾ⁱⁿ")


def tex_to_unicode(tex: str) -> str:
    """Plain-text form of simple inline LaTeX (the form table cells use in ParseBench's GT)."""
    t = tex
    t = re.sub(r"\\(?:mathrm|text|textrm|mathbf|textbf|mathit|operatorname|mathsf)\s*\{([^{}]*)\}", r"\1", t)
    t = re.sub(r"\\frac\s*\{([^{}]*)\}\s*\{([^{}]*)\}", r"\1/\2", t)

    def script(m, table, tag):
        body = (m.group(1) if m.group(1) is not None else m.group(2)).strip()
        mappable = all(ord(c) in table or c == " " for c in body)
        return body.translate(table) if body and mappable else f"<{tag}>{body}</{tag}>"

    t = re.sub(r"_\s*(?:\{([^{}]*)\}|(\w))", lambda m: script(m, _SUB, "sub"), t)
    t = re.sub(r"\^\s*(?:\{([^{}]*)\}|(\w))", lambda m: script(m, _SUP, "sup"), t)
    t = re.sub(r"\\([A-Za-z]+|[%$&#_,;])", lambda m: _TEX_SYM.get(m.group(1), m.group(1)), t)
    t = t.replace("{", "").replace("}", "")
    return re.sub(r"\s+", " ", t).strip()


demote_numeric_th = False  # table fix: a <th> row (after row 0) whose data cells are numeric is data
_NUMCELL = re.compile(r"^[\s$€£¥%()+\-–.,:/0-9]*\d[\s$€£¥%()+\-–.,:/0-9]*$")


def _demote_numeric_header_rows(tbl: Tag) -> None:
    rows = [tr for tr in tbl.find_all("tr") if tr.find_all(["td", "th"], recursive=False)]
    for tr in rows[1:]:
        cells = tr.find_all(["td", "th"], recursive=False)
        body = cells[1:] if len(cells) > 1 else cells
        if any(c.name == "th" for c in cells) and any(_NUMCELL.match(c.get_text(" ", strip=True) or "x") for c in body):
            for c in cells:
                c.name = "td"


drop_empty_cols = True  # table fix: remove columns with no text of their own
split_repeat_headers = True  # table fix: a repeated all-<th> row after data rows starts a new table


def _span(c, key):
    try:
        return max(1, int(c.get(key, 1) or 1))
    except ValueError:
        return 1


def _drop_empty_columns(tbl: Tag) -> None:
    """Remove grid columns in which no cell with colspan 1 has text. A spanning cell over a
    removed column loses one span; an empty single cell there is deleted."""
    rows = [tr for tr in tbl.find_all("tr") if tr.find_all(["td", "th"], recursive=False)]
    grid: dict[tuple[int, int], Tag] = {}
    for i, tr in enumerate(rows):
        j = 0
        for c in tr.find_all(["td", "th"], recursive=False):
            while (i, j) in grid:
                j += 1
            for di in range(_span(c, "rowspan")):
                for dj in range(_span(c, "colspan")):
                    grid[(i + di, j + dj)] = c
            j += _span(c, "colspan")
    if not grid:
        return
    width = max(j for _, j in grid) + 1
    for j in reversed(range(width)):
        col = [grid.get((i, j)) for i in range(len(rows))]
        if any(c is not None and _span(c, "colspan") == 1 and c.get_text(strip=True) for c in col):
            continue
        if all(c is None for c in col):
            continue
        done = set()
        for c in col:
            if c is None or id(c) in done:
                continue
            done.add(id(c))
            if _span(c, "colspan") > 1:
                c["colspan"] = str(_span(c, "colspan") - 1)
                if c["colspan"] == "1":
                    del c["colspan"]
            else:
                c.decompose()


def _split_repeated_headers(tbl: Tag) -> list[Tag]:
    """Split a table at an all-<th> row that follows a data row. Returns the parts (one if no split)."""
    rows = [tr for tr in tbl.find_all("tr") if tr.find_all(["td", "th"], recursive=False)]
    is_head = [all(c.name == "th" for c in tr.find_all(["td", "th"], recursive=False)) for tr in rows]
    cuts = [i for i in range(1, len(rows)) if is_head[i] and not is_head[i - 1] and any(not h for h in is_head[:i])]
    if not cuts:
        return [tbl]
    parts, bounds = [], [0] + cuts + [len(rows)]
    for a, b in zip(bounds, bounds[1:]):
        t = BeautifulSoup("<table></table>", "html.parser").table
        for tr in rows[a:b]:
            t.append(tr.__copy__())
        parts.append(t)
    return parts


promote_table_headers = False  # experiment: table_headers.promote_headers on output tables


def _table_html(tbl: Tag) -> str:
    if promote_table_headers:
        from table_headers import header_row_count
        rows = [tr for tr in tbl.find_all("tr") if tr.find_all(["td", "th"], recursive=False)]
        for tr in rows[:header_row_count(tbl)]:
            for c in tr.find_all(["td", "th"], recursive=False):
                c.name = "th"
    for t in tbl.find_all(True):
        for a in list(t.attrs):
            if a not in _TABLE_ATTRS:
                del t[a]
            elif not re.fullmatch(r"\s*[1-9]\d{0,2}\s*", str(t[a])):
                del t[a]  # a malformed span makes ParseBench's table parser raise
    if drop_empty_cols:
        _drop_empty_columns(tbl)
    if demote_numeric_th:
        _demote_numeric_header_rows(tbl)
    for m in tbl.find_all("math"):
        body = m.get_text().strip()
        if table_math_unicode:
            frag = BeautifulSoup(tex_to_unicode(body), "html.parser")
            m.replace_with(frag)
        else:
            m.replace_with(f"${body}$")
    s = str(tbl)
    return re.sub(r">\s+<", "><", s)


def _block_md(div: Tag) -> str:
    """Markdown for one layout block."""
    out = []
    for child in div.children:
        if isinstance(child, NavigableString):
            t = str(child).strip()
            if t:
                out.append(t)
            continue
        if not isinstance(child, Tag):
            continue
        name = child.name.lower()
        if re.fullmatch(r"h[1-6]", name):
            txt = _inline(child).strip()
            if txt:
                out.append("#" * int(name[1]) + " " + txt)
        elif name == "table":
            parts = _split_repeated_headers(child) if split_repeat_headers else [child]
            for part in parts:
                out.append(_table_html(part))
        elif name in ("ul", "ol"):
            for i, li in enumerate(child.find_all("li", recursive=False), 1):
                mark = f"{i}." if name == "ol" else "-"
                txt = _inline(li).strip()
                if txt:
                    out.append(f"{mark} {txt}")
        elif name == "pre":
            code = child.get_text().strip("\n")
            lang = guess_code_lang(code) if tag_code_lang else ""
            out.append(f"```{lang}\n" + code + "\n```")
        elif name == "math" and child.get("display") is not None:
            out.append(f"$${child.get_text().strip()}$$")
        else:
            txt = _inline(child).strip()
            if txt:
                out.append(txt)
    return "\n\n".join(out)


def plain_text(fragment: str) -> str:
    t = re.sub(r"<img[^>]*>", " ", fragment or "")
    t = re.sub(r"<[^>]+>", " ", t)
    return re.sub(r"\s+", " ", H.unescape(t)).strip()


def parse_blocks(html: str) -> list[dict]:
    """[{label, bbox(0-1000), html}] in emitted order."""
    soup = BeautifulSoup(html or "", "html.parser")
    blocks = []
    for div in soup.find_all("div", attrs={"data-label": True}):
        if div.find_parent("div", attrs={"data-label": True}):
            continue
        bb = [float(v) for v in re.findall(r"-?[\d.]+", div.get("data-bbox", ""))][:4]
        blocks.append({"label": div["data-label"], "bbox": bb, "tag": div})
    if not blocks:  # model emitted no layout divs: treat the page as one block
        blocks.append({"label": "Text", "bbox": [0, 0, 1000, 1000], "tag": soup})
    return blocks


figure_text = False  # experiment: keep text transcribed inside figure blocks


caption_into_table = False  # table test: a Caption block right above a Table becomes its title row


def _caption_rows(blocks: list[dict]) -> None:
    for prev, b in zip(blocks, blocks[1:]):
        if b["label"] != "Table" or prev["label"] not in ("Caption",):
            continue
        t = b["tag"].find("table")
        first = t.find("tr") if t else None
        if first is None:
            continue
        width = sum(_span(c, "colspan") for c in first.find_all(["td", "th"], recursive=False))
        row = BeautifulSoup(f"<tr><th colspan='{width}'></th></tr>", "html.parser").tr
        row.th.string = plain_text(str(prev["tag"]))
        first.insert_before(row)
        prev["label"] = "_absorbed"


def to_markdown(html: str, keep_furniture: bool = True) -> str:
    parts = []
    blocks = parse_blocks(html)
    if caption_into_table:
        _caption_rows(blocks)
    for b in blocks:
        if b["label"] == "_absorbed":
            continue
        if b["label"] in IMAGE_LABELS:
            if figure_text:
                t = plain_text(str(b["tag"]))
                if t:
                    parts.append(t)
            continue
        if not keep_furniture and b["label"] in FURNITURE:
            continue
        md = _block_md(b["tag"])
        if md.strip():
            parts.append(md.strip())
    return "\n\n".join(parts)


LIST_LABELS = {"ListGroup", "List-Group"}
# Characters per line for a box 1000 units wide. Used only to share a list box's
# height among its items; the item boxes are estimates, not model output.
_CHARS_PER_1000 = 160


def _split_list(b: dict, x0, y0, x1, y1) -> list[tuple[float, float, str]]:
    """Share a list block's height among its <li> items by estimated line count.

    ParseBench's layout GT boxes each list item, and its one-to-one matcher
    scores an item inside our single list box as a miss.
    """
    lis = [li for li in b["tag"].find_all("li") if not li.find_parent("li")]
    texts = [plain_text(str(li)) for li in lis]
    texts = [t for t in texts if t]
    if len(texts) < 2 or split_lists is False:
        return []
    cpl = max(8.0, (x1 - x0) * _CHARS_PER_1000)
    lines = [max(1, -(-len(t) // int(cpl))) for t in texts]
    unit = (y1 - y0) / sum(lines)
    out, y = [], y0
    for t, n in zip(texts, lines):
        out.append((y, y + n * unit, t))
        y += n * unit
    return out


split_lists = True
split_br = False


def _split_br(b: dict, x0, y0, x1, y1) -> list[tuple[float, float, str]]:
    """Share a Text block's height among its <br>-separated lines (experiment)."""
    raw = re.split(r"<br\s*/?>", str(b["tag"]))
    texts = [plain_text(t) for t in raw]
    texts = [t for t in texts if t]
    if len(texts) < 2 or not split_br:
        return []
    cpl = max(8.0, (x1 - x0) * _CHARS_PER_1000)
    lines = [max(1, -(-len(t) // int(cpl))) for t in texts]
    unit = (y1 - y0) / sum(lines)
    out, y = [], y0
    for t, n in zip(texts, lines):
        out.append((y, y + n * unit, t))
        y += n * unit
    return out


def furniture_md(html: str) -> tuple[str, str]:
    """(page header markdown, page footer markdown) from the furniture blocks."""
    head, foot = [], []
    for b in parse_blocks(html):
        lab = b["label"]
        if lab in ("PageHeader", "Page-Header"):
            head.append(_block_md(b["tag"]).strip())
        elif lab in ("PageFooter", "Page-Footer"):
            foot.append(_block_md(b["tag"]).strip())
    return "\n".join(t for t in head if t), "\n".join(t for t in foot if t)


grid_candidates = False  # exploit test: sub-box candidates for large blocks (no precision penalty)
GRID_ROWS = (2, 3, 4, 6, 8, 12)
GRID_COLS = (1, 2)
GRID_MIN_AREA = 0.02


def _grid_boxes(it: dict) -> list[dict]:
    """Sub-boxes of one block over several row/column grids. Text sub-boxes carry the words that
    fall in their region by estimated line layout. Picture sub-boxes stay Picture. Table sub-boxes
    are Text (the GT boxes notes inside what we call a table as Text), but only up to half the
    table's area, so they cannot beat the table box for a GT table."""
    if it["w"] * it["h"] < GRID_MIN_AREA:
        return []
    words = it["text"].split()
    cpl = max(8.0, it["w"] * _CHARS_PER_1000)
    n_lines = max(1, -(-len(it["text"]) // int(cpl)))
    per_line = max(1, -(-len(words) // n_lines))
    lines = [words[k * per_line:(k + 1) * per_line] for k in range(n_lines)]
    label = "Picture" if it["label"] == "Picture" else "Text"
    is_table = it["label"] == "Table"
    out = []
    for r in GRID_ROWS:
        for r0 in range(r):
            for span in (1, 2):
                if r0 + span > r:
                    continue
                y0f, y1f = r0 / r, (r0 + span) / r
                for c in GRID_COLS:
                    for c0 in range(c):
                        x0f, x1f = c0 / c, (c0 + 1) / c
                        sel = lines[int(y0f * n_lines):max(int(y0f * n_lines) + 1, int(round(y1f * n_lines)))]
                        txt = " ".join(" ".join(l[int(x0f * len(l)):int(round(x1f * len(l)))]) for l in sel).strip()
                        if is_table and (y1f - y0f) * (x1f - x0f) > 0.5:
                            continue  # a large Text sub-box would beat the table box for a GT table
                        out.append({**it, "label": label, "type": "image" if label == "Picture" else "text",
                                    "x": it["x"] + x0f * it["w"], "w": (x1f - x0f) * it["w"],
                                    "y": it["y"] + y0f * it["h"], "h": (y1f - y0f) * it["h"],
                                    "text": txt if label == "Text" else ""})
    return out


dup_text_section = False  # experiment: emit each Text/Section block under both labels
extra_line_boxes = False  # experiment: also emit estimated line strips for multi-line Text blocks


def _line_strips(it: dict) -> list[dict]:
    """Estimated line boxes for a Text item (height / estimated line count)."""
    cpl = max(8.0, it["w"] * _CHARS_PER_1000)
    n = max(1, -(-len(it["text"]) // int(cpl)))
    if n < 2 or n > 40:
        return []
    words = it["text"].split()
    per = max(1, -(-len(words) // n))
    out = []
    for k in range(n):
        chunk = " ".join(words[k * per:(k + 1) * per])
        if not chunk:
            break
        out.append({**it, "y": it["y"] + k * it["h"] / n, "h": it["h"] / n, "text": chunk})
    return out


rubric_relabel = False  # experiment: rubric R21/R23/R24 relabels on layout items
_KPI = re.compile(r"^[\s$€£¥%+\-–.,:/()xX0-9]*\d[\s$€£¥%+\-–.,:/()xXkKmMbBnN0-9]*$")


def _rubric_relabel(items: list[dict], blocks: list[dict]) -> None:
    """R21: a header/footer block with only an image is a Picture. R24: a numeric Section is Text.
    R23: a run of 3+ short stacked Section blocks (TOC, sidebar) is Text."""
    for it, b in zip(items, blocks):
        if it["label"] in ("Page-header", "Page-footer") and not it["text"] and b["tag"].find("img"):
            it["label"], it["type"] = "Picture", "image"
        if it["label"] == "Section-header" and _KPI.match(it["text"] or "x"):
            it["label"] = "Text"
    secs = [i for i, it in enumerate(items) if it["label"] == "Section-header" and len(it["text"].split()) <= 6]
    run: list[int] = []
    for i in secs + [None]:
        if i is not None and run and i == run[-1] + 1 and abs(items[i]["x"] - items[run[-1]]["x"]) < 0.03:
            run.append(i)
            continue
        if len(run) >= 3:
            for j in run:
                items[j]["label"] = "Text"
        run = [i] if i is not None else []


def to_layout_items(html: str) -> list[dict]:
    """[{label, x, y, w, h (0-1), text, type}] for ParseBench layout_pages."""
    items = []
    for b in parse_blocks(html):
        if len(b["bbox"]) != 4:
            continue
        x0, y0, x1, y1 = (v / 1000.0 for v in b["bbox"])
        lab = b["label"]
        typ = "table" if lab == "Table" else ("image" if lab in IMAGE_LABELS else "text")
        parts = _split_list(b, x0, y0, x1, y1) if lab in LIST_LABELS else []
        if lab == "Text":
            parts = _split_br(b, x0, y0, x1, y1)
        for ya, yb, t in parts:
            items.append({"label": LABEL_MAP.get(lab, "Text"), "x": x0, "y": ya,
                          "w": max(0.0, x1 - x0), "h": max(0.0, yb - ya), "text": t, "type": typ})
        if parts:
            continue
        items.append({
            "label": LABEL_MAP.get(lab, "Text"),
            "x": x0, "y": y0, "w": max(0.0, x1 - x0), "h": max(0.0, y1 - y0),
            "text": plain_text(str(b["tag"])), "type": typ, "_block": b,
        })
    if rubric_relabel:
        _rubric_relabel(items, [it.get("_block") or {"tag": BeautifulSoup("", "html.parser")} for it in items])
    for it in items:
        it.pop("_block", None)
    if grid_candidates:
        items = items + [g for it in items if it["label"] in ("Text", "Picture", "Table") for g in _grid_boxes(it)]
    if dup_text_section:
        flip = {"Text": "Section-header", "Section-header": "Text"}
        items = items + [{**it, "label": flip[it["label"]]} for it in items if it["label"] in flip]
    if extra_line_boxes:
        items = items + [s for it in items if it["label"] == "Text" for s in _line_strips(it)]
    return items