"""Diagnose a page's located-table geometry, stage by stage.

Runs the real grid locator on one page and prints, per flagged table, the
model's grid flags, whether it is treated as row-sharing (which confines the
tighten to the flagged columns), and the resulting left/right edges BEFORE
declash runs. Distinguishes a horizontal stub-column clip (a tighten/shares
problem) from anything declash does (which only moves top/bottom).

Usage: uv run python experiments/locate_diag.py <pdf-or-s3> <page>
"""

from __future__ import annotations

import asyncio
import sys
import tempfile
from pathlib import Path

import fitz

from quber.agents.completeness import page_words
from quber.agents.grid_locator import (
    PydanticAIGridLocator,
    column_index,
    grid_region_norm,
    overlay_grid,
    tighten_region,
)
from quber.files.cache import resolve_document


def render(src: Path, page: int, dpi: int, out: Path):
    doc = fitz.open(str(src))
    try:
        pg = doc[page - 1]
        w, h = pg.rect.width, pg.rect.height
        pg.get_pixmap(dpi=dpi).save(str(out))
    finally:
        doc.close()
    return out, w, h


async def main() -> None:
    raw, page = sys.argv[1], int(sys.argv[2])
    src = Path(resolve_document(raw))
    loc = PydanticAIGridLocator()
    with tempfile.TemporaryDirectory() as tmp:
        png, pw, ph = render(src, page, loc.dpi, Path(tmp) / "p.png")
        gridded = Path(tmp) / "g.png"
        overlay_grid(png, gridded, loc.rows, loc.cols)
        from pydantic_ai import BinaryContent

        image = BinaryContent(data=gridded.read_bytes(), media_type="image/png")
        result = await loc.run_locator_with_retry(image, page)
        flags = result.output.tables
    _pw2, _ph2, words = page_words(src, page)

    print(f"{src.name} page {page}  ({loc.rows}x{loc.cols} grid, page {pw:.0f}x{ph:.0f}pt)\n")
    for f in flags:
        shares = any(
            o is not f and not (f.row_end < o.row_start or f.row_start > o.row_end) for o in flags
        )
        c0 = column_index(f.col_start, loc.cols)
        c1 = column_index(f.col_end, loc.cols)
        if c1 < c0:
            c0, c1 = c1, c0
        coarse = grid_region_norm(f.row_start, f.row_end, c0, c1, loc.rows, loc.cols)
        tight_shared = tighten_region(coarse, words, pw, ph, True)
        tight_full = tighten_region(coarse, words, pw, ph, False)
        print(f"ord{f.ordinal} rows {f.row_start}..{f.row_end} cols {f.col_start}..{f.col_end} "
              f"(idx {c0}..{c1})  shares_rows={shares}  title={f.title!r}")
        print(f"   coarse x0={coarse[0]:.4f}")
        print(f"   tighten(shares={shares}) -> x0={tight_shared[0]:.4f}" if tight_shared else "   tighten(shared)=None")
        print(f"   tighten(full extent)     -> x0={tight_full[0]:.4f}" if tight_full else "   tighten(full)=None")
        used = (tight_shared if shares else tight_full) or coarse
        print(f"   --> USED x0 = {used[0]:.4f}  (label column kept if x0<=~0.06)\n")


if __name__ == "__main__":
    asyncio.run(main())