"""End-to-end: vision box -> guided Camelot -> grounded structure correction.

For each visual table: convert its box to a Camelot area and extract a grid
constrained to that region, then run the grounded correction step (crop image
arbitrates structure; values come only from Camelot + the page text layer; a
guard rejects any invented figure). Prints the final cleaned markdown per
table so the whole pipeline can be eyeballed on hard pages.
"""

from __future__ import annotations

import asyncio
import json
import sys
import tempfile
from pathlib import Path

import fitz

from quber.agents.llm_client import get_llm_client
from quber.core.extractors.camelot.acquire import grid_to_markdown
from quber.core.extractors.camelot.correspondence.correction import correct_structure
from quber.core.extractors.camelot.correspondence.geometry import norm_bbox_to_table_area
from quber.core.extractors.camelot.correspondence.recovery import camelot_targeted

DPI = 200


async def run(dual_json: Path, source: Path, pages: list[int]) -> None:
    data = json.loads(dual_json.read_text())
    doc = fitz.open(str(source))
    visual_by_page: dict[int, list] = {}
    for v in data["visual"]:
        visual_by_page.setdefault(v["page"], []).append(v)

    llm = get_llm_client(None)
    sem = asyncio.Semaphore(4)

    for page in pages:
        w, h = doc[page - 1].rect.width, doc[page - 1].rect.height
        vtables = sorted(visual_by_page.get(page, []), key=lambda v: v["ordinal"])
        print(f"\n{'='*80}\nPAGE {page}: {len(vtables)} visual tables")
        with tempfile.TemporaryDirectory() as tmp:
            img = Path(tmp) / "page.png"
            doc[page - 1].get_pixmap(dpi=DPI).save(str(img))
            for v in vtables:
                area = norm_bbox_to_table_area(tuple(v["region"]), w, h)
                try:
                    cand = camelot_targeted(str(source), page, area, v["ordinal"])
                except Exception as exc:
                    print(f"\n  V{v['ordinal']} {v['title'][:45]!r}: camelot raised {exc}")
                    continue
                if cand is None:
                    print(f"\n  V{v['ordinal']} {v['title'][:45]!r}: camelot found nothing")
                    continue
                raw_md = grid_to_markdown(cand.cells)
                correction = await correct_structure(
                    cand.cells, img, cand.bbox, str(source), page, llm, sem, DPI
                )
                final_md = correction.markdown if correction else raw_md
                corrected = bool(correction and correction.llm_corrected)
                print(f"\n  V{v['ordinal']} {v['title'][:45]!r}  acc={cand.accuracy:.0f} "
                      f"raw={len(cand.cells)}rows corrected={corrected}")
                if correction:
                    print(f"     title={correction.title[:55]!r} footnotes={len(correction.footnotes)}")
                print("     --- FINAL markdown (first 8 lines) ---")
                for line in final_md.splitlines()[:8]:
                    print("     " + line[:100])


def main() -> None:
    jobs = "/home/mande/.claude/jobs/a95a5f59/tmp"
    targets = {
        "daloopa": (f"{jobs}/93239979.dual.json", ".cache/s3/qubera-docs/daloopa/93239979/93239979.pdf",
                    [10, 23]),
        "bhe": (f"{jobs}/BHE_991.dual.json", "documents/BHE_991.pdf", [1]),
    }
    which = sys.argv[1] if len(sys.argv) > 1 else "daloopa"
    js, pdf, pages = targets[which]
    asyncio.run(run(Path(js), Path(pdf), pages))


if __name__ == "__main__":
    main()