"""Tests for the over-merged Set-of-Mark region split post-pass.

`split_table` refines one captured table into the separate tables it fused,
gated on agreement between a supplied count and an internal vision count probe.
These tests exercise the geometry (`subdivide`) directly and drive `split_table`
with a mock LLM and monkeypatched render / Camelot / correction I/O, so no real
PDF, rendering, Camelot, or network is needed.
"""

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

import pytest
from _pytest.monkeypatch import MonkeyPatch

from quber.agents.llm_client import FootnoteDef, MockLLMClient
from quber.core.extractors.base import ExtractedTable
from quber.core.extractors.camelot.acquire import CamelotCandidate
from quber.core.extractors.camelot.correspondence.correction import StructureCorrection
from quber.core.extractors.set_of_mark import split as split_module
from quber.core.extractors.set_of_mark.split import split_table, subdivide

NormBox = Tuple[float, float, float, float]

REGION: NormBox = (0.1, 0.2, 0.9, 0.8)
# Two tight table boxes stacked inside REGION; each leaves a blank gap below it
# where a footnote sits. The seam must land between them so neither footnote is
# cut off from its table.
BOX_TOP: NormBox = (0.1, 0.25, 0.9, 0.40)
BOX_BOTTOM: NormBox = (0.1, 0.55, 0.9, 0.70)


def fused_table() -> ExtractedTable:
    return ExtractedTable(
        table_id="VISA_991_Q126-p11-t3",
        title="Fused",
        markdown="| a | b |\n|---|---|\n| 1 | 2 |",
        page=11,
        bbox=(72.0, 100.0, 540.0, 700.0),
        source="documents/VISA_991_Q126.pdf",
        flavor="stream",
        som_region=REGION,
        camelot_accuracy=97.0,
    )


def run(coro):
    return asyncio.run(coro)


# --- subdivide: deterministic geometry --------------------------------------


def test_subdivide_tiles_region_into_full_width_bands() -> None:
    subs = subdivide(REGION, [BOX_TOP, BOX_BOTTOM])
    assert subs is not None
    assert len(subs) == 2
    # Every band spans the region's full width.
    for x1, _y1, x2, _y2 in subs:
        assert (x1, x2) == (0.1, 0.9)


def test_subdivide_union_equals_original_region() -> None:
    subs = subdivide(REGION, [BOX_TOP, BOX_BOTTOM])
    assert subs is not None
    # The bands abut with no gap and no overlap, from the region top to its bottom.
    assert subs[0][1] == pytest.approx(0.2)
    assert subs[-1][3] == pytest.approx(0.8)
    assert subs[0][3] == pytest.approx(subs[1][1])


def test_subdivide_seam_lands_in_the_blank_gap_between_boxes() -> None:
    subs = subdivide(REGION, [BOX_TOP, BOX_BOTTOM])
    assert subs is not None
    seam = subs[0][3]
    # The seam sits in the gap, below the top box and above the bottom box, so a
    # footnote under the top box (~0.42) stays in the first band, not the second.
    assert 0.40 < seam < 0.55
    assert seam == pytest.approx(0.475)


def test_subdivide_orders_unsorted_boundaries_top_to_bottom() -> None:
    subs = subdivide(REGION, [BOX_BOTTOM, BOX_TOP])
    assert subs is not None
    assert subs[0][3] == pytest.approx(0.475)


def test_subdivide_rejects_boundary_centered_outside_region() -> None:
    outside = (0.1, 0.85, 0.9, 0.95)
    assert subdivide(REGION, [BOX_TOP, outside]) is None


# --- split_table: cheap no-op gates (no rendering, no LLM) -------------------


def test_single_table_count_is_passthrough() -> None:
    table = fused_table()
    out = run(split_table(table, count=1, boundaries=[BOX_TOP], llm=MockLLMClient()))
    assert out == [table]


def test_count_without_matching_boundaries_is_passthrough() -> None:
    table = fused_table()
    out = run(split_table(table, count=2, boundaries=[BOX_TOP], llm=MockLLMClient()))
    assert out == [table]


def test_missing_source_is_passthrough() -> None:
    table = fused_table()
    table.source = None
    out = run(split_table(table, count=2, boundaries=[BOX_TOP, BOX_BOTTOM], llm=MockLLMClient()))
    assert out == [table]


def test_missing_region_is_passthrough() -> None:
    table = fused_table()
    table.som_region = None
    out = run(split_table(table, count=2, boundaries=[BOX_TOP, BOX_BOTTOM], llm=MockLLMClient()))
    assert out == [table]


# --- split_table: confluence gate and the happy path ------------------------


def patch_render_and_crop(monkeypatch: MonkeyPatch) -> None:
    """Stub page render and region crop so the split path runs without a PDF."""

    def fake_render(source: Path, page: int, dpi: int, out_dir: Path) -> Tuple[Path, float, float]:
        img = out_dir / f"page-{page:04d}.png"
        img.write_bytes(b"\x89PNG")
        return img, 612.0, 792.0

    monkeypatch.setattr(split_module, "render_page", fake_render)
    monkeypatch.setattr(split_module, "crop_region_png", lambda page_image, region, dpi: b"\x89PNG")


def patch_reextract(monkeypatch: MonkeyPatch, *, fail_ordinal: Optional[int] = None) -> None:
    """Stub Camelot replay and structure correction per sub-region.

    `fail_ordinal` makes that sub-region's Camelot replay find no grid, so the
    split is abandoned and the fused region is kept.
    """

    def fake_camelot(source_str: str, page: int, area: str, ordinal: int) -> Optional[CamelotCandidate]:
        if fail_ordinal is not None and ordinal == fail_ordinal:
            return None
        return CamelotCandidate(
            candidate_id=f"r-{ordinal}",
            flavor="stream",
            page=page,
            bbox=(72.0, 100.0, 540.0, 300.0),
            accuracy=98.0,
            cells=[["Region", "2025"], ["US", "1,637"]],
            markdown="| Region | 2025 |\n|---|---|\n| US | 1,637 |",
        )

    async def fake_correct(cells, page_image, bbox, source, page, llm, sem, dpi) -> StructureCorrection:
        return StructureCorrection(
            title=f"Reconciliation {page}-{bbox[1]:.0f}",
            caption="Three Months Ended December 31, 2025",
            footnotes=[FootnoteDef(marker="(1)", text="See note.")],
            markdown="| Region | 2025 |\n|---|---|\n| US | 1,637 |",
            llm_corrected=True,
            header_rows=1,
        )

    monkeypatch.setattr(split_module, "camelot_targeted", fake_camelot)
    monkeypatch.setattr(split_module, "correct_structure", fake_correct)


def test_vision_disagreement_keeps_camelot(monkeypatch: MonkeyPatch) -> None:
    patch_render_and_crop(monkeypatch)
    patch_reextract(monkeypatch)
    table = fused_table()
    # Supplied count 2, vision sees 1 -> disagreement -> region stands.
    out = run(split_table(table, count=2, boundaries=[BOX_TOP, BOX_BOTTOM], llm=MockLLMClient(table_count=1)))
    assert out == [table]


def test_agreement_on_two_splits_into_two_subregions(monkeypatch: MonkeyPatch) -> None:
    patch_render_and_crop(monkeypatch)
    patch_reextract(monkeypatch)
    table = fused_table()
    out: List[ExtractedTable] = run(
        split_table(table, count=2, boundaries=[BOX_TOP, BOX_BOTTOM], llm=MockLLMClient(table_count=2))
    )
    assert len(out) == 2
    # Each emitted table carries its own sub-region; together they tile REGION.
    assert out[0].som_region == pytest.approx((0.1, 0.2, 0.9, 0.475))
    assert out[1].som_region == pytest.approx((0.1, 0.475, 0.9, 0.8))
    # Per-table title, caption and footnote are recovered, values kept.
    assert out[0].caption == "Three Months Ended December 31, 2025"
    assert out[0].footnotes == [FootnoteDef(marker="(1)", text="See note.")]
    assert all("1,637" in t.markdown for t in out)
    # Each sub-table extends the parent's address with its own sub-index, and
    # carries the fingerprint of its own re-extracted grid.
    assert out[0].table_id == "VISA_991_Q126-p11-t3-s1"
    assert out[1].table_id == "VISA_991_Q126-p11-t3-s2"
    assert all(len(t.content_fingerprint) == 8 for t in out)


def test_units_forward_fill_across_split_siblings(monkeypatch: MonkeyPatch) -> None:
    # Only the top sub-table's crop sees the group caption; the bottom one comes
    # back empty. The forward fill must carry the top caption down to it.
    patch_render_and_crop(monkeypatch)

    def fake_camelot(source_str: str, page: int, area: str, ordinal: int) -> Optional[CamelotCandidate]:
        # Distinct y per sub-region so the correction stub can tell them apart.
        y = 100.0 if ordinal == 1 else 400.0
        return CamelotCandidate(
            candidate_id=f"r-{ordinal}",
            flavor="stream",
            page=page,
            bbox=(72.0, y, 540.0, y + 200.0),
            accuracy=98.0,
            cells=[["Region", "2025"], ["US", "1,637"]],
            markdown="| Region | 2025 |\n|---|---|\n| US | 1,637 |",
        )

    async def fake_correct(cells, page_image, bbox, source, page, llm, sem, dpi) -> StructureCorrection:
        # The top sub-region (y==100) captured the caption; the lower one did not.
        units = "(in millions, except percentages and per share data)" if bbox[1] == 100.0 else ""
        return StructureCorrection(
            title=f"Reconciliation {bbox[1]:.0f}",
            caption="",
            footnotes=[],
            markdown="| Region | 2025 |\n|---|---|\n| US | 1,637 |",
            llm_corrected=True,
            units=units,
            header_rows=1,
        )

    monkeypatch.setattr(split_module, "camelot_targeted", fake_camelot)
    monkeypatch.setattr(split_module, "correct_structure", fake_correct)

    out = run(
        split_table(
            fused_table(), count=2, boundaries=[BOX_TOP, BOX_BOTTOM], llm=MockLLMClient(table_count=2)
        )
    )
    assert len(out) == 2
    caption = "(in millions, except percentages and per share data)"
    # The top captured it; the empty bottom sibling inherits it via the fill.
    assert out[0].units == caption
    assert out[1].units == caption


def test_failed_subregion_reextraction_keeps_fused_region(monkeypatch: MonkeyPatch) -> None:
    patch_render_and_crop(monkeypatch)
    patch_reextract(monkeypatch, fail_ordinal=2)
    table = fused_table()
    # Vision confirms the split, but the second sub-region yields no grid; rather
    # than ship a half table, the original fused region is returned unchanged.
    out = run(split_table(table, count=2, boundaries=[BOX_TOP, BOX_BOTTOM], llm=MockLLMClient(table_count=2)))
    assert out == [table]