"""Demand-driven footnote rendering: snapped boxes only on candidate pages.

For each corrected table (from a `quber table` run's *.tables.json), the
reference markers are read off the CORRECTED markdown with REFERENCE_MARKER /
REFERENCE_SYMBOL. A page is a footnote CANDIDATE only when at least one of its
tables references a marker. Only candidate pages run the footnote finder; every
other page is skipped (no LLM call, no box).

On candidate pages the finder's footnote blocks are grounded through the SAME
production geometry chain the tables use (tighten_region -> pad_boxes ->
declash), so the drawn footnote box is snapped to the page text layer. Each box
is labelled with the marker read off the snapped region and the table ordinals
that reference it.

Output: an annotated PDF per doc (blue = table som_region, orange dashed =
snapped footnote candidate) plus PNGs of the candidate pages for quick review.

Run:
    uv run python experiments/que270/demand_driven_annotate.py
"""

from __future__ import annotations

import asyncio
import glob
import json
import re
import sys
import tempfile
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Tuple

import fitz
from loguru import logger

sys.path.insert(0, str(Path(__file__).resolve().parent))
import annotate_regions as ar  # noqa: E402
import footnote_finder_probe as ffp  # noqa: E402

from quber.agents.completeness import page_words  # noqa: E402
from quber.agents.grid_locator import (  # noqa: E402
    DEFAULT_COLS,
    DEFAULT_ROWS,
    LocatedTable,
    overlay_grid,
)

ar.GROUNDING = "prod"  # snap footnote boxes through the production geometry chain

CORPUS = Path("experiments/que270/out/strip_test_corpus")
OUT_DIR = Path("experiments/que270/out/demand_driven")
DPI = 200

# A marker REFERENCE attached to a label ("Income Tax Provision(1)"); the
# lookbehind excludes a bare parenthesized negative value ("(11)").
REFERENCE_MARKER = re.compile(r"(?<=[A-Za-z%)])\((\d{1,2})\)")
REFERENCE_SYMBOL = re.compile(r"(?<=[A-Za-z%)])([*†‡§])")


def referenced_markers(corrected_markdown: str) -> set:
    """Markers a CORRECTED table references in its header/cell text."""
    m = set(REFERENCE_MARKER.findall(corrected_markdown))
    m.update(REFERENCE_SYMBOL.findall(corrected_markdown))
    return m


def _located(region, ordinal: int) -> LocatedTable:
    """Wrap a stored som_region as a LocatedTable for the grounding chain."""
    return LocatedTable(
        ordinal=ordinal,
        title="",
        region=tuple(region),
        grid_rows=(1, 1),
        grid_cols=(0, 0),
        tightened=True,
    )


async def annotate_doc(fn_agent: object, tables_json: Path) -> dict:
    tables = json.loads(tables_json.read_text())
    if not tables:
        return {"doc": tables_json.stem, "candidate_pages": []}
    pdf = Path(tables[0]["source"])

    # Group tables by page; keep region + per-table referenced markers.
    by_page: Dict[int, List[dict]] = defaultdict(list)
    for t in tables:
        by_page[t["page"]].append(t)

    doc = fitz.open(str(pdf))
    candidate_pages: List[int] = []
    findings: List[dict] = []

    for page in sorted(by_page):
        page_tables = by_page[page]
        located = [
            _located(t["som_region"], i + 1) for i, t in enumerate(page_tables) if t.get("som_region")
        ]
        # Draw the table regions (blue) on every page that has them.
        tbl_boxes = [(i + 1, t.get("title", ""), t["som_region"]) for i, t in enumerate(page_tables) if t.get("som_region")]

        # Demand-driven gate: which markers do this page's tables reference?
        page_refs: set = set()
        ref_by_table: Dict[int, set] = {}
        for i, t in enumerate(page_tables):
            r = referenced_markers(t.get("markdown", ""))
            ref_by_table[i + 1] = r
            page_refs |= r

        fn_boxes: List[Tuple[List[str], tuple]] = []
        if page_refs and located:
            candidate_pages.append(page)
            pw, ph, words = await asyncio.to_thread(page_words, pdf, page)
            with tempfile.TemporaryDirectory(prefix="dd-") as tmp:
                raw = Path(tmp) / "p.png"
                _render_page(doc, page, raw)
                gridded = Path(tmp) / "g.png"
                await asyncio.to_thread(overlay_grid, raw, gridded, DEFAULT_ROWS, DEFAULT_COLS)
                from pydantic_ai import BinaryContent

                image = BinaryContent(data=gridded.read_bytes(), media_type="image/png")
                try:
                    result = await fn_agent.run(["Report using the grid.", image])  # type: ignore[attr-defined]
                    blocks = result.output.blocks
                except Exception as exc:
                    logger.warning("finder failed {} p{}: {}", pdf.name, page, exc)
                    blocks = []
            fn_boxes = ar._ground_footnotes(blocks, located, words, pw, ph)
            findings.append(
                _attribute(pdf.name, page, page_refs, ref_by_table, fn_boxes, words, pw, ph)
            )

        _draw_page(doc, page, tbl_boxes, fn_boxes, ref_by_table)

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    out_pdf = OUT_DIR / f"{tables_json.stem}.dd.pdf"
    doc.save(str(out_pdf))
    # PNGs of candidate pages for quick review.
    for page in candidate_pages:
        doc[page - 1].get_pixmap(dpi=120).save(str(OUT_DIR / f"{tables_json.stem}_p{page}.png"))
    doc.close()
    logger.info("{}: candidate pages {}", pdf.name, candidate_pages)
    return {"doc": tables_json.stem, "candidate_pages": candidate_pages, "findings": findings}


def _attribute(doc, page, page_refs, ref_by_table, fn_boxes, words, pw, ph) -> dict:
    """Match each grounded footnote block to the table ordinals referencing the
    markers it defines; report satisfied/missing. A block may define several
    markers ((1)(2)(3)...), so every definition-line marker inside the box is
    collected, not just the leading one."""
    found_markers: set = set()
    blocks_out = []
    for markers, region in fn_boxes:
        text, box_markers = _read_region(region, words, pw, ph)
        found_markers |= box_markers
        owners = [ord_ for ord_, refs in ref_by_table.items() if refs & box_markers]
        blocks_out.append({"markers": sorted(box_markers), "owner_tables": owners, "text": text[:90]})
    return {
        "doc": doc,
        "page": page,
        "referenced": sorted(page_refs),
        "definitions_found": sorted(found_markers),
        "satisfied": sorted(page_refs & found_markers),
        "missing": sorted(page_refs - found_markers),
        "extra": sorted(found_markers - page_refs),
        "blocks": blocks_out,
    }


def _read_region(region, words, pw, ph) -> Tuple[str, set]:
    """Region text plus EVERY definition marker that begins a line inside it."""
    x0, y0, x1, y1 = (min(region[0], region[2]), min(region[1], region[3]),
                      max(region[0], region[2]), max(region[1], region[3]))
    inside = [
        w for w in words
        if x0 * pw - 2 <= (w[0] + w[2]) / 2 <= x1 * pw + 2 and y0 * ph - 2 <= (w[1] + w[3]) / 2 <= y1 * ph + 2
    ]
    markers: set = set()
    for line in ffp._group_lines(inside):
        toks = [w[4].strip() for w in line]
        if toks:
            m = ffp._LEADING_MARKER.match(toks[0])
            if m:
                markers.add(next((g for g in m.groups() if g), None))
    inside.sort(key=lambda w: (round((w[1] + w[3]) / 2 / 4), w[0]))
    text = " ".join(w[4] for w in inside)
    return text, markers


def _render_page(doc: fitz.Document, page: int, out: Path) -> None:
    zoom = DPI / 72.0
    doc[page - 1].get_pixmap(matrix=fitz.Matrix(zoom, zoom)).save(str(out))


def _draw_page(doc, page_no, tbl_boxes, fn_boxes, ref_by_table) -> None:
    page = doc[page_no - 1]
    w, h = page.rect.width, page.rect.height
    for ordinal, title, region in tbl_boxes:
        rect = ar._rect(region, w, h)
        page.draw_rect(rect, color=ar.TABLE_RGB, width=1.4)
        refs = ref_by_table.get(ordinal) or set()
        lab = f"T{page_no}.{ordinal}" + (f" refs={sorted(refs)}" if refs else "")
        page.insert_text((rect.x0 + 2, max(8.0, rect.y0 - 3)), lab[:80], fontsize=6, color=ar.TABLE_RGB)
    for markers, region in fn_boxes:
        rect = ar._rect(region, w, h)
        page.draw_rect(rect, color=ar.FOOTNOTE_RGB, width=1.4, dashes="[3 2] 0")
        lab = "fn(candidate) " + " ".join(markers) if markers else "fn(candidate)"
        page.insert_text((rect.x0 + 2, min(h - 2, rect.y1 + 7)), lab[:60], fontsize=6, color=ar.FOOTNOTE_RGB)


async def main() -> None:
    fn_agent = ffp._build_agent(ffp.footnote_only_prompt(DEFAULT_ROWS, DEFAULT_COLS), ffp.FootnoteFinderResult)
    files = sorted(glob.glob(str(CORPUS / "*" / "*.tables.json")))
    index = []
    for f in files:
        index.append(await annotate_doc(fn_agent, Path(f)))
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    (OUT_DIR / "index.json").write_text(json.dumps(index, indent=2))
    # Aggregate
    cand = [d for d in index if d.get("candidate_pages")]
    findings = [fd for d in cand for fd in d.get("findings", [])]
    print(f"docs scanned: {len(index)}   docs with candidate pages: {len(cand)}")
    print(f"candidate pages searched: {sum(len(d['candidate_pages']) for d in index)}")
    print(f"referenced markers satisfied: {sum(len(fd['satisfied']) for fd in findings)}")
    print(f"missing (referenced, no def found): {sum(len(fd['missing']) for fd in findings)}")
    print(f"extra (def found, unreferenced): {sum(len(fd['extra']) for fd in findings)}")
    print(f"\nOutput under {OUT_DIR}")


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