"""QUE-270 probe: can the Set-of-Mark vision model find footnote blocks?

Question this answers
---------------------
Today the per-table structure-correction step captures a footnote only when the
definition block falls inside that table's crop, and `share_footnotes` then
spreads it page-wide by marker string alone. That marker-only sharing mis-attributes
when a page carries two independent footnote namespaces. This probe tests a
grounded alternative: detect footnote BLOCKS visually as their own page regions,
ground each to the PDF text layer, read its leading marker deterministically, and
associate it to a table by GEOMETRY (the block sits directly below the table's
column band), confirmed by the marker the table references.

Two detector shapes are run on the SAME gridded page image so the visual model's
capability can be compared directly:

  A. Separate footnote-only locator pass (one extra vision call; dedicated prompt).
  B. Extended locator that returns tables AND footnote blocks in one call.

Nothing in the production package is modified. The probe reuses the prod scaffold
read-only: `overlay_grid`, the grid geometry, `tighten_region`, `page_words`, and
the real `PydanticAIGridLocator` for table regions in approach A.

Run
---
    uv run python experiments/que270/footnote_finder_probe.py            # Visa p11
    uv run python experiments/que270/footnote_finder_probe.py <pdf> <page>

Outputs a side-by-side report to stdout and writes the gridded page image plus a
JSON dump of both runs under experiments/que270/out/.
"""

from __future__ import annotations

import asyncio
import json
import re
import sys
import tempfile
from pathlib import Path
from typing import Any, List, Optional, Sequence, Tuple

from pydantic import BaseModel, Field

from quber.agents.completeness import page_words
from quber.agents.grid_locator import (
    DEFAULT_COLS,
    DEFAULT_ROWS,
    PydanticAIGridLocator,
    column_index,
    column_labels,
    grid_region_norm,
    overlay_grid,
    tighten_region,
)
from quber.core.extractors.camelot.acquire import render_pages
from quber.settings import DEFAULT_LLM_MODEL, get_settings

REPO = Path(__file__).resolve().parents[2]
DEFAULT_PDF = REPO / ".cache/s3/qubera-docs/visa/Q1-2026-Earnings-Release_vF.pdf"
DEFAULT_PAGE = 11
DPI = 200
OUT_DIR = Path(__file__).resolve().parent / "out"

# A footnote definition's leading marker, read off the grounded PDF text.
_LEADING_MARKER = re.compile(r"^\s*(?:\((\d{1,2})\)|(\d{1,2})\.|([*†‡§]))")


# --- model I/O schemas -------------------------------------------------------


class FootnoteBlockFlag(BaseModel):
    ordinal: int = Field(ge=1, description="1-based reading order of the footnote block")
    markers: List[str] = Field(
        default_factory=list,
        description="The marker each footnote line leads with, e.g. ['(1)','(2)','*']",
    )
    row_start: int = Field(description="Printed grid row at the block's TOP edge")
    row_end: int = Field(description="Printed grid row at the block's BOTTOM edge")
    col_start: str = Field(description="Printed grid column LETTER at the block's LEFT edge")
    col_end: str = Field(description="Printed grid column LETTER at the block's RIGHT edge")
    text_preview: str = Field(default="", description="First few words of the block, verbatim")


class FootnoteFinderResult(BaseModel):
    blocks: List[FootnoteBlockFlag] = Field(default_factory=list)


class TableFlag(BaseModel):
    ordinal: int = Field(ge=1)
    title: str = Field(default="")
    row_start: int
    row_end: int
    col_start: str
    col_end: str


class CombinedResult(BaseModel):
    tables: List[TableFlag] = Field(default_factory=list)
    blocks: List[FootnoteBlockFlag] = Field(default_factory=list)


# --- prompts -----------------------------------------------------------------


def _grid_preamble(rows: int, cols: int) -> str:
    last = column_labels(cols)[-1]
    return (
        f"The page image has a grid overlaid: {rows} numbered horizontal rows (1 at the "
        f"top to {rows} at the bottom, labeled on BOTH side margins) and {cols} lettered "
        f"columns (A at the left to {last} at the right, labeled top and bottom). Read the "
        f"printed row numbers and column letters off the grid; never estimate coordinates."
    )


_FOOTNOTE_DEFINITION = (
    "A FOOTNOTE BLOCK is one or more short notes that explain a marker used elsewhere on "
    "the page. Each line leads with a marker: a parenthesized or superscript number "
    "((1), (2), a raised 1/2), an asterisk *, or a dagger / double-dagger / section sign. "
    "Footnote blocks are visually distinct: smaller or lighter text, set off below a table "
    "by whitespace or a thin rule, usually near the bottom of the page.\n\n"
    "Do NOT report any of the following as a footnote block:\n"
    "- a bulleted list (lines led by a round bullet, dash, or square) — that is body text;\n"
    "- a numbered list that is the document's main content (full sentences, not tied to a "
    "marker that appears up in a table);\n"
    "- a table row, a column header, a title, a units/scale caption, or a prose paragraph;\n"
    "- page furniture (page number, running header, URL).\n"
    "The distinguishing test: a footnote's marker also appears attached to a header or cell "
    "somewhere in a table above it. A bullet's marker does not."
)


def footnote_only_prompt(rows: int, cols: int) -> str:
    return f"""\
{_grid_preamble(rows, cols)}

{_FOOTNOTE_DEFINITION}

Find every FOOTNOTE BLOCK on the page. For each, in reading order, report:
- ordinal (1 = first, top-to-bottom)
- markers: the marker each line in the block leads with, in order (e.g. ["(1)","(2)"])
- row_start / row_end: printed grid rows at the block's top and bottom edges
- col_start / col_end: printed grid column letters at the block's left and right edges
- text_preview: the first several words of the block, copied verbatim

If the page has no footnote block, return an empty list. Never invent a footnote
from a marker that has no distinct note text below the table.

Return ONLY a JSON object conforming to this schema (no prose, no code fences):

{json.dumps(FootnoteFinderResult.model_json_schema(), indent=2)}
"""


def combined_prompt(rows: int, cols: int) -> str:
    return f"""\
{_grid_preamble(rows, cols)}

Report TWO things on this page: its tables, and its footnote blocks.

TABLES: a block where values line up in two or more vertical columns across rows.
Bound the whole table (title, headers, row labels, body, totals). A bulleted or
numbered list, a single column of figures, and prose are NOT tables. For each
table report ordinal, title, row_start, row_end, col_start, col_end.

{_FOOTNOTE_DEFINITION}

For each footnote block report ordinal, markers, row_start, row_end, col_start,
col_end, text_preview (as in the footnote schema). Footnotes are NOT tables and
must not appear in the tables list.

Return ONLY a JSON object conforming to this schema (no prose, no code fences):

{json.dumps(CombinedResult.model_json_schema(), indent=2)}
"""


# --- agent plumbing (mirrors PydanticAIGridLocator) --------------------------


def _build_agent(system_prompt: str, output_type: Any) -> Any:
    from pydantic_ai import Agent
    from pydantic_ai.settings import ModelSettings

    llm = get_settings().llm
    model = llm.model or DEFAULT_LLM_MODEL
    if not llm.anthropic_auth_token:
        raise RuntimeError("ANTHROPIC_AUTH_TOKEN not set; this probe needs the real model.")
    from quber.agents._oauth_gate import make_oauth_anthropic_model

    anth_model = make_oauth_anthropic_model(model, llm.anthropic_auth_token)
    return Agent(
        anth_model,
        output_type=output_type,
        system_prompt=system_prompt,
        model_settings=ModelSettings(temperature=0.0),
    )


async def _run_agent(system_prompt: str, output_type: Any, gridded: Path) -> Any:
    from pydantic_ai import BinaryContent

    agent = _build_agent(system_prompt, output_type)
    image = BinaryContent(data=gridded.read_bytes(), media_type="image/png")
    result = await agent.run(["Report using the grid.", image])
    return result.output


# --- grounding + association (deterministic, reuses prod geometry) -----------


def _ground_block(
    flag: FootnoteBlockFlag,
    words: Sequence[Tuple[float, float, float, float, str]],
    page_w: float,
    page_h: float,
    rows: int,
    cols: int,
) -> Tuple[Optional[Tuple[float, float, float, float]], str, Optional[str]]:
    """Coarse grid flag -> tight box, grounded text, and leading marker from PDF words."""
    c0 = column_index(flag.col_start, cols)
    c1 = column_index(flag.col_end, cols)
    if c1 < c0:
        c0, c1 = c1, c0
    coarse = grid_region_norm(flag.row_start, flag.row_end, c0, c1, rows, cols)
    # The model's vertical span is approximate (grid quantization), and a footnote
    # block butts directly against the table's total row, so a blunt geometric pad
    # cannot separate them. Anchor instead on the MARKER token: within a couple of
    # rows of the model's span, find the highest word that IS a bare marker ("(1)",
    # "*"). That line is the block's true top; the table row above has no marker.
    # Group nearby words into text lines, then keep only FOOTNOTE-DEFINITION lines:
    # a line whose first token is a marker AND whose next token is alphabetic prose.
    # A negative-value row ("(39) — 9 30") leads with a marker-shaped token but is
    # followed by numbers, so it is rejected — the markers-vs-negatives hazard that
    # also bites page-wide marker sharing.
    search_top = (coarse[1] - 2.5 / rows) * page_h
    search_bot = (coarse[3] + 1.5 / rows) * page_h
    band = [w for w in words if search_top <= (w[1] + w[3]) / 2 <= search_bot]
    def_line_tops: List[float] = []
    for line in _group_lines(band):
        toks = [w[4].strip() for w in line]
        if len(toks) >= 2 and _LEADING_MARKER.match(toks[0]) and re.match(r"^[A-Za-z]", toks[1]):
            def_line_tops.append(min(w[1] for w in line))
    top_y = (min(def_line_tops) - 1.0) if def_line_tops else coarse[1] * page_h
    bot_y = search_bot
    inside = [w for w in words if top_y <= (w[1] + w[3]) / 2 <= bot_y]
    box = (coarse[0], top_y / page_h, coarse[2], bot_y / page_h)
    inside.sort(key=lambda w: (round((w[1] + w[3]) / 2 / 4), w[0]))
    text = " ".join(w[4] for w in inside)
    m = _LEADING_MARKER.match(text)
    marker = next((g for g in m.groups() if g), None) if m else None
    return box, text, marker


def _group_lines(
    words: Sequence[Tuple[float, float, float, float, str]],
) -> List[List[Tuple[float, float, float, float, str]]]:
    """Cluster words into text lines by vertical center, each sorted left-to-right."""
    lines: List[List[Tuple[float, float, float, float, str]]] = []
    for w in sorted(words, key=lambda w: (w[1] + w[3]) / 2):
        yc = (w[1] + w[3]) / 2
        if lines and abs(yc - (lines[-1][0][1] + lines[-1][0][3]) / 2) <= 4.0:
            lines[-1].append(w)
        else:
            lines.append([w])
    for line in lines:
        line.sort(key=lambda w: w[0])
    return lines


def _associate(
    block_box: Tuple[float, float, float, float],
    tables: Sequence[Tuple[int, Tuple[float, float, float, float]]],
) -> Optional[int]:
    """Ordinal of the nearest table whose column band overlaps and sits ABOVE the block."""
    bx0, by0, bx1, _ = (min(block_box[0], block_box[2]), min(block_box[1], block_box[3]),
                        max(block_box[0], block_box[2]), max(block_box[1], block_box[3]))
    best: Optional[int] = None
    best_gap = 1e9
    for ordinal, (tx0, ty0, tx1, ty1) in tables:
        t_x0, t_x1, t_bottom = min(tx0, tx1), max(tx0, tx1), max(ty0, ty1)
        x_overlap = min(bx1, t_x1) - max(bx0, t_x0)
        if x_overlap <= 0:
            continue
        gap = by0 - t_bottom  # block top minus table bottom; >= 0 means below
        if gap >= -0.02 and gap < best_gap:
            best_gap, best = gap, ordinal
    return best


# --- driver ------------------------------------------------------------------


async def run(pdf: Path, page: int) -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    print(f"PDF : {pdf}")
    print(f"page: {page}   model: {get_settings().llm.model or DEFAULT_LLM_MODEL}   dpi: {DPI}")
    print("=" * 78)

    page_w, page_h, words = page_words(pdf, page)

    with tempfile.TemporaryDirectory(prefix="que270-") as tmp:
        images = render_pages(pdf, DPI, Path(tmp))
        page_img = images[page - 1]
        gridded = OUT_DIR / f"{pdf.stem}-p{page}-gridded.png"
        overlay_grid(page_img, gridded, DEFAULT_ROWS, DEFAULT_COLS)

        # Tables: the real locator (re-overlays the raw page itself) for approach A.
        locator = PydanticAIGridLocator()
        located = await locator.locate(page_img, pdf, page)
        real_tables = [(t.ordinal, t.region) for t in located]

        print(f"\nReal grid locator found {len(real_tables)} table(s):")
        for ordinal, region in real_tables:
            print(f"  table {ordinal}: region={_fmt(region)}")

        # Approach A: dedicated footnote-only pass over the gridded image.
        print("\n" + "-" * 78)
        print("APPROACH A — separate footnote-only locator pass")
        print("-" * 78)
        a = await _run_agent(footnote_only_prompt(DEFAULT_ROWS, DEFAULT_COLS), FootnoteFinderResult, gridded)
        _report_blocks(a.blocks, words, page_w, page_h, real_tables)

        # Approach B: combined tables + footnotes in one call.
        print("\n" + "-" * 78)
        print("APPROACH B — extended locator (tables + footnotes, one call)")
        print("-" * 78)
        b = await _run_agent(combined_prompt(DEFAULT_ROWS, DEFAULT_COLS), CombinedResult, gridded)
        b_tables = [(t.ordinal, _flag_region(t, DEFAULT_ROWS, DEFAULT_COLS)) for t in b.tables]
        print(f"combined call reported {len(b.tables)} table(s), {len(b.blocks)} footnote block(s)")
        _report_blocks(b.blocks, words, page_w, page_h, b_tables)

    dump = OUT_DIR / f"{pdf.stem}-p{page}-runs.json"
    dump.write_text(
        json.dumps(
            {
                "approach_a": a.model_dump(),
                "approach_b": b.model_dump(),
                "real_tables": [{"ordinal": o, "region": r} for o, r in real_tables],
            },
            indent=2,
        )
    )
    print(f"\nWrote {dump}")
    print(f"Gridded image: {gridded}")


def _flag_region(flag: Any, rows: int, cols: int) -> Tuple[float, float, float, float]:
    c0, c1 = column_index(flag.col_start, cols), column_index(flag.col_end, cols)
    if c1 < c0:
        c0, c1 = c1, c0
    return grid_region_norm(flag.row_start, flag.row_end, c0, c1, rows, cols)


def _report_blocks(
    blocks: Sequence[FootnoteBlockFlag],
    words: Sequence[Tuple[float, float, float, float, str]],
    page_w: float,
    page_h: float,
    tables: Sequence[Tuple[int, Tuple[float, float, float, float]]],
) -> None:
    if not blocks:
        print("  (no footnote blocks reported)")
        return
    for blk in blocks:
        box, text, marker = _ground_block(blk, words, page_w, page_h, DEFAULT_ROWS, DEFAULT_COLS)
        owner = _associate(box, tables) if box else None
        preview = (text[:100] + "...") if len(text) > 100 else text
        print(f"  block {blk.ordinal}: model markers={blk.markers}")
        print(f"    grounded region : {_fmt(box)}")
        print(f"    grounded marker : {marker!r}  (read from PDF text, not the model)")
        print(f"    associated table: {owner}  (nearest table above with column overlap)")
        print(f"    grounded text   : {preview}")


def _fmt(region: Optional[Tuple[float, float, float, float]]) -> str:
    if region is None:
        return "None"
    return "(" + ", ".join(f"{v:.3f}" for v in region) + ")"


if __name__ == "__main__":
    pdf_arg = Path(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PDF
    page_arg = int(sys.argv[2]) if len(sys.argv) > 2 else DEFAULT_PAGE
    asyncio.run(run(pdf_arg, page_arg))