"""Unit tests for haiku.rag.converters.pdf_split.

The integration test that pins split-and-merge against a real-PDF baseline
lives in tests/test_converters.py — that path requires docling installed and
is gated accordingly. These tests cover the byte-level split mechanism and
the docling-core concatenate contract in isolation.
"""

import io
from pathlib import Path

import pypdfium2 as pdfium
import pytest

from haiku.rag.converters.pdf_split import (
    convert_pdf_with_splitting,
    iter_pdf_slices,
)


def _make_pdf(page_count: int, tmp_path: Path) -> Path:
    """Synthesize a minimal valid multi-page PDF with pypdfium2. Each page is
    an A4-sized blank. Returns the path."""
    doc = pdfium.PdfDocument.new()
    try:
        for _ in range(page_count):
            doc.new_page(width=595.0, height=842.0)  # A4 in points
        out = tmp_path / "synth.pdf"
        with open(out, "wb") as f:
            doc.save(f)
        return out
    finally:
        doc.close()


def test_iter_pdf_slices_partitions_pages(tmp_path):
    src = _make_pdf(7, tmp_path)
    slices = list(iter_pdf_slices(src, slice_size=3))

    assert [(s, e) for s, e, _ in slices] == [(1, 3), (4, 6), (7, 7)]

    page_counts = []
    for _, _, pdf_bytes in slices:
        d = pdfium.PdfDocument(io.BytesIO(pdf_bytes))
        try:
            page_counts.append(len(d))
        finally:
            d.close()
    assert page_counts == [3, 3, 1]
    assert sum(page_counts) == 7


def test_iter_pdf_slices_single_slice_when_doc_fits(tmp_path):
    src = _make_pdf(4, tmp_path)
    slices = list(iter_pdf_slices(src, slice_size=10))
    assert len(slices) == 1
    start, end, pdf_bytes = slices[0]
    assert (start, end) == (1, 4)
    d = pdfium.PdfDocument(io.BytesIO(pdf_bytes))
    try:
        assert len(d) == 4
    finally:
        d.close()


def test_iter_pdf_slices_rejects_zero_slice_size(tmp_path):
    src = _make_pdf(2, tmp_path)
    with pytest.raises(ValueError, match="slice_size must be >= 1"):
        list(iter_pdf_slices(src, slice_size=0))


@pytest.mark.asyncio
async def test_convert_aborts_and_cleans_up_on_mid_stream_slice_failure(
    tmp_path, monkeypatch
):
    """When converting slice 2 of 3 fails, the whole call must raise
    ValueError naming the failed slice's page range, every tempfile created
    along the way must be deleted, and the source PDF handle must be closed.
    """
    src = _make_pdf(7, tmp_path)

    # Pin tempfiles to a per-test dir so we can list leaks deterministically.
    monkeypatch.setattr("tempfile.tempdir", str(tmp_path))

    calls: list[Path] = []

    class _FlakyConverter:
        async def convert_file(self, path: Path, *, source_uri):
            calls.append(path)
            if len(calls) == 2:
                raise RuntimeError("docling exploded")
            from docling_core.types.doc.document import DoclingDocument

            return DoclingDocument(name="slice")

    with pytest.raises(ValueError, match="pages 4-6"):
        await convert_pdf_with_splitting(
            _FlakyConverter(),  # ty: ignore[invalid-argument-type]
            src,
            source_uri=None,
            slice_size=3,
        )

    # All converter inputs were under tmp_path (the pinned tempdir)…
    assert all(str(p).startswith(str(tmp_path)) for p in calls)
    # …and none of them are still on disk.
    leftover = [p for p in calls if p.exists()]
    assert leftover == [], f"tempfiles leaked: {leftover}"
    # We aborted after slice 2; slice 3 was never attempted.
    assert len(calls) == 2


def test_concatenate_shifts_page_nos_and_unique_self_refs():
    """Pins the docling-core contract we rely on: when two docs (each with
    items on page 1) are concatenated, the second doc's items move to page 2
    and self_refs across both stay unique."""
    pytest.importorskip("docling_core")
    from docling_core.types.doc.base import BoundingBox, CoordOrigin
    from docling_core.types.doc.document import (
        DoclingDocument,
        PageItem,
        ProvenanceItem,
        Size,
    )
    from docling_core.types.doc.labels import DocItemLabel

    def _make_one_page_doc(name: str, text: str) -> DoclingDocument:
        d = DoclingDocument(name=name)
        d.pages[1] = PageItem(page_no=1, size=Size(width=595.0, height=842.0))
        d.add_text(
            label=DocItemLabel.TEXT,
            text=text,
            prov=ProvenanceItem(
                page_no=1,
                bbox=BoundingBox(
                    l=0.0, t=0.0, r=100.0, b=20.0, coord_origin=CoordOrigin.TOPLEFT
                ),
                charspan=(0, len(text)),
            ),
        )
        return d

    a = _make_one_page_doc("a", "alpha")
    b = _make_one_page_doc("b", "beta")

    merged = DoclingDocument.concatenate([a, b])

    assert len(merged.texts) == 2
    refs = [t.self_ref for t in merged.texts]
    assert len(set(refs)) == 2, f"self_refs collided: {refs}"

    page_nos = sorted({p.page_no for t in merged.texts for p in t.prov})
    assert page_nos == [1, 2], f"expected b's page 1 to shift to page 2, got {page_nos}"

    assert sorted(merged.pages.keys()) == [1, 2]