"""Tighten Camelot cell rectangles onto the words each cell prints.

Camelot reports each cell as structural grid geometry: the column's x-extent
crossed with the row band. On dense filings that geometry hugs the print, but
on sparse layouts the separators land mid-whitespace, so a cell's rectangle
runs far wider than its value and can sit half a line off the glyphs — a
highlight drawn from it cuts through the number instead of surrounding it.

The PDF itself records the rectangle of every word it prints (the text
layer). A cell whose text matches words near its structural rectangle takes
the extent of those words as its box (`tighten_box`). That box can reach past
the rectangle, because the words are gathered from a margin around it: a row
band can sit off its glyph line, and a glued `$` can print in the next
column.

A cell whose text matches no nearby words is clamped to the words whose
centers sit inside its rectangle (`ink_extent`). The clamp only shrinks the
rectangle. A cell with no box or no text keeps the structural rectangle, and
so does an unmatched cell with no word inside it, such as a cell of an
image-based table with no text layer. No box is ever invented.

Words arrive as `quber.agents.completeness.page_words` returns them:
(x0, y0, x1, y1, text) in PDF points, top-left origin. They are read from the
same text layer Camelot parses, through PyMuPDF, whose word splits can differ
from Camelot's. `tighten_box` works in that frame. `tighten_cell_boxes`
accepts and returns Camelot's frame — PDF points, bottom-left origin — so the
acquisition layer stores exactly what Camelot's contract promises, only
tighter.
"""

from __future__ import annotations

from typing import List, Optional, Sequence, Tuple

Rect = Tuple[float, float, float, float]  # x1, y1, x2, y2
Word = Tuple[float, float, float, float, str]

# The least horizontal slack, in points, when gathering candidate words near a
# cell: separators can land a hair inside the print. `_candidates` widens it
# to the row band's height when that is larger.
PAD_X = 4.0
# Height in points of the bands word centers are rounded into when sorting
# into reading order. Words whose centers fall in one band read as one line.
LINE_QUANT = 4.0

# Typographic variants that make identical text read as different strings:
# curly quotes against straight ones, the dash family against the hyphen.
# The cell text and the text layer can disagree on these, so both sides
# normalize before comparing.
PUNCT_VARIANTS = str.maketrans(
    {"‘": "'", "’": "'", "“": '"', "”": '"', "‐": "-", "‑": "-", "‒": "-", "–": "-", "—": "-", "−": "-"}
)


def squeeze(text: str) -> str:
    """Comparison form: whitespace removed, typographic punctuation folded."""
    return "".join(text.split()).translate(PUNCT_VARIANTS)


def _union(run: Sequence[Word]) -> Rect:
    return (
        min(w[0] for w in run),
        min(w[1] for w in run),
        max(w[2] for w in run),
        max(w[3] for w in run),
    )


def _candidates(box: Rect, words: Sequence[Word]) -> List[Tuple[Word, str]]:
    """Words near the structural box, in reading order, with comparison text.

    The margin is one box height vertically (camelot's row band drifts off
    the glyph line) and one box height horizontally, never less than `PAD_X`
    (a glued '$' can print in the neighboring column's band).
    """
    x1, y1, x2, y2 = box
    band = y2 - y1
    margin_x = max(PAD_X, band)
    cand = [
        (w, t)
        for w, t in ((w, squeeze(w[4])) for w in words)
        if t
        and x1 - margin_x <= (w[0] + w[2]) / 2 <= x2 + margin_x
        and y1 - band <= (w[1] + w[3]) / 2 <= y2 + band
    ]
    cand.sort(key=lambda wt: (round((wt[0][1] + wt[0][3]) / 2 / LINE_QUANT), wt[0][0]))
    return cand


def _token_runs(cand: List[Tuple[Word, str]], token: str) -> List[Tuple[int, int]]:
    """Every contiguous candidate run whose text equals `token`,
    as (first, last) index pairs."""
    out: List[Tuple[int, int]] = []
    for i in range(len(cand)):
        joined = ""
        for j in range(i, len(cand)):
            joined += cand[j][1]
            if len(joined) >= len(token):
                if joined == token:
                    out.append((i, j))
                break
    return out


def tighten_box(box: Rect, text: str, words: Sequence[Word]) -> Optional[Rect]:
    """The printed extent of `text` near the structural `box`, or None.

    Two passes over the nearby words. The first looks for the whole cell text
    as one contiguous run in reading order, which is the normal case. Of
    several matching runs, the one nearest the box center wins. Failing that,
    and when the text has at least two whitespace-separated tokens, each token
    is matched as its own contiguous run and takes the run nearest the box
    center, in any order. That grounds glued cells whose fragments print apart
    ('$' left-aligned, its value right-aligned) and labels whose superscript
    markers sort off their line. The result is the union of the matched words,
    so it can extend past `box`. No pass matching means None.
    """
    target = squeeze(text)
    if not target:
        return None
    cand = _candidates(box, words)

    runs: List[List[Word]] = []
    for i in range(len(cand)):
        joined = ""
        for j in range(i, len(cand)):
            joined += cand[j][1]
            if len(joined) >= len(target):
                if joined == target:
                    runs.append([wt[0] for wt in cand[i : j + 1]])
                break
    if runs:
        x1, y1, x2, y2 = box
        cx, cy = (x1 + x2) / 2, (y1 + y2) / 2

        def offset(run: Sequence[Word]) -> float:
            rx = (min(w[0] for w in run) + max(w[2] for w in run)) / 2
            ry = (min(w[1] for w in run) + max(w[3] for w in run)) / 2
            return (rx - cx) ** 2 + (ry - cy) ** 2

        return _union(min(runs, key=offset))

    tokens = [squeeze(t) for t in text.split()]
    tokens = [t for t in tokens if t]
    if len(tokens) < 2:
        return None
    # No ordering constraint between tokens: a superscript marker sorts off
    # its label's line, and a glued '$' prints a column away. Each token takes
    # its run nearest the box center; the candidate window bounds the damage
    # a repeated token could do.
    bx1, by1, bx2, by2 = box
    cx, cy = (bx1 + bx2) / 2, (by1 + by2) / 2
    picked: List[Word] = []
    for token in tokens:
        found = _token_runs(cand, token)
        if not found:
            return None

        def run_offset(ij: Tuple[int, int]) -> float:
            run = [wt[0] for wt in cand[ij[0] : ij[1] + 1]]
            rx = (min(w[0] for w in run) + max(w[2] for w in run)) / 2
            ry = (min(w[1] for w in run) + max(w[3] for w in run)) / 2
            return (rx - cx) ** 2 + (ry - cy) ** 2

        i, j = min(found, key=run_offset)
        picked.extend(wt[0] for wt in cand[i : j + 1])
    return _union(picked)


def ink_extent(box: Rect, words: Sequence[Word]) -> Optional[Rect]:
    """The extent of the words printed inside the structural `box`, or None.

    The last resort for a cell whose text has no match in the text layer. The
    cell cannot be grounded to specific words, but it should never claim more
    of the page than the ink it holds, so it shrinks to the words whose centers
    sit inside it, on both axes. A box holding no words, as in an image table,
    returns None.
    """
    x1, y1, x2, y2 = box
    inside = [
        w for w in words if x1 <= (w[0] + w[2]) / 2 <= x2 and y1 <= (w[1] + w[3]) / 2 <= y2 and w[4].strip()
    ]
    if not inside:
        return None
    ux1, uy1, ux2, uy2 = _union(inside)
    # Never grow: the clamp intersects the band, it does not replace it.
    return (max(ux1, x1), max(uy1, y1), min(ux2, x2), min(uy2, y2))


def tighten_cell_boxes(
    cells: Sequence[Sequence[str]],
    cell_boxes: List[List[Optional[Rect]]],
    words: Sequence[Word],
    page_height: float,
) -> List[List[Optional[Rect]]]:
    """Cell boxes moved onto the words each cell prints.

    `cell_boxes` arrive and return in Camelot's frame (PDF points,
    bottom-left origin), shaped like `cells`. Each cell takes `tighten_box`,
    falling back to `ink_extent`. A cell with no box or no text, or one both
    return None for, keeps its entry unchanged.
    """
    out: List[List[Optional[Rect]]] = []
    for i, row in enumerate(cell_boxes):
        out_row: List[Optional[Rect]] = []
        for j, box in enumerate(row):
            text = cells[i][j] if i < len(cells) and j < len(cells[i]) else ""
            if box is None or not text.strip():
                out_row.append(box)
                continue
            bx1, by1, bx2, by2 = box
            flipped = (bx1, page_height - by2, bx2, page_height - by1)
            tight = tighten_box(flipped, text, words) or ink_extent(flipped, words)
            if tight is None:
                out_row.append(box)
            else:
                tx1, ty1, tx2, ty2 = tight
                out_row.append((tx1, page_height - ty2, tx2, page_height - ty1))
        out.append(out_row)
    return out