"""Tests for the dpt-3 track: digestion, the corrected graft, footnote resolution.

The responses are synthetic but carry the shape a live dpt-3 scan returns — one
long markdown string and a structure tree whose nodes point at stretches of it
by codepoint range, boxes keyed xmin/ymin/xmax/ymax, metadata reporting
model_version and billing.total_credits. The builders assemble the string and
the ranges together so the two can never disagree.

The dpt-2 path's tests live in `test_figures.py` and are untouched — that is
the frozen-path guarantee. A final test digests real saved responses from live
runs when the local data directory holds them.
"""

from __future__ import annotations

import asyncio
import json
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import pytest
from docling_core.types.doc.base import BoundingBox, CoordOrigin, Size
from docling_core.types.doc.document import (
    DoclingDocument,
    PictureClassificationMetaField,
    PictureClassificationPrediction,
    PictureMeta,
    ProvenanceItem,
)
from docling_core.types.doc.labels import DocItemLabel

from quber.agents.figure_correction import Marker, MockFigureCorrector, Note
from quber.core.figures import dpt3
from quber.core.figures.correct import correct_figures
from quber.core.figures.dpt3.digest import UnrecognizedResponse, digest_page, figure_text, project_scan
from quber.core.figures.graft import FOOTNOTE_MARKS_FIELD, FOOTNOTES_FIELD
from quber.core.figures.models import PageScan
from quber.core.figures.orchestrator import _page_response
from quber.core.figures.scan import page_scan

PAGE_W, PAGE_H = 720.0, 540.0
NormBox = Tuple[float, float, float, float]

MODEL = "dpt-3-pro-latest"


# --- response builders -------------------------------------------------------


def _v2_box(box: NormBox) -> Dict[str, float]:
    x1, y1, x2, y2 = box
    return {"xmin": x1, "ymin": y1, "xmax": x2, "ymax": y2}


def _v2_figure(
    kind: str,
    description: str,
    rows: Sequence[Sequence[str]] = (),
    box: NormBox = (0.1, 0.2, 0.9, 0.8),
    node_id: str = "figure-0",
) -> Dict[str, Any]:
    table = ""
    if rows:
        table = (
            "<table>"
            + "".join("<tr>" + "".join(f"<td>{cell}</td>" for cell in row) + "</tr>" for row in rows)
            + "</table>"
        )
    content = f'<figure type="{kind}">\n\n<description>{description}</description>{table}</figure>'
    return {"type": "figure", "id": node_id, "content": content, "box": box}


def _v2_text(kind: str, text: str, box: NormBox, node_id: str = "text-0") -> Dict[str, Any]:
    return {"type": kind, "id": node_id, "content": text, "box": box}


def _v2_table(
    rows: Sequence[Sequence[Tuple[str, Optional[NormBox]]]],
    box: NormBox = (0.1, 0.2, 0.9, 0.8),
    node_id: str = "table-0",
    spans: Optional[Dict[Tuple[int, int], Tuple[int, int]]] = None,
) -> Dict[str, Any]:
    """A table node spec: rows of (cell text, cell box), spans keyed (row, col)."""
    return {"type": "table", "id": node_id, "rows": rows, "box": box, "spans": spans or {}}


def _v2_response(
    nodes: Sequence[Dict[str, Any]],
    credits: float = 0.8,
    page_status: str = "ok",
    page_reason: Optional[str] = None,
) -> Dict[str, Any]:
    """Assemble the markdown string and the tree together, ranges computed."""
    markdown = ""
    children: List[Dict[str, Any]] = []
    for spec in nodes:
        if markdown:
            markdown += "\n\n"
        start = len(markdown)
        if spec["type"] == "table":
            cell_nodes: List[Dict[str, Any]] = []
            markdown += "<table>\n"
            for r, row in enumerate(spec["rows"]):
                markdown += "<tr>"
                for c, (text, cbox) in enumerate(row):
                    markdown += "<td>"
                    cell_start = len(markdown)
                    markdown += text
                    rowspan, colspan = spec["spans"].get((r, c), (1, 1))
                    cell_nodes.append(
                        {
                            "type": "table_cell",
                            "id": f"{spec['id']}-cell-{r}-{c}",
                            "row": r,
                            "col": c,
                            "rowspan": rowspan,
                            "colspan": colspan,
                            "grounding": {
                                "page": 1,
                                "range": {"start": cell_start, "end": len(markdown)},
                                "box": _v2_box(cbox) if cbox else None,
                            },
                        }
                    )
                    markdown += "</td>"
                markdown += "</tr>\n"
            markdown += "</table>"
            children.append(
                {
                    "type": "table",
                    "id": spec["id"],
                    "grounding": {
                        "page": 1,
                        "range": {"start": start, "end": len(markdown)},
                        "box": _v2_box(spec["box"]),
                    },
                    "children": cell_nodes,
                }
            )
        else:
            markdown += spec["content"]
            children.append(
                {
                    "type": spec["type"],
                    "id": spec["id"],
                    "grounding": {
                        "page": 1,
                        "range": {"start": start, "end": len(markdown)},
                        "box": _v2_box(spec["box"]),
                    },
                    "children": [],
                }
            )
    return {
        "markdown": markdown,
        "structure": {
            "type": "document",
            "children": [
                {
                    "type": "page",
                    "status": page_status,
                    "reason": page_reason,
                    "grounding": {
                        "page": 1,
                        "range": {"start": 0, "end": len(markdown)},
                        "box": {"xmin": 0.0, "ymin": 0.0, "xmax": 1.0, "ymax": 1.0},
                    },
                    "children": children,
                }
            ],
        },
        "grounding": None,
        "metadata": {
            "job_id": "parse-01synthetic",
            "model_version": "dpt-3-pro-20260710",
            "page_count": 1,
            "billing": {"service_tier": "standard", "total_credits": credits},
            "failed_pages": [],
        },
    }


# --- document builders -------------------------------------------------------


def _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 _document(pictures: Sequence[Tuple[int, NormBox, str]]) -> DoclingDocument:
    doc = DoclingDocument(name="synthetic")
    for page in sorted({p for p, _box, _cls in pictures}):
        doc.add_page(page_no=page, size=Size(width=PAGE_W, height=PAGE_H))
    for page, box, cls in pictures:
        picture = doc.add_picture(prov=ProvenanceItem(page_no=page, bbox=_bottomleft(box), charspan=(0, 0)))
        if cls:
            picture.meta = PictureMeta(
                classification=PictureClassificationMetaField(
                    predictions=[PictureClassificationPrediction(class_name=cls, confidence=0.53)]
                )
            )
    return doc


def _page_dims(pages: Sequence[int]) -> Dict[int, Tuple[float, float]]:
    return dict.fromkeys(pages, (PAGE_W, PAGE_H))


def _blank_pdf(dest: Path, pages: int) -> Path:
    import fitz

    doc = fitz.open()
    for _ in range(pages):
        doc.new_page(width=PAGE_W, height=PAGE_H)
    doc.save(str(dest))
    doc.close()
    return dest


# --- digestion ---------------------------------------------------------------


def test_a_figure_node_digests_into_kind_description_and_value_table():
    response = _v2_response(
        [
            _v2_figure(
                "CHART",
                "Two pie charts of portfolio composition.",
                rows=[["Industrial: 26%", "Sunbelt: 25%"], ["Multifamily: 25%", "UK: 17%"]],
                box=(0.05, 0.2, 0.95, 0.85),
            )
        ]
    )
    digest = digest_page(response, 9, MODEL, "deck.p9.ade.json")

    assert digest.page == 9
    assert digest.version == "dpt-3-pro-20260710"
    assert digest.credits == 0.8
    figure = digest.figures[0]
    assert figure.kind == "chart"
    assert figure.description == "Two pie charts of portfolio composition."
    assert figure.box == {"left": 0.05, "top": 0.2, "right": 0.95, "bottom": 0.85}
    assert [[cell.text for cell in row] for row in figure.values[0].rows] == [
        ["Industrial: 26%", "Sunbelt: 25%"],
        ["Multifamily: 25%", "UK: 17%"],
    ]
    # Digestion classifies nothing: printed versus estimated is a later step.
    assert all(cell.provenance == "unclassified" for row in figure.values[0].rows for cell in row)


def test_a_description_only_figure_digests_without_values():
    response = _v2_response([_v2_figure("DIAGRAM", "A process flow with three stages.")])
    figure = digest_page(response, 20, MODEL).figures[0]
    assert figure.kind == "diagram"
    assert figure.values == []
    assert figure_text(figure) == "A process flow with three stages."


def test_a_table_node_digests_into_a_dense_grid_with_cell_boxes():
    response = _v2_response(
        [
            _v2_table(
                rows=[
                    [("Header", (0.1, 0.2, 0.5, 0.25)), ("Value", (0.5, 0.2, 0.9, 0.25))],
                    [("Loans", (0.1, 0.25, 0.5, 0.3)), ("$1.2B", (0.5, 0.25, 0.9, 0.3))],
                ],
                box=(0.1, 0.2, 0.9, 0.3),
            )
        ]
    )
    digest = digest_page(response, 31, MODEL)
    table = digest.tables[0]
    assert table.cells == [["Header", "Value"], ["Loans", "$1.2B"]]
    assert table.cell_boxes[1][1] == {"left": 0.5, "top": 0.25, "right": 0.9, "bottom": 0.3}
    assert table.box == {"left": 0.1, "top": 0.2, "right": 0.9, "bottom": 0.3}


def test_a_spanning_cell_is_written_once_and_covered_positions_stay_blank():
    response = _v2_response(
        [
            _v2_table(
                rows=[[("Spans both", (0.1, 0.2, 0.9, 0.25))], [("a", None), ("b", None)]],
                spans={(0, 0): (1, 2)},
            )
        ]
    )
    table = digest_page(response, 1, MODEL).tables[0]
    assert table.cells == [["Spans both", ""], ["a", "b"]]
    assert table.cell_boxes[0][1] is None


def test_text_and_page_furniture_digest_into_context():
    response = _v2_response(
        [
            _v2_text("text", "Portfolio Overview(1)", (0.1, 0.05, 0.5, 0.1), "text-0"),
            _v2_text("marginalia", "(1) As of quarter end.", (0.1, 0.9, 0.5, 0.95), "marginalia-0"),
            _v2_text("logo", "STARWOOD", (0.8, 0.02, 0.95, 0.08), "logo-0"),
        ]
    )
    digest = digest_page(response, 4, MODEL)
    assert digest.figures == [] and digest.tables == []
    assert [(c.kind, c.text) for c in digest.context] == [
        ("text", "Portfolio Overview(1)"),
        ("marginalia", "(1) As of quarter end."),
        ("logo", "STARWOOD"),
    ]


def test_projection_carries_figures_tables_and_context_to_the_page_scan():
    response = _v2_response(
        [
            _v2_figure("CHART", "A bar chart.", rows=[["2025: 560"]], box=(0.1, 0.2, 0.9, 0.8)),
            _v2_text("marginalia", "Note: unaudited.", (0.1, 0.9, 0.5, 0.95), "marginalia-0"),
        ]
    )
    digest = digest_page(response, 7, MODEL, "deck.p7.ade.json")
    scan = project_scan(digest, ["bar_chart"], "deck.p7.ade.json", reused=False)

    assert scan.status == "figures"
    assert scan.model == MODEL
    assert scan.version == "dpt-3-pro-20260710"
    assert scan.credits == 0.8
    record = scan.figures[0]
    assert record.chunk_id == "figure-0"
    assert record.box == {"left": 0.1, "top": 0.2, "right": 0.9, "bottom": 0.8}
    # The projected text is the description followed by the value rows.
    assert "A bar chart." in record.text
    assert "| 2025: 560 |" in record.text
    assert [c.kind for c in scan.context] == ["marginalia"]


def test_a_page_of_only_tables_projects_as_tables_status():
    response = _v2_response([_v2_table(rows=[[("a", None)]])])
    scan = project_scan(digest_page(response, 2, MODEL), [])
    assert scan.status == "tables"
    assert scan.tables[0].cells == [["a"]]


def test_a_recognized_response_with_no_figure_or_table_is_an_empty_page():
    response = _v2_response([_v2_text("text", "Prose only.", (0.1, 0.1, 0.9, 0.2))])
    scan = project_scan(digest_page(response, 4, MODEL), [])
    assert scan.status == "empty"


def test_an_unrecognized_format_raises_naming_the_file_and_model():
    with pytest.raises(UnrecognizedResponse) as excinfo:
        digest_page({"unexpected": True}, 5, MODEL, "deck.p5.ade.json")
    assert "deck.p5.ade.json" in str(excinfo.value)
    assert MODEL in str(excinfo.value)


def test_a_failed_page_status_raises_with_the_responses_reason():
    response = _v2_response([], page_status="failed", page_reason="render error")
    with pytest.raises(ValueError, match="render error"):
        digest_page(response, 3, MODEL, "deck.p3.ade.json")


# --- the corrected graft -------------------------------------------------------


def _grafted_page(
    texts: Sequence[Tuple[str, DocItemLabel]],
) -> Tuple[DoclingDocument, PageScan, dpt3.PageDigest]:
    """One page: a picture with `texts` filed under it, grafted by the track.

    The scan comes from a digested response whose figure covers the picture.
    """
    doc = _document([(7, (0.08, 0.15, 0.48, 0.64), "bar_chart")])
    picture = doc.pictures[0]
    for text, label in texts:
        doc.add_text(
            label=label,
            text=text,
            prov=ProvenanceItem(page_no=7, bbox=_bottomleft((0.1, 0.2, 0.3, 0.25)), charspan=(0, len(text))),
            parent=picture,
        )
    response = _v2_response(
        [
            _v2_figure(
                "CHART",
                "Undepreciated book value per share by year.",
                rows=[["2024: $21.13"], ["2025: $21.38"]],
                box=(0.07, 0.14, 0.49, 0.65),
            )
        ]
    )
    digest = digest_page(response, 7, MODEL)
    scan = project_scan(digest, ["bar_chart"])
    refined, errors = dpt3.graft_figures(doc, [scan], _page_dims([7]))
    assert errors == []
    return refined, scan, digest


def test_the_grafts_reading_replaces_content_text_under_the_picture():
    refined, scan, _digest = _grafted_page([("$21.13", DocItemLabel.TEXT), ("$21.38", DocItemLabel.TEXT)])
    assert [t.text for t in refined.texts] == []
    picture = refined.pictures[0]
    assert picture.meta is not None and picture.meta.description is not None
    assert "Undepreciated book value" in picture.meta.description.text
    assert scan.figures[0].picture_ref == picture.self_ref


def test_footnote_and_caption_labelled_text_survives_the_graft():
    refined, _scan, _digest = _grafted_page(
        [
            ("$21.13", DocItemLabel.TEXT),
            ("(1) Based on undepreciated book equity.", DocItemLabel.FOOTNOTE),
            ("Book Value Per Share", DocItemLabel.CAPTION),
        ]
    )
    assert [t.text for t in refined.texts] == [
        "(1) Based on undepreciated book equity.",
        "Book Value Per Share",
    ]


# --- footnote resolution -------------------------------------------------------


def _mark(picture, marks: Sequence[Dict[str, str]], notes: Sequence[Dict[str, str]]) -> None:
    """What the correction agent leaves on a picture, written directly."""
    base = picture.meta or PictureMeta()
    picture.meta = base.model_copy(update={FOOTNOTE_MARKS_FIELD: list(marks), FOOTNOTES_FIELD: list(notes)})


def test_agreeing_sources_corroborate_the_note():
    refined, scan, digest = _grafted_page(
        [("(1) Based on undepreciated book equity.", DocItemLabel.FOOTNOTE)]
    )
    _mark(
        refined.pictures[0],
        marks=[{"marker": "(1)", "kind": "footnote", "label": "2024: $21.13"}],
        notes=[{"marker": "(1)", "text": "Based on undepreciated book equity."}],
    )
    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert flags == []
    note = digest.figures[0].footnotes[0]
    assert note.corroborated is True
    assert sorted(note.sources) == ["agent", "document"]
    assert note.text == "Based on undepreciated book equity."


def test_disagreeing_sources_resolve_with_a_review_flag():
    refined, scan, digest = _grafted_page([("(1) Excludes the CMBS portfolio.", DocItemLabel.FOOTNOTE)])
    _mark(
        refined.pictures[0],
        marks=[{"marker": "(1)", "kind": "footnote", "label": "2024: $21.13"}],
        notes=[{"marker": "(1)", "text": "Based on undepreciated book equity."}],
    )
    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert len(flags) == 1
    assert "reads differently" in flags[0]
    note = digest.figures[0].footnotes[0]
    assert note.corroborated is False


def test_the_documents_own_text_resolves_when_the_agent_read_nothing():
    """A failed agent call no longer costs the page its footnotes."""
    refined, scan, digest = _grafted_page(
        [("(1) Based on undepreciated book equity.", DocItemLabel.FOOTNOTE)]
    )
    _mark(
        refined.pictures[0],
        marks=[{"marker": "(1)", "kind": "footnote", "label": "2024: $21.13"}],
        notes=[],
    )
    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert flags == []
    note = digest.figures[0].footnotes[0]
    assert note.sources == ["document"]
    assert note.text == "Based on undepreciated book equity."


def test_a_marker_no_source_defines_is_flagged_never_dropped():
    refined, scan, digest = _grafted_page([])
    _mark(
        refined.pictures[0],
        marks=[{"marker": "(9)", "kind": "footnote", "label": "2024: $21.13"}],
        notes=[],
    )
    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert len(flags) == 1
    assert "(9)" in flags[0] and "defined by no note" in flags[0]


def test_a_note_no_marker_points_at_is_flagged():
    refined, scan, digest = _grafted_page([])
    _mark(
        refined.pictures[0],
        marks=[],
        notes=[{"marker": "(2)", "text": "An orphaned definition."}],
    )
    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert len(flags) == 1
    assert "no marker" in flags[0]


def test_a_resolved_note_is_inserted_after_the_picture_in_the_document():
    refined, scan, digest = _grafted_page([])
    _mark(
        refined.pictures[0],
        marks=[{"marker": "(1)", "kind": "footnote", "label": "2024: $21.13"}],
        notes=[{"marker": "(1)", "text": "Based on undepreciated book equity."}],
    )
    dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    picture = refined.pictures[0]
    assert picture.parent is not None
    body = [ref.cref for ref in picture.parent.resolve(refined).children]
    inserted = [t for t in refined.texts if t.label == DocItemLabel.FOOTNOTE]
    assert inserted[0].text == "(1) Based on undepreciated book equity."
    # Right after the picture, so the marker and its note read together.
    assert body.index(inserted[0].self_ref) == body.index(picture.self_ref) + 1


def test_resolving_twice_does_not_double_the_inserted_notes():
    """Reusing stored scans reruns the workflow over the enriched parse; the
    notes it inserted the first time must not double."""
    refined, scan, digest = _grafted_page([])
    _mark(
        refined.pictures[0],
        marks=[{"marker": "(1)", "kind": "footnote", "label": "2024"}],
        notes=[{"marker": "(1)", "text": "Based on undepreciated book equity."}],
    )
    dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    inserted = [t for t in refined.texts if t.label == DocItemLabel.FOOTNOTE]
    assert len(inserted) == 1


def test_an_unmarked_general_note_lands_on_its_figure():
    refined, scan, digest = _grafted_page([])
    _mark(
        refined.pictures[0],
        marks=[],
        notes=[{"marker": "", "text": "Note: As of 03/31/2026."}],
    )
    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert flags == []
    assert digest.figures[0].footnotes[0].marker == ""
    inserted = [t for t in refined.texts if t.label == DocItemLabel.FOOTNOTE]
    assert inserted[0].text == "Note: As of 03/31/2026."


def test_a_section_marker_resolves_to_a_heading_pointer():
    refined, scan, digest = _grafted_page([])
    refined.add_text(
        label=DocItemLabel.SECTION_HEADER,
        text="Note 16. Commitments and Contingencies",
        prov=ProvenanceItem(page_no=7, bbox=_bottomleft((0.1, 0.02, 0.6, 0.06)), charspan=(0, 38)),
    )
    _mark(
        refined.pictures[0],
        marks=[{"marker": "Note 16", "kind": "section", "label": "2024: $21.13"}],
        notes=[],
    )
    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert flags == []
    note = digest.figures[0].footnotes[0]
    assert note.sources == ["headings"]
    assert "Note 16. Commitments and Contingencies" in note.text


# --- value-level placement -----------------------------------------------------


def test_a_note_attaches_to_the_value_its_markers_label_names():
    refined, scan, digest = _grafted_page([])
    _mark(
        refined.pictures[0],
        marks=[{"marker": "(1)", "kind": "footnote", "label": "2024"}],
        notes=[{"marker": "(1)", "text": "Based on undepreciated book equity."}],
    )
    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert flags == []
    rows = digest.figures[0].values[0].rows
    placed = [(cell.text, cell.footnotes) for row in rows for cell in row if cell.footnotes]
    assert len(placed) == 1
    assert placed[0][0] == "2024: $21.13"
    assert placed[0][1][0].marker == "(1)"
    assert digest.figures[0].footnotes[0].placed is True


def test_a_marker_on_two_labels_of_one_figure_resolves_once_and_places_by_either():
    """One printed note line, one resolution, one inserted node — the chart-title
    label misses the value table, the series label hits it, and neither raises
    an unplaced flag because the note found its value."""
    refined, scan, digest = _grafted_page([])
    _mark(
        refined.pictures[0],
        marks=[
            {"marker": "(1)", "kind": "footnote", "label": "Book Value Per Share(1)"},
            {"marker": "(1)", "kind": "footnote", "label": "2024"},
        ],
        notes=[{"marker": "(1)", "text": "Based on undepreciated book equity."}],
    )
    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert flags == []
    assert len(digest.figures[0].footnotes) == 1
    assert digest.figures[0].footnotes[0].placed is True
    inserted = [t for t in refined.texts if t.label == DocItemLabel.FOOTNOTE]
    assert len(inserted) == 1


def test_an_unresolvable_marker_shared_by_two_figures_flags_once():
    doc = _document([(7, (0.05, 0.2, 0.45, 0.8), "pie_chart"), (7, (0.55, 0.2, 0.95, 0.8), "pie_chart")])
    response = _v2_response(
        [
            _v2_figure("CHART", "Left.", box=(0.05, 0.2, 0.45, 0.8), node_id="figure-0"),
            _v2_figure("CHART", "Right.", box=(0.55, 0.2, 0.95, 0.8), node_id="figure-1"),
        ]
    )
    digest = digest_page(response, 7, MODEL)
    scan = project_scan(digest, ["pie_chart"])
    refined, errors = dpt3.graft_figures(doc, [scan], _page_dims([7]))
    assert errors == []
    for picture in refined.pictures:
        _mark(picture, marks=[{"marker": "(k)", "kind": "footnote", "label": "Total"}], notes=[])
    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert len(flags) == 1
    assert "(k)" in flags[0]


def test_a_marker_whose_label_matches_no_value_is_flagged_unplaced():
    refined, scan, digest = _grafted_page([])
    _mark(
        refined.pictures[0],
        marks=[{"marker": "(1)", "kind": "footnote", "label": "Woodstar Fund Debt"}],
        notes=[{"marker": "(1)", "text": "Based on undepreciated book equity."}],
    )
    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert len(flags) == 1
    assert "no value under that label" in flags[0]
    # The note still reaches the picture and the figure record.
    assert digest.figures[0].footnotes[0].placed is False
    assert [t for t in refined.texts if t.label == DocItemLabel.FOOTNOTE]


# --- figure identity into the correction ---------------------------------------


def test_the_correction_agent_is_briefed_with_the_digested_figures_in_order():
    """Marker binding is by figure index, so the list the agent sees must be the
    digested figures, in digestion order, with their text."""
    doc = _document([(7, (0.05, 0.2, 0.45, 0.8), "pie_chart"), (7, (0.55, 0.2, 0.95, 0.8), "pie_chart")])
    response = _v2_response(
        [
            _v2_figure(
                "CHART", "Left donut.", rows=[["A: 60%"]], box=(0.05, 0.2, 0.45, 0.8), node_id="figure-0"
            ),
            _v2_figure(
                "CHART", "Right donut.", rows=[["B: 40%"]], box=(0.55, 0.2, 0.95, 0.8), node_id="figure-1"
            ),
        ]
    )
    digest = digest_page(response, 7, MODEL)
    scan = project_scan(digest, ["pie_chart"])
    refined, errors = dpt3.graft_figures(doc, [scan], _page_dims([7]))
    assert errors == []

    finder = MockFigureCorrector(
        markers=[Marker(figure=1, marker="(1)", label="B")],
        notes=[Note(figure=1, marker="(1)", text="Right-hand note.")],
    )
    with tempfile.TemporaryDirectory() as tmp:
        source = _blank_pdf(Path(tmp) / "doc.pdf", pages=7)
        asyncio.run(correct_figures(refined, [scan], _page_dims([7]), source, finder))

    briefed = finder.figures[0]
    assert [f.index for f in briefed] == [0, 1]
    assert "Left donut." in briefed[0].description
    assert "Right donut." in briefed[1].description

    flags = dpt3.resolve_figure_footnotes(refined, [scan], {7: digest})
    assert flags == []
    # The marker filed under the right-hand figure, and its value took the note.
    right = digest.figures[1]
    assert right.footnotes[0].text == "Right-hand note."
    assert right.values[0].rows[0][0].footnotes[0].marker == "(1)"
    assert digest.figures[0].footnotes == []


# --- text blocks -----------------------------------------------------------------


def _panel_page() -> Tuple[DoclingDocument, PageScan, dpt3.PageDigest]:
    """A stat-panel page: four value/label atoms in two blocks, one prose line
    outside every block, and a footnote inside a block that must keep its role."""
    doc = _document([])
    doc.add_page(page_no=4, size=Size(width=PAGE_W, height=PAGE_H))
    for text, box, label in [
        ("$78B", (0.10, 0.45, 0.20, 0.50), DocItemLabel.TEXT),
        ("Investor capital", (0.10, 0.52, 0.30, 0.56), DocItemLabel.TEXT),
        ("165+", (0.60, 0.45, 0.70, 0.50), DocItemLabel.TEXT),
        ("BREDS professionals", (0.60, 0.52, 0.80, 0.56), DocItemLabel.TEXT),
        ("BREDS invests across Core and Core+ strategies.", (0.10, 0.05, 0.60, 0.10), DocItemLabel.TEXT),
        ("(f) Based on originations since 2019.", (0.12, 0.60, 0.40, 0.64), DocItemLabel.FOOTNOTE),
    ]:
        doc.add_text(
            label=label,
            text=text,
            prov=ProvenanceItem(page_no=4, bbox=_bottomleft(box), charspan=(0, len(text))),
        )
    response = _v2_response(
        [
            _v2_text("text", "Global Scale \n$78B \nInvestor capital ", (0.05, 0.40, 0.45, 0.80), "text-0"),
            _v2_text("text", "165+ \nBREDS professionals ", (0.55, 0.40, 0.95, 0.80), "text-1"),
        ]
    )
    digest = digest_page(response, 4, MODEL)
    scan = project_scan(digest, [])
    return doc, scan, digest


def test_scanned_text_groups_under_the_blocks_the_scan_drew():
    doc, scan, digest = _panel_page()
    written = dpt3.group_scanned_text(doc, [scan], {4: digest}, _page_dims([4]))
    assert written == 2

    groups = [g for g in doc.groups if g.name == dpt3.SCAN_BLOCK]
    assert len(groups) == 2
    left = [ref.resolve(doc).text for ref in groups[0].children]
    right = [ref.resolve(doc).text for ref in groups[1].children]
    assert left == ["$78B", "Investor capital"]
    assert right == ["165+", "BREDS professionals"]
    # The items keep their identity; only their parent changed.
    assert all(ref.resolve(doc).parent.cref == groups[0].self_ref for ref in groups[0].children)


def test_text_outside_every_block_and_labelled_text_stay_where_they_were():
    doc, scan, digest = _panel_page()
    dpt3.group_scanned_text(doc, [scan], {4: digest}, _page_dims([4]))
    body_texts = [t.text for t in doc.texts if t.parent is not None and t.parent.resolve(doc) is doc.body]
    assert "BREDS invests across Core and Core+ strategies." in body_texts
    assert "(f) Based on originations since 2019." in body_texts


def test_a_group_keeps_the_first_members_reading_order_position():
    doc, scan, digest = _panel_page()
    dpt3.group_scanned_text(doc, [scan], {4: digest}, _page_dims([4]))
    refs = [ref.cref for ref in doc.body.children]
    group_pos = refs.index(doc.groups[0].self_ref)
    # The first member was the first text added: the group stands in its place.
    assert group_pos == 0


def test_a_block_with_one_item_gets_no_group():
    doc = _document([])
    doc.add_page(page_no=4, size=Size(width=PAGE_W, height=PAGE_H))
    doc.add_text(
        label=DocItemLabel.TEXT,
        text="$78B",
        prov=ProvenanceItem(page_no=4, bbox=_bottomleft((0.10, 0.45, 0.20, 0.50)), charspan=(0, 4)),
    )
    response = _v2_response([_v2_text("text", "Global Scale", (0.05, 0.40, 0.45, 0.80), "text-0")])
    digest = digest_page(response, 4, MODEL)
    written = dpt3.group_scanned_text(doc, [project_scan(digest, [])], {4: digest}, _page_dims([4]))
    assert written == 0
    assert doc.groups == []


def test_grouping_twice_changes_nothing():
    doc, scan, digest = _panel_page()
    dpt3.group_scanned_text(doc, [scan], {4: digest}, _page_dims([4]))
    written_again = dpt3.group_scanned_text(doc, [scan], {4: digest}, _page_dims([4]))
    assert written_again == 0
    assert len([g for g in doc.groups if g.name == dpt3.SCAN_BLOCK]) == 2


def test_a_grouped_document_round_trips_through_serialization():
    doc, scan, digest = _panel_page()
    dpt3.group_scanned_text(doc, [scan], {4: digest}, _page_dims([4]))
    reloaded = DoclingDocument.model_validate(doc.export_to_dict())
    groups = [g for g in reloaded.groups if g.name == dpt3.SCAN_BLOCK]
    assert [[ref.resolve(reloaded).text for ref in g.children] for g in groups] == [
        ["$78B", "Investor capital"],
        ["165+", "BREDS professionals"],
    ]


def test_a_table_nodes_cells_are_blocks_and_bind_value_to_label():
    """A panel collage the scan read as a table still partitions the page: each
    cell is a block, so a value and its label bind at the cell and two cells'
    contents never merge."""
    doc = _document([])
    doc.add_page(page_no=3, size=Size(width=PAGE_W, height=PAGE_H))
    for text, box in [
        ("$1.6 billion", (0.10, 0.42, 0.28, 0.47)),
        ("Total Financing Capacity", (0.10, 0.50, 0.30, 0.55)),
        ("~ 67%", (0.60, 0.42, 0.70, 0.47)),
        ("Non-Mark-toMarket Borrowings", (0.60, 0.50, 0.85, 0.55)),
    ]:
        doc.add_text(
            label=DocItemLabel.TEXT,
            text=text,
            prov=ProvenanceItem(page_no=3, bbox=_bottomleft(box), charspan=(0, len(text))),
        )
    response = _v2_response(
        [
            _v2_table(
                rows=[
                    [
                        ("$1.6 billion Total Financing Capacity", (0.05, 0.40, 0.45, 0.60)),
                        ("~67% Non-Mark-to-Market Borrowings", (0.55, 0.40, 0.95, 0.60)),
                    ]
                ],
                box=(0.05, 0.35, 0.95, 0.65),
            )
        ]
    )
    digest = digest_page(response, 3, MODEL)
    scan = project_scan(digest, [])
    written = dpt3.group_scanned_text(doc, [scan], {3: digest}, _page_dims([3]))
    assert written == 2
    groups = [g for g in doc.groups if g.name == dpt3.SCAN_BLOCK]
    assert [[ref.resolve(doc).text for ref in g.children] for g in groups] == [
        ["$1.6 billion", "Total Financing Capacity"],
        ["~ 67%", "Non-Mark-toMarket Borrowings"],
    ]


def test_a_scanned_table_nothing_homed_is_flagged_never_silent():
    doc = _document([])
    doc.add_page(page_no=3, size=Size(width=PAGE_W, height=PAGE_H))
    response = _v2_response([_v2_table(rows=[[("a", None)]], box=(0.1, 0.2, 0.9, 0.6), node_id="table-1")])
    scan = project_scan(digest_page(response, 3, MODEL), [])
    flags = dpt3.unhomed_tables([scan], doc, _page_dims([3]))
    assert len(flags) == 1
    assert "table-1" in flags[0] and "no table or picture" in flags[0]

    scan.tables[0].table_id = "extracted"
    assert dpt3.unhomed_tables([scan], doc, _page_dims([3])) == []


def test_a_scanned_table_over_a_document_table_is_redundancy_not_a_flag():
    """Set-of-Mark owns text-layer tables; the scan re-reading one is by design."""
    from docling_core.types.doc.document import TableData

    doc = _document([])
    doc.add_page(page_no=3, size=Size(width=PAGE_W, height=PAGE_H))
    doc.add_table(
        data=TableData(table_cells=[], num_rows=0, num_cols=0),
        prov=ProvenanceItem(page_no=3, bbox=_bottomleft((0.1, 0.2, 0.9, 0.6)), charspan=(0, 0)),
    )
    response = _v2_response([_v2_table(rows=[[("a", None)]], box=(0.1, 0.2, 0.9, 0.6), node_id="table-1")])
    scan = project_scan(digest_page(response, 3, MODEL), [])
    assert dpt3.unhomed_tables([scan], doc, _page_dims([3])) == []


def test_a_stored_dpt3_response_beyond_nomination_is_read(tmp_path):
    from quber.core.figures.orchestrator import _stored_track_pages

    (tmp_path / "deck.p4.ade.json").write_text(json.dumps(_v2_response([])), encoding="utf-8")
    (tmp_path / "deck.p7.ade.json").write_text(json.dumps(_dpt2_response()), encoding="utf-8")
    (tmp_path / "deck.p9.ade.json").write_text(json.dumps(_v2_response([])), encoding="utf-8")
    (tmp_path / "deck.p9.ade.dpt2.json").write_text(json.dumps(_dpt2_response()), encoding="utf-8")

    extra = _stored_track_pages("deck", str(tmp_path), nominated={9})
    # Page 4's dpt-3 response joins; page 7 is dpt-2 (rescanning it would bill a
    # page nobody nominated); page 9 is already nominated; aside files never match.
    assert [n.page for n in extra] == [4]


# --- orphan tables ---------------------------------------------------------------


def _legend_response() -> Dict[str, Any]:
    """A map-legend table: printed values the parse's text layer never captured."""
    return _v2_response(
        [
            _v2_table(
                rows=[
                    [("Northeast", (0.74, 0.24, 0.84, 0.28)), ("17 %", (0.86, 0.24, 0.94, 0.28))],
                    [("Southwest", (0.74, 0.30, 0.84, 0.34)), ("17 %", (0.86, 0.30, 0.94, 0.34))],
                ],
                box=(0.73, 0.23, 0.96, 0.42),
                node_id="table-0",
            )
        ]
    )


def test_a_grid_over_nothing_is_captured_and_inserted_as_a_table():
    from typing import cast

    from quber.agents.cell_reader import MockCellReader
    from quber.core.extractors.set_of_mark.assemble import TableAssembly
    from quber.core.figures.capture import capture_tables

    doc = _document([])
    doc.add_page(page_no=5, size=Size(width=PAGE_W, height=PAGE_H))
    digest = digest_page(_legend_response(), 5, MODEL)
    scan = project_scan(digest, [])

    with tempfile.TemporaryDirectory() as tmp:
        source = _blank_pdf(Path(tmp) / "doc.pdf", pages=5)
        assembly = TableAssembly(source, cast(Any, None), asyncio.Semaphore(1), 200)
        tables, errors = asyncio.run(
            capture_tables(doc, [scan], source, _page_dims([5]), assembly, MockCellReader(), orphans=True)
        )
    assert errors == []
    assert len(tables) == 1
    assert scan.tables[0].table_id is not None

    insert_errors = dpt3.insert_orphan_tables(doc, [scan], tables, _page_dims([5]))
    assert insert_errors == []
    assert len(doc.tables) == 1
    assert [c.text for c in doc.tables[0].data.table_cells][:2] == ["Northeast", "17 %"]
    assert scan.tables[0].table_ref == doc.tables[0].self_ref
    # Homed now, so the unhomed check has nothing to flag.
    assert dpt3.unhomed_tables([scan], doc, _page_dims([5])) == []


def test_a_grid_over_parse_text_is_not_an_orphan():
    """A stat panel's strings the parse does hold belong to the block grouping;
    inserting the scan's grid would state the content twice."""
    from typing import cast

    from quber.agents.cell_reader import MockCellReader
    from quber.core.extractors.set_of_mark.assemble import TableAssembly
    from quber.core.figures.capture import capture_tables

    doc = _document([])
    doc.add_page(page_no=5, size=Size(width=PAGE_W, height=PAGE_H))
    doc.add_text(
        label=DocItemLabel.TEXT,
        text="Northeast",
        prov=ProvenanceItem(page_no=5, bbox=_bottomleft((0.74, 0.24, 0.84, 0.28)), charspan=(0, 9)),
    )
    scan = project_scan(digest_page(_legend_response(), 5, MODEL), [])

    with tempfile.TemporaryDirectory() as tmp:
        source = _blank_pdf(Path(tmp) / "doc.pdf", pages=5)
        assembly = TableAssembly(source, cast(Any, None), asyncio.Semaphore(1), 200)
        tables, _errors = asyncio.run(
            capture_tables(doc, [scan], source, _page_dims([5]), assembly, MockCellReader(), orphans=True)
        )
    assert tables == []
    assert scan.tables[0].table_id is None


def test_without_the_orphan_tier_an_unclaimed_grid_stays_unclaimed():
    """The dpt-2 path's behavior, unchanged: capture with default arguments
    leaves a grid over nothing alone."""
    from typing import cast

    from quber.agents.cell_reader import MockCellReader
    from quber.core.extractors.set_of_mark.assemble import TableAssembly
    from quber.core.figures.capture import capture_tables

    doc = _document([])
    doc.add_page(page_no=5, size=Size(width=PAGE_W, height=PAGE_H))
    scan = project_scan(digest_page(_legend_response(), 5, MODEL), [])

    with tempfile.TemporaryDirectory() as tmp:
        source = _blank_pdf(Path(tmp) / "doc.pdf", pages=5)
        assembly = TableAssembly(source, cast(Any, None), asyncio.Semaphore(1), 200)
        tables, _errors = asyncio.run(
            capture_tables(doc, [scan], source, _page_dims([5]), assembly, MockCellReader())
        )
    assert tables == []
    assert scan.tables[0].table_id is None


# --- model-keyed reuse ----------------------------------------------------------


def _dpt2_response() -> Dict[str, Any]:
    return {
        "markdown": "a chart",
        "chunks": [
            {
                "id": "c1",
                "type": "figure",
                "markdown": "a chart",
                "grounding": {"box": {"left": 0.1, "top": 0.2, "right": 0.9, "bottom": 0.8}, "page": 0},
            }
        ],
        "metadata": {"credit_usage": 3.0, "job_id": "job-v1", "version": "dpt-2-20260410"},
    }


def _fake_scan(response: Dict[str, Any]):
    calls: List[int] = []

    async def scan(source: Path, page: int, model: str = "unused") -> Dict[str, Any]:
        calls.append(page)
        return response

    return scan, calls


def test_a_matching_stored_response_is_reused_without_a_scan(tmp_path, monkeypatch):
    (tmp_path / "deck.p7.ade.json").write_text(json.dumps(_v2_response([])), encoding="utf-8")
    fake, calls = _fake_scan(_v2_response([]))
    monkeypatch.setattr("quber.core.figures.orchestrator.scan_page", fake)

    response, _artifact, reused = asyncio.run(
        _page_response(Path("doc.pdf"), "deck", 7, MODEL, str(tmp_path))
    )

    assert reused is True
    assert calls == []
    assert response["structure"]["type"] == "document"


def test_a_dpt2_response_under_a_dpt3_request_is_moved_aside_and_rescanned(tmp_path, monkeypatch):
    (tmp_path / "deck.p7.ade.json").write_text(json.dumps(_dpt2_response()), encoding="utf-8")
    fresh = _v2_response([])
    fake, calls = _fake_scan(fresh)
    monkeypatch.setattr("quber.core.figures.orchestrator.scan_page", fake)

    response, _artifact, reused = asyncio.run(
        _page_response(Path("doc.pdf"), "deck", 7, MODEL, str(tmp_path))
    )

    assert reused is False
    assert calls == [7]
    assert response["structure"]["type"] == "document"
    # The paid dpt-2 response is aside, never destroyed; the main name holds dpt-3.
    aside = json.loads((tmp_path / "deck.p7.ade.dpt2.json").read_text(encoding="utf-8"))
    assert "chunks" in aside
    stored = json.loads((tmp_path / "deck.p7.ade.json").read_text(encoding="utf-8"))
    assert "structure" in stored


def test_a_dpt3_response_under_a_dpt2_request_is_moved_aside_and_rescanned(tmp_path, monkeypatch):
    (tmp_path / "deck.p7.ade.json").write_text(json.dumps(_v2_response([])), encoding="utf-8")
    fake, calls = _fake_scan(_dpt2_response())
    monkeypatch.setattr("quber.core.figures.orchestrator.scan_page", fake)

    response, _artifact, reused = asyncio.run(
        _page_response(Path("doc.pdf"), "deck", 7, "dpt-2", str(tmp_path))
    )

    assert reused is False and calls == [7]
    assert "chunks" in response
    assert (tmp_path / "deck.p7.ade.dpt3.json").exists()


def test_a_previously_asided_response_is_found_again_without_a_scan(tmp_path, monkeypatch):
    """A later run under the old model finds its paid response by the same check."""
    (tmp_path / "deck.p7.ade.json").write_text(json.dumps(_v2_response([])), encoding="utf-8")
    (tmp_path / "deck.p7.ade.dpt2.json").write_text(json.dumps(_dpt2_response()), encoding="utf-8")
    fake, calls = _fake_scan(_dpt2_response())
    monkeypatch.setattr("quber.core.figures.orchestrator.scan_page", fake)

    response, _artifact, reused = asyncio.run(
        _page_response(Path("doc.pdf"), "deck", 7, "dpt-2", str(tmp_path))
    )

    assert reused is True
    assert calls == []
    assert "chunks" in response
    # The two generations swapped places: main holds dpt-2, aside holds dpt-3.
    assert "chunks" in json.loads((tmp_path / "deck.p7.ade.json").read_text(encoding="utf-8"))
    assert "structure" in json.loads((tmp_path / "deck.p7.ade.dpt3.json").read_text(encoding="utf-8"))


def test_a_stored_response_in_neither_format_raises(tmp_path, monkeypatch):
    (tmp_path / "deck.p7.ade.json").write_text(json.dumps({"nonsense": 1}), encoding="utf-8")
    fake, calls = _fake_scan(_dpt2_response())
    monkeypatch.setattr("quber.core.figures.orchestrator.scan_page", fake)

    with pytest.raises(ValueError, match="neither"):
        asyncio.run(_page_response(Path("doc.pdf"), "deck", 7, "dpt-2", str(tmp_path)))
    assert calls == []


# --- real saved responses --------------------------------------------------------

_REAL = Path(__file__).resolve().parents[2] / "data" / "ade-compare" / "QUE-341"


@pytest.mark.skipif(not _REAL.exists(), reason="local ADE compare artifacts not present")
def test_real_saved_responses_digest_under_both_generations():
    """The same live page, read by each generation's own path — no mocks, no network."""
    v2_path = next((_REAL / "p9-dpt3").glob("*.ade.json"))
    v1_path = next((_REAL / "p9-dpt2").glob("*.ade.json"))

    digest = digest_page(json.loads(v2_path.read_text(encoding="utf-8")), 9, MODEL, v2_path.name)
    v2_scan = project_scan(digest, ["pie_chart"])
    assert v2_scan.status == "figures"
    assert digest.figures[0].kind == "chart"
    assert digest.figures[0].values, "the live chart carries a value table"
    assert v2_scan.version and v2_scan.credits

    v1_scan = page_scan(
        json.loads(v1_path.read_text(encoding="utf-8")), 9, "dpt-2", ["pie_chart"], v1_path.name
    )
    assert v1_scan.status == "figures"
    # Both generations return the page's figures; the track adds the structure.
    assert len(v2_scan.figures) >= 1 and len(v1_scan.figures) >= 1