"""In-place tightening of stored table-cell grounding boxes.

Documents already in the library were extracted before cell rectangles were
tightened at acquisition, so their stored boxes are still structural grid
geometry — the column's x-extent crossed with the row band. Re-extracting
would cost LLM passes and re-ingesting would wipe curation, so this backfill
corrects the one thing that is wrong: every tableCell grounding with a box
and text has its bbox moved onto the staged PDF's text layer. A cell whose
text matches nearby words takes those words' printed extent. A cell whose
text does not match shrinks to the words printed inside its box, and counts
as tightened too. Only a cell with no words inside its box keeps its stored
box. No other column is touched.

Run:
  uv run python -m quber.playground.backfill_boxes [doc_key ...]

With no arguments every document in the library is processed. A document
whose staged PDF is missing is reported and skipped, never silently.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List, Tuple

import psycopg
from psycopg.types.json import Jsonb

from quber.agents.completeness import page_words
from quber.core.extractors.camelot.tighten import ink_extent, tighten_box
from quber.playground import db
from quber.settings import get_settings


def backfill_document(conn: psycopg.Connection, doc_id: int, pdf_path: Path) -> Tuple[int, int]:
    """Tighten one document's cell boxes in place. Returns (tightened, kept)."""
    rows = conn.execute(
        """SELECT ref_id, page, bbox, cell_text FROM ade_playground.groundings
           WHERE document_id = %s AND ref_type = 'tableCell'
             AND bbox IS NOT NULL AND coalesce(cell_text, '') <> ''""",
        (doc_id,),
    ).fetchall()
    by_page: Dict[int, List[Tuple[Any, ...]]] = {}
    for r in rows:
        by_page.setdefault(r[1], []).append(r)

    tightened = kept = 0
    for page0, page_rows in by_page.items():
        try:
            w, h, words = page_words(pdf_path, page0 + 1)
        except IndexError:
            kept += len(page_rows)
            continue
        for ref_id, _, bbox, cell_text in page_rows:
            box = bbox if isinstance(bbox, dict) else json.loads(bbox)
            points = (box["left"] * w, box["top"] * h, box["right"] * w, box["bottom"] * h)
            tight = tighten_box(points, cell_text, words) or ink_extent(points, words)
            if tight is None:
                kept += 1
                continue
            x1, y1, x2, y2 = tight
            conn.execute(
                """UPDATE ade_playground.groundings SET bbox = %s
                   WHERE document_id = %s AND ref_id = %s""",
                (
                    Jsonb({"left": x1 / w, "top": y1 / h, "right": x2 / w, "bottom": y2 / h}),
                    doc_id,
                    ref_id,
                ),
            )
            tightened += 1
    return tightened, kept


def main() -> int:
    parser = argparse.ArgumentParser(description="Tighten stored table-cell grounding boxes in place")
    parser.add_argument("doc_keys", nargs="*", help="doc_key values to backfill; default all")
    args = parser.parse_args()

    data_dir = get_settings().playground.data_dir
    total_tightened = total_kept = 0
    missing: List[str] = []
    with db.connect() as conn:
        docs = conn.execute(
            "SELECT id, doc_key, coalesce(title, filename) FROM ade_playground.documents ORDER BY id"
        ).fetchall()
        for doc_id, doc_key, label in docs:
            if args.doc_keys and doc_key not in args.doc_keys:
                continue
            pdf_path = data_dir / f"{doc_key}.pdf"
            if not pdf_path.exists():
                missing.append(f"{doc_key} ({label})")
                continue
            tightened, kept = backfill_document(conn, doc_id, pdf_path)
            total_tightened += tightened
            total_kept += kept
            print(f"  {doc_key}  tightened {tightened:>5}  kept {kept:>4}  {label}")

    print(f"\ntotal tightened {total_tightened}, kept {total_kept}")
    if missing:
        print(f"\nSKIPPED — staged PDF missing for {len(missing)} document(s):")
        for m in missing:
            print(f"  {m}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())