"""Tests for the docling-spine / Camelot reconciliation.

The matcher and graft are exercised against synthetic DoclingDocuments built with
docling-core's own API, so the box-frame conversions and the docling-core
mutation calls are tested for real (no docling pipeline run).
"""

from __future__ import annotations

from typing import List, Literal, Optional, Sequence, Tuple

import pytest
from docling_core.types.doc.base import BoundingBox, CoordOrigin
from docling_core.types.doc.document import (
    DoclingDocument,
    PictureClassificationClass,
    PictureClassificationData,
    ProvenanceItem,
    TableData,
)
from docling_core.types.doc.labels import DocItemLabel

from quber.agents.llm_client import FootnoteDef
from quber.core.extractors.base import ExtractedTable
from quber.core.fusion import build_unified_document, markdown_to_table_data, match_tables
from quber.core.parsers.result import PageParse, ParseResult, TableProvenance

PAGE_W, PAGE_H = 612.0, 792.0
NormBox = Tuple[float, float, float, float]

SAMPLE_MD = "| A | B |\n| --- | --- |\n| 1 | 2 |\n| 3 | 4 |"


def _norm_to_bottomleft(box: NormBox) -> BoundingBox:
    x1, y1, x2, y2 = box
    return BoundingBox(
        l=x1 * PAGE_W,
        r=x2 * PAGE_W,
        t=(1.0 - y1) * PAGE_H,
        b=(1.0 - y2) * PAGE_H,
        coord_origin=CoordOrigin.BOTTOMLEFT,
    )


def _make_parse(
    tables: Sequence[Tuple[NormBox, Literal["native", "ocr", "empty"]]] = (),
    pictures: Sequence[Tuple[NormBox, List[str]]] = (),
) -> ParseResult:
    """A ParseResult with docling tables (box, verdict) and pictures (box, classes)."""
    doc = DoclingDocument(name="synthetic")
    provenance: List[TableProvenance] = []
    for box, verdict in tables:
        prov = ProvenanceItem(page_no=1, bbox=_norm_to_bottomleft(box), charspan=(0, 0))
        item = doc.add_table(data=TableData(table_cells=[], num_rows=0, num_cols=0), prov=prov)
        provenance.append(
            TableProvenance(
                self_ref=item.self_ref,
                page_no=1,
                cell_count=10,
                from_ocr_fraction=1.0 if verdict == "ocr" else 0.0,
                mean_ocr_confidence=0.98 if verdict == "ocr" else 0.0,
                verdict=verdict,
            )
        )
    for box, classes in pictures:
        prov = ProvenanceItem(page_no=1, bbox=_norm_to_bottomleft(box), charspan=(0, 0))
        pic = doc.add_picture(prov=prov)
        if classes:
            # The live pipeline populates meta.classification; this exercises the
            # deprecated-annotations fallback path that _picture_classes also reads.
            pic.annotations.append(
                PictureClassificationData(
                    provenance="test",
                    predicted_classes=[
                        PictureClassificationClass(class_name=c, confidence=0.9 - i * 0.1)
                        for i, c in enumerate(classes)
                    ],
                )
            )
    pages = [PageParse(page_no=1, width=PAGE_W, height=PAGE_H, cells=[])]
    return ParseResult(document=doc, page_scores=[], pages=pages, table_provenance=provenance)


def _som(box: NormBox, markdown: str = SAMPLE_MD, accuracy: float = 95.0, title: str = "") -> ExtractedTable:
    return ExtractedTable(
        markdown=markdown,
        page=1,
        som_region=box,
        camelot_accuracy=accuracy,
        title=title,
        source="x.pdf",
    )


def _kind_of(parse: ParseResult, soms: List[ExtractedTable], som_idx: Optional[int] = None) -> str:
    matches = match_tables(parse, soms)
    if som_idx is None:
        return matches[0].kind
    for m in matches:
        if som_idx in m.som_indices:
            return m.kind
    return "none"


# --- markdown_to_table_data ------------------------------------------------


def test_markdown_to_table_data_grid_and_header():
    data = markdown_to_table_data(SAMPLE_MD)
    assert data.num_rows == 3
    assert data.num_cols == 2
    header = [c for c in data.table_cells if c.column_header]
    assert {c.text for c in header} == {"A", "B"}
    body = {(c.start_row_offset_idx, c.start_col_offset_idx): c.text for c in data.table_cells}
    assert body[(1, 0)] == "1" and body[(2, 1)] == "4"


def test_markdown_to_table_data_pads_ragged_rows():
    data = markdown_to_table_data("| A | B | C |\n| --- | --- | --- |\n| 1 |")
    assert data.num_cols == 3
    assert len([c for c in data.table_cells if c.start_row_offset_idx == 1]) == 3


def test_markdown_to_table_data_empty():
    data = markdown_to_table_data("")
    assert data.num_rows == 0 and data.table_cells == []


# --- matcher classification ------------------------------------------------


def test_match_replace_one_to_one():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    assert _kind_of(parse, [_som(box)]) == "replace"


def test_match_docling_miss():
    parse = _make_parse(tables=[((0.1, 0.1, 0.9, 0.4), "native")])
    # SoM table far from the docling table -> no overlap.
    assert _kind_of(parse, [_som((0.1, 0.6, 0.9, 0.9))], som_idx=0) == "docling_miss"


def test_match_chart_over_picture():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(pictures=[(box, ["bar_chart"])])
    matches = match_tables(parse, [_som(box)])
    assert matches[0].kind == "chart"
    assert matches[0].picture_classes == ["bar_chart"]


def test_match_image_table_no_text_layer():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "ocr")])
    # Camelot empty (image table): no markdown, zero accuracy.
    assert _kind_of(parse, [_som(box, markdown="", accuracy=0.0)]) == "image_table"


def test_match_docling_undercount():
    region = (0.1, 0.1, 0.9, 0.5)
    parse = _make_parse(tables=[(region, "native")])
    # Two SoM tables tile one docling table -> Camelot found more.
    soms = [_som((0.1, 0.12, 0.9, 0.28)), _som((0.1, 0.32, 0.9, 0.48))]
    assert _kind_of(parse, soms, som_idx=0) == "docling_undercount"


def test_match_som_merged():
    big = (0.1, 0.1, 0.9, 0.5)
    parse = _make_parse(tables=[((0.1, 0.12, 0.9, 0.28), "native"), ((0.1, 0.32, 0.9, 0.48), "native")])
    # One SoM table spans two docling tables -> docling found more.
    assert _kind_of(parse, [_som(big)]) == "som_merged"


def test_match_som_miss_reported():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    matches = match_tables(parse, [])  # no SoM tables at all
    assert [m.kind for m in matches] == ["som_miss"]


# --- graft -----------------------------------------------------------------


def test_graft_replace_rebuilds_body_and_keeps_source_pure():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    soms = [_som(box)]
    matches = match_tables(parse, soms)
    page_dims = {1: (PAGE_W, PAGE_H)}

    unified, errors = build_unified_document(parse.document, soms, matches, page_dims)
    assert errors == []
    # Source document untouched: still the empty 0x0 table.
    assert parse.document.tables[0].data.num_rows == 0
    # Unified table now carries the Camelot body.
    assert unified.tables[0].data.num_rows == 3
    assert unified.tables[0].data.num_cols == 2


def test_graft_carries_units_and_footnotes_onto_table():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    som = _som(box)
    som.units = "($ in millions)"
    som.footnotes = [
        FootnoteDef(marker="(1)", text="Restated for FX."),
        FootnoteDef(marker="(2)", text="Unaudited."),
    ]
    matches = match_tables(parse, [som])
    unified, errors = build_unified_document(parse.document, [som], matches, {1: (PAGE_W, PAGE_H)})
    assert errors == []

    table = unified.tables[0]
    assert table.caption_text(unified) == "($ in millions)"
    # Footnotes are body nodes after the table (docling's markdown ignores
    # table.footnotes refs), so they render in the exported markdown.
    md = unified.export_to_markdown()
    assert "(1) Restated for FX." in md
    assert "(2) Unaudited." in md


def _caption_texts(unified: DoclingDocument, table) -> List[str]:
    return [unified.texts[int(r.cref.split("/")[-1])].text for r in table.captions]


def test_graft_carries_caption_title_units_as_ordered_captions():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    som = _som(box, title="Visa Debit Programs")
    som.caption = "The following table presents payment volume by program:"
    som.units = "($ in millions)"
    matches = match_tables(parse, [som])
    unified, _ = build_unified_document(parse.document, [som], matches, {1: (PAGE_W, PAGE_H)})
    assert _caption_texts(unified, unified.tables[0]) == [
        "The following table presents payment volume by program:",
        "Visa Debit Programs",
        "($ in millions)",
    ]


def test_graft_older_artifact_subtitle_still_carried():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    som = _som(box, title="For the 3 Months Ended December 31, 2025")
    som.subtitle = "Visa Debit Programs"
    matches = match_tables(parse, [som])
    unified, _ = build_unified_document(parse.document, [som], matches, {1: (PAGE_W, PAGE_H)})
    assert _caption_texts(unified, unified.tables[0]) == [
        "For the 3 Months Ended December 31, 2025",
        "Visa Debit Programs",
    ]


def test_graft_caption_takes_the_whole_paragraph_for_both_side_by_side_tables():
    # One sentence introduces two tables printed side by side. Each vetting
    # call copied a different part of it; the graft attaches the docling
    # paragraph that contains each copy, so both tables carry the same whole
    # sentence, bold lead-in included.
    left, right = (0.05, 0.2, 0.48, 0.6), (0.52, 0.2, 0.95, 0.6)
    parse = _make_parse(tables=[(left, "native"), (right, "native")])
    sentence = (
        "Concentration of Credit Risk - The following tables present the geographies and "
        "property types of collateral underlying the loans' principal amounts:"
    )
    parse.document.add_text(
        label=DocItemLabel.TEXT,
        text=sentence,
        prov=ProvenanceItem(page_no=1, bbox=_norm_to_bottomleft((0.05, 0.12, 0.95, 0.15)), charspan=(0, 0)),
    )
    som_left = _som(left)
    som_left.caption = "The following tables present the geographies and property types of collateral underlying the loans' principal amounts:"
    som_right = _som(right)
    som_right.caption = "Concentration of Credit Risk — The following tables present the geographies"
    matches = match_tables(parse, [som_left, som_right])
    unified, errors = build_unified_document(
        parse.document, [som_left, som_right], matches, {1: (PAGE_W, PAGE_H)}
    )
    assert errors == []
    for table in unified.tables:
        assert _caption_texts(unified, table) == [sentence]


def test_graft_caption_stands_when_no_paragraph_holds_it():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    som = _som(box)
    som.caption = "The table below sets forth additional information relating to our portfolio:"
    matches = match_tables(parse, [som])
    unified, _ = build_unified_document(parse.document, [som], matches, {1: (PAGE_W, PAGE_H)})
    assert _caption_texts(unified, unified.tables[0]) == [som.caption]


def test_graft_never_writes_page_furniture_as_a_header_line():
    # "Table of Contents" is printed at the top of every page, so a copied
    # caption or title of that text passes the printed-text check; the parse
    # labels it a page header, and that label rules it out.
    box = (0.1, 0.2, 0.9, 0.5)
    parse = _make_parse(tables=[(box, "native")])
    parse.document.add_text(
        label=DocItemLabel.PAGE_HEADER,
        text="Table of Contents",
        prov=ProvenanceItem(page_no=1, bbox=_norm_to_bottomleft((0.1, 0.01, 0.3, 0.03)), charspan=(0, 0)),
    )
    som = _som(box, title="Table of Contents")
    som.caption = "Table of Contents"
    som.units = "($ in millions)"
    matches = match_tables(parse, [som])
    unified, _ = build_unified_document(parse.document, [som], matches, {1: (PAGE_W, PAGE_H)})
    assert _caption_texts(unified, unified.tables[0]) == ["($ in millions)"]


def test_graft_title_that_repeats_the_caption_is_written_once():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    som = _som(box, title="Activity")
    som.caption = (
        "Activity — For the three months ended March 31, 2026, the loan portfolio activity was as follows:"
    )
    matches = match_tables(parse, [som])
    unified, _ = build_unified_document(parse.document, [som], matches, {1: (PAGE_W, PAGE_H)})
    assert _caption_texts(unified, unified.tables[0]) == [som.caption]


def test_graft_dedups_in_region_text_keeps_spine():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    doc = parse.document

    def _text(label, text, norm_box):
        return doc.add_text(
            label=label,
            text=text,
            prov=ProvenanceItem(page_no=1, bbox=_norm_to_bottomleft(norm_box), charspan=(0, 0)),
        )

    _text(DocItemLabel.TEXT, "$539", (0.2, 0.2, 0.3, 0.25))  # loose cell fragment, in region
    _text(DocItemLabel.TEXT, "(millions)", (0.4, 0.15, 0.5, 0.17))  # in-region label
    _text(DocItemLabel.SECTION_HEADER, "1. Branded Volume", (0.1, 0.11, 0.6, 0.13))  # in region but spine
    _text(DocItemLabel.TEXT, "Footnote section body", (0.1, 0.6, 0.9, 0.65))  # out of region

    som = _som(box, markdown="| Volume | 2025 |\n| --- | --- |\n| Branded | $539 |")
    som.units = "(millions)"
    matches = match_tables(parse, [som])
    unified, _ = build_unified_document(doc, [som], matches, {1: (PAGE_W, PAGE_H)})
    # Only docling's own items carry a page box; the attribution nodes the graft
    # attaches (units, footnotes) have none and are the table's representation.
    texts = [t.text for t in unified.texts if t.prov]
    assert "$539" not in texts  # in-region cell fragment removed
    assert "(millions)" not in texts  # in-region label removed
    assert "1. Branded Volume" in texts  # section header kept even in region
    assert "Footnote section body" in texts  # out-of-region text kept


def test_graft_dedup_keeps_in_region_text_the_table_does_not_carry():
    # The locator's region is drawn past the grid to enclose the footnote lines,
    # so prose printed under the table sits inside it too. Only text the grafted
    # table carries is a duplicate; position alone never deletes.
    box = (0.1, 0.1, 0.9, 0.5)
    parse = _make_parse(tables=[(box, "native")])
    doc = parse.document

    def _text(label, text, norm_box):
        return doc.add_text(
            label=label,
            text=text,
            prov=ProvenanceItem(page_no=1, bbox=_norm_to_bottomleft(norm_box), charspan=(0, 0)),
        )

    _text(DocItemLabel.TEXT, "(1) Includes the impact of interest rate floors.", (0.1, 0.36, 0.9, 0.38))
    # The table's copy of a footnote cut at the page break; docling's full text is longer.
    _text(
        DocItemLabel.TEXT,
        "(2) Net of deferred origination fees and other items as of June 30, 2026, "
        "recognized in unrealized gain (loss) on securities.",
        (0.1, 0.39, 0.9, 0.41),
    )
    # A rendering difference of one token is still the same printed line.
    _text(DocItemLabel.TEXT, "First mortgage loans’ balance", (0.1, 0.2, 0.5, 0.22))
    _text(
        DocItemLabel.TEXT,
        "As of June 30, 2026, $2.5 billion, or 88.9%, of the outstanding face amount "
        "were at variable interest rates linked to Term SOFR.",
        (0.1, 0.43, 0.9, 0.47),
    )

    som = _som(
        box, markdown="| | Face Amount |\n| --- | --- |\n| First mortgage loans' balance | $ 2,807,310 |"
    )
    som.footnotes = [
        FootnoteDef(marker="(1)", text="Includes the impact of interest rate floors."),
        FootnoteDef(
            marker="(2)", text="Net of deferred origination fees and other items as of June 30, 2026,"
        ),
    ]
    matches = match_tables(parse, [som])
    unified, _ = build_unified_document(doc, [som], matches, {1: (PAGE_W, PAGE_H)})
    texts = [t.text for t in unified.texts if t.prov]
    assert not any(t.startswith("(1) Includes") for t in texts)  # footnote the table carries
    assert not any(t.startswith("First mortgage") for t in texts)  # one-token apostrophe noise
    assert any(t.startswith("(2) Net of") for t in texts)  # longer than the table's copy: kept
    assert any(t.startswith("As of June 30, 2026, $2.5 billion") for t in texts)  # prose: kept


def test_graft_no_attribution_when_som_has_none():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    som = _som(box)  # no units, no footnotes
    matches = match_tables(parse, [som])
    unified, _ = build_unified_document(parse.document, [som], matches, {1: (PAGE_W, PAGE_H)})
    assert unified.tables[0].captions == []
    assert unified.tables[0].footnotes == []


def test_graft_insert_adds_missed_table():
    matched = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(matched, "native")])
    # One SoM table matches the docling table (replace); one is in a region
    # docling missed (insert).
    soms = [_som(matched), _som((0.1, 0.6, 0.9, 0.9))]
    matches = match_tables(parse, soms)
    unified, errors = build_unified_document(parse.document, soms, matches, {1: (PAGE_W, PAGE_H)})
    assert errors == []
    assert len(unified.tables) == 2  # replaced docling table + inserted miss


def test_graft_insert_empty_body_is_surfaced():
    parse = _make_parse(tables=[((0.1, 0.1, 0.9, 0.4), "native")])
    soms = [_som((0.1, 0.6, 0.9, 0.9), markdown="", accuracy=0.0, title="Lost table")]
    matches = match_tables(parse, soms)
    unified, errors = build_unified_document(parse.document, soms, matches, {1: (PAGE_W, PAGE_H)})
    assert len(unified.tables) == 1  # nothing inserted
    assert any("Lost table" in e for e in errors)


def test_graft_som_miss_surfaced_as_error():
    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    matches = match_tables(parse, [])
    unified, errors = build_unified_document(parse.document, [], matches, {1: (PAGE_W, PAGE_H)})
    assert len(unified.tables) == 1  # untouched
    assert any("no SoM match" in e for e in errors)


def test_graft_rewrites_report_refs_to_unified_numbering():
    # An undercount graft deletes the merged docling table and appends the two
    # finer SoM tables, renumbering every later table. The report must come out
    # referencing the unified document, not the pre-graft one: a consumer that
    # resolves a stale ref lands on a neighboring table and inherits its
    # identity (section heading, reading-order position).
    merged = (0.1, 0.05, 0.9, 0.35)
    kept = (0.1, 0.5, 0.9, 0.7)
    missed = (0.1, 0.8, 0.9, 0.95)
    parse = _make_parse(tables=[(merged, "native"), (kept, "native"), (missed, "native")])
    soms = [
        _som((0.1, 0.07, 0.9, 0.19)),  # upper half of the merged region
        _som((0.1, 0.21, 0.9, 0.33)),  # lower half of the merged region
        _som(kept),  # 1:1 replace
    ]
    matches = match_tables(parse, soms)
    by_kind = {m.kind: m for m in matches}
    assert set(by_kind) == {"docling_undercount", "replace", "som_miss"}
    assert by_kind["docling_undercount"].docling_table_refs == ["#/tables/0"]
    assert by_kind["replace"].docling_table_refs == ["#/tables/1"]
    assert by_kind["som_miss"].docling_table_refs == ["#/tables/2"]

    unified, _ = build_unified_document(parse.document, soms, matches, {1: (PAGE_W, PAGE_H)})

    # Deleting the merged table shifted the survivors down; the two inserted
    # tables took the end of the array.
    assert [t.self_ref for t in unified.tables] == [f"#/tables/{i}" for i in range(4)]
    by_ref = {t.self_ref: t for t in unified.tables}

    def _assert_resolves_to(match, boxes: List[NormBox]) -> None:
        assert len(match.docling_table_refs) == len(boxes)
        for ref, box in zip(match.docling_table_refs, boxes, strict=False):
            got = by_ref[ref].prov[0].bbox
            want = _norm_to_bottomleft(box)
            assert (got.l, got.t, got.r, got.b) == pytest.approx((want.l, want.t, want.r, want.b))

    _assert_resolves_to(by_kind["replace"], [kept])
    _assert_resolves_to(by_kind["som_miss"], [missed])
    _assert_resolves_to(by_kind["docling_undercount"], [(0.1, 0.07, 0.9, 0.19), (0.1, 0.21, 0.9, 0.33)])


# --- cell provenance and identity binding ------------------------------------


def _grounded(rows):
    """A cell_grid / corrected_grid from [(text, box), ...] rows; boxes normalized top-left."""
    from quber.core.extractors.base import GroundedCell

    return [[GroundedCell(text=t, box=b) for t, b in row] for row in rows]


def test_graft_binds_identity_and_carries_resolved_cell_geometry():
    # The graft is pure carriage: it reads the corrected_grid the grounding
    # stage resolved and writes meta + per-cell bboxes; no matching happens here.
    import json

    box = (0.1, 0.1, 0.9, 0.4)
    parse = _make_parse(tables=[(box, "native")])
    som = _som(box)
    som.table_id = "67609452-p1-t1"
    som.content_fingerprint = "ab12cd34"
    som.corrected_grid = _grounded(
        [
            [("A", (0.10, 0.12, 0.20, 0.14)), ("B", (0.50, 0.12, 0.60, 0.14))],
            [("1", (0.10, 0.20, 0.20, 0.22)), ("2", None)],  # one honest gap
            [("3", (0.10, 0.30, 0.20, 0.32)), ("4", (0.50, 0.30, 0.60, 0.32))],
        ]
    )
    matches = match_tables(parse, [som])
    unified, errors = build_unified_document(parse.document, [som], matches, {1: (PAGE_W, PAGE_H)})
    assert errors == []

    table = unified.tables[0]
    # Identity rides docling's own meta extension point.
    assert table.meta is not None
    dumped = table.meta.model_dump()
    assert dumped["quber__table_id"] == "67609452-p1-t1"
    assert dumped["quber__content_fingerprint"] == "ab12cd34"

    # Boxes are carried exactly, converted to bottom-left points; the grounding
    # stage's gap stays a gap.
    by_addr = {(c.start_row_offset_idx, c.start_col_offset_idx): c for c in table.data.table_cells}
    cell = by_addr[(1, 0)]  # "1"
    assert cell.bbox is not None
    assert cell.bbox.coord_origin == CoordOrigin.BOTTOMLEFT
    assert cell.bbox.l == 0.10 * PAGE_W
    assert cell.bbox.t == (1.0 - 0.20) * PAGE_H
    assert by_addr[(1, 1)].bbox is None

    # The binding survives a full serialize/deserialize round trip.
    payload = unified.model_dump(mode="json")
    revived = DoclingDocument.model_validate(payload)
    assert revived.tables[0].meta is not None
    assert revived.tables[0].meta.model_dump()["quber__table_id"] == "67609452-p1-t1"
    assert "67609452-p1-t1" in json.dumps(payload)