"""The pages=all union: every page scanned, nomination info intact."""

from __future__ import annotations

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

from quber.core.figures.nominate import nominate_pages
from quber.core.figures.orchestrator import widen_to_all_pages

PAGE_W, PAGE_H = 720.0, 540.0


def _document(page_count: int, chart_pages: list[int]) -> DoclingDocument:
    doc = DoclingDocument(name="synthetic")
    for page in range(1, page_count + 1):
        doc.add_page(page_no=page, size=Size(width=PAGE_W, height=PAGE_H))
    for page in chart_pages:
        bbox = BoundingBox(l=72, r=648, t=486, b=54, coord_origin=CoordOrigin.BOTTOMLEFT)
        picture = doc.add_picture(prov=ProvenanceItem(page_no=page, bbox=bbox, charspan=(0, 0)))
        picture.meta = PictureMeta(
            classification=PictureClassificationMetaField(
                predictions=[PictureClassificationPrediction(class_name="bar_chart", confidence=0.9)]
            )
        )
    return doc


def test_all_mode_covers_every_page_and_keeps_nomination_refs():
    doc = _document(page_count=6, chart_pages=[2, 5])
    nominated = nominate_pages(doc)
    assert [n.page for n in nominated] == [2, 5]

    widened = widen_to_all_pages(doc, nominated)
    assert [n.page for n in widened] == [1, 2, 3, 4, 5, 6]

    by_page = {n.page: n for n in widened}
    # The nominated pages keep what nominated them; the joined pages are bare.
    assert by_page[2].picture_refs and by_page[2].picture_classes == ["bar_chart"]
    assert by_page[5].picture_refs
    for bare in (1, 3, 4, 6):
        assert not by_page[bare].picture_refs
        assert not by_page[bare].table_refs


def test_all_mode_on_a_fully_nominated_document_changes_nothing():
    doc = _document(page_count=2, chart_pages=[1, 2])
    nominated = nominate_pages(doc)
    assert widen_to_all_pages(doc, nominated) == nominated