"""Tests for the grounded structure correction wired into
CamelotCorrespondenceExtractor.

`apply_structure_correction` crops the table region (structural arbiter), pulls
the region text layer (value well), calls the LLM's `vet_structure`, and then
applies a grounding guard: any output value absent from BOTH Camelot's grid and
the page text layer is rejected. These tests exercise that logic with a mock LLM
and monkeypatched PDF/image I/O (no Camelot, rendering, or network), plus
`make_extracted` directly.
"""

import asyncio
from pathlib import Path
from typing import Optional, Sequence, Tuple

import pytest
from _pytest.monkeypatch import MonkeyPatch

from quber.agents.completeness import MockCompleteness
from quber.agents.detector import DetectedTable, DetectorResult
from quber.agents.llm_client import FootnoteDef, LLMTableCorrection, MockLLMClient
from quber.core.extractors.camelot.acquire import CamelotCandidate, grid_to_markdown
from quber.core.extractors.camelot.correspondence import (
    CamelotCorrespondenceExtractor,
)
from quber.core.extractors.camelot.correspondence import (
    correction as correction_module,
)
from quber.core.extractors.camelot.correspondence.correction import StructureCorrection

BBOX: Tuple[float, float, float, float] = (0.0, 0.0, 100.0, 100.0)
SOURCE = "doc.pdf"
PAGE = 1


def patch_io(
    monkeypatch: MonkeyPatch,
    words: Sequence[Tuple[float, float, float, float, str]] = (),
    blocks: Sequence[str] = (),
) -> None:
    """Stub the PDF text-layer reads and the page crop so the correction logic
    runs without a real PDF. `words` seeds the page's word layer and `blocks`
    its text blocks, the two readings a copied caption or title is checked against."""
    monkeypatch.setattr(correction_module, "page_words", lambda source, page: (612.0, 792.0, list(words)))
    monkeypatch.setattr(correction_module, "page_blocks", lambda source, page: list(blocks))
    monkeypatch.setattr(correction_module, "crop_region_png", lambda page_image, region, dpi: b"\x89PNG")


def run(ext: CamelotCorrespondenceExtractor, cells: list[list[str]], image: Optional[Path]):
    return asyncio.run(ext.apply_structure_correction(cells, image, BBOX, SOURCE, PAGE))


# A fragmented multi-row header (rows 0-2) above clean data rows, modelled on
# the real Visa Q1FY25 case the ticket cites: the header is split across rows
# and the data carries currency, thousands separators, and a parenthesized
# (negative) percent that exact-literal handling must preserve byte-for-byte.
FRAGMENTED_GRID = [
    ["Branded", "Volume", "and"],
    ["Transactions", "", ""],
    ["Region", "2025", "2024"],
    ["US", "1,637", "1,500"],
    ["EU", "(5.0%)", "10.0%"],
]

MERGED_MARKDOWN = "| Region | 2025 | 2024 |\n|---|---|---|\n| US | 1,637 | 1,500 |\n| EU | (5.0%) | 10.0% |"


class StubDetector:
    """No-op TableDetector; never called by these unit paths."""

    async def detect(self, page_image: Path) -> DetectorResult:
        _ = page_image
        return DetectorResult(tables=[])


def make_extractor(
    llm: Optional[object] = None,
    run_llm_correction: bool = True,
) -> CamelotCorrespondenceExtractor:
    return CamelotCorrespondenceExtractor(
        detector=StubDetector(),
        completeness=MockCompleteness(),
        run_box_repair=False,
        llm=llm,  # type: ignore[arg-type]
        run_llm_correction=run_llm_correction,
    )


def candidate(cells: list[list[str]]) -> CamelotCandidate:
    return CamelotCandidate(
        candidate_id="lattice-p1-0",
        flavor="lattice",
        page=1,
        accuracy=95.0,
        cells=cells,
        markdown=grid_to_markdown(cells),
    )


def page_image(tmp_path: Path) -> Path:
    p = tmp_path / "page-0001.png"
    p.write_bytes(b"\x89PNG\r\n\x1a\n")
    return p


# --- apply_structure_correction --------------------------------------------


def _line(y: float, text: str) -> list[Tuple[float, float, float, float, str]]:
    """One printed line of the page text layer, one word box per word."""
    return [(10.0 * i, y, 10.0 * i + 9.0, y + 8.0, w) for i, w in enumerate(text.split())]


def test_correction_vets_structure(tmp_path: Path, monkeypatch: MonkeyPatch):
    patch_io(
        monkeypatch,
        words=_line(20.0, "Branded Volume and Transactions")
        + _line(32.0, "For the 3 Months Ended September 30, 2025"),
    )
    correction = LLMTableCorrection(
        title="Branded Volume and Transactions",
        caption="For the 3 Months Ended September 30, 2025",
        markdown=MERGED_MARKDOWN,
        footnotes=[FootnoteDef(marker="(1)", text="Preliminary.")],
        header_rows=1,
    )
    ext = make_extractor(llm=MockLLMClient(correction=correction))

    result = run(ext, FRAGMENTED_GRID, page_image(tmp_path))

    assert isinstance(result, StructureCorrection)
    assert result.markdown == MERGED_MARKDOWN
    assert result.llm_corrected is True
    assert result.title == "Branded Volume and Transactions"
    assert result.caption == "For the 3 Months Ended September 30, 2025"
    assert result.footnotes == [FootnoteDef(marker="(1)", text="Preliminary.")]


def test_correction_drops_title_and_caption_the_page_does_not_print(tmp_path: Path, monkeypatch: MonkeyPatch):
    # A composed name or a paraphrased sentence is not on the page, so it is
    # blanked; the printed one, matched across dash and apostrophe styles, stays.
    patch_io(
        monkeypatch, words=_line(20.0, "Concentration of Credit Risk — KREF’s loans (“REO”) as a percentage:")
    )
    correction = LLMTableCorrection(
        title="Computation of Basic and Diluted EPS",
        caption='Concentration of Credit Risk - KREF\'s loans ("REO") as a percentage:',
        markdown=MERGED_MARKDOWN,
        footnotes=[],
        header_rows=1,
    )
    ext = make_extractor(llm=MockLLMClient(correction=correction))
    result = run(ext, FRAGMENTED_GRID, page_image(tmp_path))
    assert isinstance(result, StructureCorrection)
    assert result.title == ""
    assert result.caption == 'Concentration of Credit Risk - KREF\'s loans ("REO") as a percentage:'
    # The structural correction itself is unaffected by a dropped header line.
    assert result.markdown == MERGED_MARKDOWN


def test_located_title_kept_only_when_printed(monkeypatch: MonkeyPatch):
    # When no correction ran, the grid locator's title is the only one on
    # offer; it passes the same printed-text check a copied title does.
    from quber.core.extractors.camelot.correspondence.correction import printed_title

    patch_io(monkeypatch, blocks=["Portfolio Information", "as of March 31, 2026"])
    assert asyncio.run(printed_title("Portfolio Information as of March 31, 2026", SOURCE, PAGE)) == ""
    patch_io(monkeypatch, blocks=["Portfolio Information as of\nMarch 31, 2026"])
    assert (
        asyncio.run(printed_title("Portfolio Information as of March 31, 2026", SOURCE, PAGE))
        == "Portfolio Information as of March 31, 2026"
    )
    assert asyncio.run(printed_title("", SOURCE, PAGE)) == ""


def test_correction_hands_the_agent_the_whole_page_text(tmp_path: Path, monkeypatch: MonkeyPatch):
    # The agent gets the page as its text blocks, the units the layout keeps
    # together, beside the region text.
    seen: dict[str, str] = {}

    class Spy(MockLLMClient):
        async def vet_structure(self, image_png, markdown, region_text, page_text):
            seen["page_text"] = page_text
            return await super().vet_structure(image_png, markdown, region_text, page_text)

    patch_io(
        monkeypatch,
        blocks=[
            "Concentration of Credit Risk — The following tables present:",
            "Geography(A)\nCalifornia 17.6",
        ],
    )
    ext = make_extractor(llm=Spy())
    run(ext, FRAGMENTED_GRID, page_image(tmp_path))
    assert seen["page_text"] == (
        "Concentration of Credit Risk — The following tables present:\n\nGeography(A)\nCalifornia 17.6"
    )


def test_correction_keeps_a_title_printed_over_two_lines_or_with_a_superscript(
    tmp_path: Path, monkeypatch: MonkeyPatch
):
    # A slide prints a title over two lines beside another column, so the
    # page's lines interleave the two; the text block keeps it whole. A
    # footnote digit set as a superscript lands as its own word.
    patch_io(
        monkeypatch,
        words=_line(20.0, "GAAP Net Income to Summary")
        + _line(32.0, "Distributable Earnings Reconciliation Income"),
        blocks=[
            "GAAP Net Income to\nDistributable Earnings Reconciliation",
            "Summary\nIncome",
            "Cash Available for Distribution\n2",
        ],
    )
    correction = LLMTableCorrection(
        title="GAAP Net Income to Distributable Earnings Reconciliation",
        caption="Cash Available for Distribution2",
        markdown=MERGED_MARKDOWN,
        footnotes=[],
        header_rows=1,
    )
    ext = make_extractor(llm=MockLLMClient(correction=correction))
    result = run(ext, FRAGMENTED_GRID, page_image(tmp_path))
    assert isinstance(result, StructureCorrection)
    assert result.title == "GAAP Net Income to Distributable Earnings Reconciliation"
    assert result.caption == "Cash Available for Distribution2"


def test_correction_carries_units_attribution(tmp_path: Path, monkeypatch: MonkeyPatch):
    # The scale/currency caption the vet model routes into `units` must survive
    # onto the StructureCorrection so the live pipeline can carry it to the table.
    patch_io(monkeypatch)
    correction = LLMTableCorrection(
        title="Branded Volume and Transactions",
        caption="",
        markdown=MERGED_MARKDOWN,
        footnotes=[],
        units="(in millions, except percentages and per share data)",
        header_rows=1,
    )
    ext = make_extractor(llm=MockLLMClient(correction=correction))

    result = run(ext, FRAGMENTED_GRID, page_image(tmp_path))

    assert isinstance(result, StructureCorrection)
    assert result.units == "(in millions, except percentages and per share data)"


def test_correction_units_defaults_empty(tmp_path: Path, monkeypatch: MonkeyPatch):
    # A table with no caption carries an empty `units` -- never an invented one.
    patch_io(monkeypatch)
    correction = LLMTableCorrection(
        title="", caption="", markdown=MERGED_MARKDOWN, footnotes=[], header_rows=1
    )
    ext = make_extractor(llm=MockLLMClient(correction=correction))

    result = run(ext, FRAGMENTED_GRID, page_image(tmp_path))

    assert isinstance(result, StructureCorrection)
    assert result.units == ""


def test_correction_carries_footnote_refs(tmp_path: Path, monkeypatch: MonkeyPatch):
    # The footnote reference markers the vet model reads off the table image must
    # survive onto the StructureCorrection so the pipeline can carry them to the
    # table; the demand-driven footnote lookup keys off these.
    patch_io(monkeypatch)
    correction = LLMTableCorrection(
        title="",
        caption="",
        markdown=MERGED_MARKDOWN,
        footnotes=[],
        footnote_refs=["(1)", "*"],
        header_rows=1,
    )
    ext = make_extractor(llm=MockLLMClient(correction=correction))

    result = run(ext, FRAGMENTED_GRID, page_image(tmp_path))

    assert isinstance(result, StructureCorrection)
    assert result.footnote_refs == ["(1)", "*"]


def test_correction_footnote_refs_default_empty(tmp_path: Path, monkeypatch: MonkeyPatch):
    # A table carrying no markers gets an empty footnote_refs -- never invented.
    patch_io(monkeypatch)
    correction = LLMTableCorrection(
        title="", caption="", markdown=MERGED_MARKDOWN, footnotes=[], header_rows=1
    )
    ext = make_extractor(llm=MockLLMClient(correction=correction))

    result = run(ext, FRAGMENTED_GRID, page_image(tmp_path))

    assert isinstance(result, StructureCorrection)
    assert result.footnote_refs == []


def test_correction_passthrough_marks_not_corrected(tmp_path: Path, monkeypatch: MonkeyPatch):
    # MockLLMClient with no canned correction echoes the input markdown, so the
    # presentation is unchanged and llm_corrected must stay False.
    patch_io(monkeypatch)
    ext = make_extractor(llm=MockLLMClient())

    result = run(ext, FRAGMENTED_GRID, page_image(tmp_path))

    assert result is not None
    assert result.markdown == grid_to_markdown(FRAGMENTED_GRID)
    assert result.llm_corrected is False


def test_correction_rejected_on_ungrounded_value(tmp_path: Path, monkeypatch: MonkeyPatch):
    # '99,999' is in neither Camelot's grid nor the (empty) page text layer, so
    # the correction is rejected and the grid markdown is kept.
    patch_io(monkeypatch)
    invented = MERGED_MARKDOWN + "\n| ZZ | 99,999 | |"
    correction = LLMTableCorrection(title="", caption="", markdown=invented, footnotes=[], header_rows=1)
    ext = make_extractor(llm=MockLLMClient(correction=correction))
    assert run(ext, FRAGMENTED_GRID, page_image(tmp_path)) is None


def test_correction_accepts_value_grounded_in_text_layer(tmp_path: Path, monkeypatch: MonkeyPatch):
    # A value absent from Camelot's grid but present in the page text layer is
    # grounded -> the recovery is accepted.
    patch_io(monkeypatch, words=[(0.0, 0.0, 10.0, 10.0, "240,083")])
    recovered = MERGED_MARKDOWN + "\n| ZZ | 240,083 | |"
    correction = LLMTableCorrection(title="", caption="", markdown=recovered, footnotes=[], header_rows=1)
    ext = make_extractor(llm=MockLLMClient(correction=correction))
    result = run(ext, FRAGMENTED_GRID, page_image(tmp_path))
    assert result is not None
    assert "240,083" in result.markdown


def test_correction_blanks_invented_placeholder_dash(tmp_path: Path, monkeypatch: MonkeyPatch):
    # The agent wrote an em-dash into a printed-blank cell: the dash exists in
    # neither Camelot's grid nor the text layer, so it is normalized to the
    # blank the page shows. The rest of the correction is untouched.
    patch_io(monkeypatch)
    dashed = MERGED_MARKDOWN + "\n| ZZ | — | 1,637 |"
    correction = LLMTableCorrection(title="", caption="", markdown=dashed, footnotes=[], header_rows=1)
    ext = make_extractor(llm=MockLLMClient(correction=correction))
    result = run(ext, FRAGMENTED_GRID, page_image(tmp_path))
    assert result is not None
    assert "—" not in result.markdown
    assert "| ZZ |" in result.markdown and "1,637" in result.markdown


def test_correction_keeps_printed_nil_dash(tmp_path: Path, monkeypatch: MonkeyPatch):
    # A dash present in the table's text layer is a printed nil marker, not an
    # invented placeholder: it stays. The word must sit inside the table region
    # (BBOX bottom-left -> top-left points puts the region at y 686..798).
    patch_io(monkeypatch, words=[(10.0, 700.0, 20.0, 710.0, "—")])
    dashed = MERGED_MARKDOWN + "\n| ZZ | — | 1,637 |"
    correction = LLMTableCorrection(title="", caption="", markdown=dashed, footnotes=[], header_rows=1)
    ext = make_extractor(llm=MockLLMClient(correction=correction))
    result = run(ext, FRAGMENTED_GRID, page_image(tmp_path))
    assert result is not None
    assert "—" in result.markdown


def test_correction_skipped_when_disabled(tmp_path: Path):
    ext = make_extractor(llm=MockLLMClient(), run_llm_correction=False)
    assert run(ext, FRAGMENTED_GRID, page_image(tmp_path)) is None


def test_correction_skipped_when_no_llm(tmp_path: Path):
    ext = make_extractor(llm=None)
    assert run(ext, FRAGMENTED_GRID, page_image(tmp_path)) is None


def test_correction_skipped_without_page_image():
    ext = make_extractor(llm=MockLLMClient())
    assert run(ext, FRAGMENTED_GRID, None) is None


def test_correction_skipped_on_empty_grid(tmp_path: Path):
    ext = make_extractor(llm=MockLLMClient())
    assert run(ext, [], page_image(tmp_path)) is None


def test_correction_failure_falls_back_to_grid(tmp_path: Path, monkeypatch: MonkeyPatch):
    patch_io(monkeypatch)
    llm = MockLLMClient()

    async def boom(image_png: bytes, markdown: str, region_text: str) -> Optional[LLMTableCorrection]:
        _ = (image_png, markdown, region_text)
        raise RuntimeError("llm exploded")

    monkeypatch.setattr(llm, "vet_structure", boom)
    ext = make_extractor(llm=llm)

    # A raised vetting call is caught and reported as no correction, so the
    # caller keeps the deterministic grid markdown.
    assert run(ext, FRAGMENTED_GRID, page_image(tmp_path)) is None


# --- make_extracted ---------------------------------------------------------


def test_make_extracted_without_correction_uses_grid_defaults():
    ext = make_extractor()
    cand = candidate(FRAGMENTED_GRID)
    detected = DetectedTable(ordinal=1, description="Branded volume table")

    table = ext.make_extracted(
        cand=cand,
        cells=FRAGMENTED_GRID,
        detected=detected,
        flavor="lattice",
        source="doc.pdf",
        source_ids=[cand.candidate_id],
        status="extracted",
        correction=None,
    )

    assert table.title == "Branded volume table"
    assert table.subtitle == ""
    assert table.footnotes == []
    assert table.markdown == grid_to_markdown(FRAGMENTED_GRID)
    assert table.llm_corrected is False


def test_make_extracted_with_correction_populates_fields():
    ext = make_extractor()
    cand = candidate(FRAGMENTED_GRID)
    detected = DetectedTable(ordinal=1, description="detector description")
    correction = StructureCorrection(
        title="Branded Volume and Transactions",
        caption="For the 3 Months Ended September 30, 2025",
        footnotes=[FootnoteDef(marker="(1)", text="Preliminary.")],
        markdown=MERGED_MARKDOWN,
        llm_corrected=True,
        header_rows=1,
    )

    table = ext.make_extracted(
        cand=cand,
        cells=FRAGMENTED_GRID,
        detected=detected,
        flavor="lattice",
        source="doc.pdf",
        source_ids=[cand.candidate_id],
        status="extracted",
        correction=correction,
    )

    assert table.title == "Branded Volume and Transactions"
    assert table.caption == "For the 3 Months Ended September 30, 2025"
    assert table.footnotes == [FootnoteDef(marker="(1)", text="Preliminary.")]
    assert table.markdown == MERGED_MARKDOWN
    assert table.llm_corrected is True


def test_make_extracted_empty_llm_title_falls_back_to_detector():
    ext = make_extractor()
    cand = candidate(FRAGMENTED_GRID)
    detected = DetectedTable(ordinal=1, description="detector description")
    correction = StructureCorrection(
        title="",
        caption="",
        footnotes=[],
        markdown=MERGED_MARKDOWN,
        llm_corrected=True,
        header_rows=1,
    )

    table = ext.make_extracted(
        cand=cand,
        cells=FRAGMENTED_GRID,
        detected=detected,
        flavor="lattice",
        source="doc.pdf",
        source_ids=[cand.candidate_id],
        status="extracted",
        correction=correction,
    )

    assert table.title == "detector description"


if __name__ == "__main__":
    raise SystemExit(pytest.main([__file__, "-v"]))


# --- the grounding guard ----------------------------------------------------


def test_a_figure_keys_the_same_however_it_was_rendered():
    from quber.core.extractors.camelot.correspondence.correction import numeric_keys

    assert numeric_keys("$ 8,273.04") == numeric_keys("8273.04") == numeric_keys("8,273.04%")


def test_a_value_neither_the_grid_nor_the_page_carries_is_ungrounded():
    from quber.core.extractors.camelot.correspondence.correction import ungrounded_values

    assert ungrounded_values(
        corrected="| Cash | 99,999,999 |",
        grid_md="| Cash | 21,249,107 |",
        page_text="Cash 21,249,107",
    ) == {"99999999"}


def test_a_value_the_page_carries_is_grounded_even_where_the_grid_missed_it():
    from quber.core.extractors.camelot.correspondence.correction import ungrounded_values

    assert (
        ungrounded_values(
            corrected="| Cash | 21,249,107 |",
            grid_md="| Cash | |",
            page_text="Cash 21,249,107",
        )
        == set()
    )


def test_a_table_read_off_an_image_is_not_value_checked(tmp_path: Path, monkeypatch: MonkeyPatch):
    """The page prints no text under such a table, so the only thing the check
    could compare against is the grid it was asked to correct. A thousands
    separator restored from a misread has to stand, and the flag is what decides."""
    patch_io(monkeypatch)
    repaired = MERGED_MARKDOWN + "\n| ZZ | 19,543,903 | |"
    llm = MockLLMClient(
        correction=LLMTableCorrection(title="", caption="", markdown=repaired, footnotes=[], header_rows=1)
    )

    def correct(ground_values: bool):
        return asyncio.run(
            correction_module.correct_structure(
                FRAGMENTED_GRID,
                page_image(tmp_path),
                BBOX,
                SOURCE,
                PAGE,
                llm,
                asyncio.Semaphore(1),
                200,
                ground_values,
            )
        )

    kept = correct(ground_values=False)
    assert kept is not None
    assert "19,543,903" in kept.markdown
    assert correct(ground_values=True) is None