tests/datasets/test_extraction.py

from __future__ import annotations

import hashlib
import io
import json
from contextlib import contextmanager
from datetime import UTC, datetime
from pathlib import Path

import httpx
from PIL import Image
from reportlab.lib.pdfencrypt import StandardEncryption
from reportlab.lib.utils import ImageReader
from reportlab.pdfgen import canvas

from gcmd_classifier.config import DatasetExtractionSettings
from gcmd_classifier.datasets.cmr import resolve_cmr_collection, validate_cmr_url
from gcmd_classifier.datasets.documents import RetrievedREADME, retrieve_selected_readme
from gcmd_classifier.datasets.extraction import extract_pdf_manifest
from gcmd_classifier.datasets.models import (
    BlindCollectionView,
    CMRRetrievedSource,
    DatasetExtractionManifest,
    DatasetIdentity,
    ExtractionPreclassificationFailure,
    SourceArtifactReference,
)
from gcmd_classifier.datasets.readme_discovery import (
    discover_readme_candidates,
    select_readme_candidate,
)
from gcmd_classifier.datasets.schemas import schema_for
from tests.datasets.fixture_manifest import load_fixture_manifest

NOW = datetime(2026, 8, 13, 16, 0, tzinfo=UTC)
PUBLIC = "93.184.216.34"
URL = "https://docs.example.test/readme.pdf"


def _pdf_bytes(
    pages: list[list[tuple[float, float, str]]],
    *,
    encrypted: bool = False,
    image_only: bool = False,
) -> bytes:
    output = io.BytesIO()
    encryption = StandardEncryption("secret") if encrypted else None
    writer = canvas.Canvas(output, encrypt=encryption)
    for entries in pages:
        if image_only:
            image = Image.new("RGB", (20, 20), "black")
            writer.drawImage(ImageReader(image), 72, 700, 20, 20)
        for x, y, text in entries:
            writer.drawString(x, y, text)
        writer.showPage()
    writer.save()
    return output.getvalue()


def _validated(pdf_bytes: bytes, tmp_path: Path, *, url: str = URL) -> RetrievedREADME:
    identity = DatasetIdentity(
        concept_id="C1-P",
        native_id="TARGET_001",
        short_name="TARGET",
        version="001",
        cmr_revision_id=3,
    )
    view = BlindCollectionView(
        identity=identity,
        derived_native_id=identity.native_id,
        related_urls=({"Subtype": "READ-ME", "URL": url},),
        blind_view_sha256="a" * 64,
    )
    discovery = discover_readme_candidates(view, cmr_source_sha256="b" * 64)
    selection = select_readme_candidate(discovery, [(0, url)], utc_now=lambda: NOW)

    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200, headers={"content-type": "application/pdf"}, content=pdf_bytes, request=request
        )

    with httpx.Client(transport=httpx.MockTransport(handler), follow_redirects=False) as client:
        return retrieve_selected_readme(
            selection,
            client=client,
            resolver=lambda host: [PUBLIC],
            artifact_directory=tmp_path,
            utc_now=lambda: NOW,
        )


def _success(result: object) -> DatasetExtractionManifest:
    assert isinstance(result, DatasetExtractionManifest), result
    return result


def _failure(result: object, code: str) -> ExtractionPreclassificationFailure:
    assert isinstance(result, ExtractionPreclassificationFailure), result
    assert result.error.code == code
    assert result.classification_attempted is False
    assert result.model_calls == 0
    assert result.usable_for_classification is False
    return result


def test_one_page_exact_provenance_stability_and_reconstruction(tmp_path: Path) -> None:
    source = _validated(
        _pdf_bytes([[(72, 740, "SCIENCE OVERVIEW"), (72, 720, "Exact source text.")]]), tmp_path
    )
    first = _success(extract_pdf_manifest(source))
    second = _success(extract_pdf_manifest(source))
    assert first.model_dump(mode="json") == second.model_dump(mode="json")
    assert first.document_id == f"sha256:{source.record.artifact.sha256}"
    assert first.pages[0].page_number == 1 and first.pages[0].source_order == 0
    assert first.pages[0].text_sha256 == hashlib.sha256(first.pages[0].text.encode()).hexdigest()
    assert "".join(block.text for block in first.blocks) == first.pages[0].text
    assert first.sections[0].label_source == "literal"
    heading = first.sections[0]
    assert (
        first.pages[0].text[heading.heading_character_start : heading.heading_character_end]
        == heading.heading
    )
    assert first.excluded_separators == ()
    assert "ScienceKeywords" not in first.model_dump_json()


def test_multiple_pages_repeated_headings_boundary_and_cross_page_section(tmp_path: Path) -> None:
    data = _pdf_bytes(
        [
            [(72, 740, "1. OVERVIEW"), (72, 720, "Page one body")],
            [(72, 740, "Continuation without heading")],
            [(72, 740, "1. OVERVIEW"), (72, 720, "Repeated heading body")],
        ]
    )
    manifest = _success(extract_pdf_manifest(_validated(data, tmp_path)))
    assert [page.page_number for page in manifest.pages] == [1, 2, 3]
    literal = [section for section in manifest.sections if section.label_source == "literal"]
    assert len(literal) == 2 and literal[0].section_id != literal[1].section_id
    assert literal[0].page_start == 1 and literal[0].page_end == 2
    assert [block.source_order for block in manifest.blocks] == list(range(len(manifest.blocks)))
    for page in manifest.pages:
        assert (
            "".join(block.text for block in manifest.blocks if block.page_start == page.page_number)
            == page.text
        )


def test_no_heading_uses_derived_label_and_long_text_splits_deterministically(
    tmp_path: Path,
) -> None:
    text = "word " * 200
    source = _validated(_pdf_bytes([[(72, 740, text)]]), tmp_path)
    settings = DatasetExtractionSettings(max_block_characters=80)
    manifest = _success(extract_pdf_manifest(source, settings=settings))
    assert manifest.sections[0].label_source == "derived"
    assert manifest.sections[0].heading_character_start is None
    assert len(manifest.blocks) > 1
    assert all(len(block.text) <= 80 for block in manifest.blocks)
    assert "".join(block.text for block in manifest.blocks) == manifest.pages[0].text


def test_blank_page_warning_and_fully_blank_failure(tmp_path: Path) -> None:
    mixed = _pdf_bytes([[], [(72, 740, "Text page")]])
    manifest = _success(extract_pdf_manifest(_validated(mixed, tmp_path / "mixed")))
    assert manifest.pages[0].extraction_status == "blank"
    assert any(warning.code == "BLANK_PAGE" for warning in manifest.pages[0].warnings)
    blank = _pdf_bytes([[]])
    _failure(extract_pdf_manifest(_validated(blank, tmp_path / "blank")), "PDF_NO_EXTRACTABLE_TEXT")


def test_image_only_page_and_document(tmp_path: Path) -> None:
    image_pdf = _pdf_bytes([[]], image_only=True)
    _failure(extract_pdf_manifest(_validated(image_pdf, tmp_path)), "PDF_IMAGE_ONLY")


def test_encrypted_and_malformed_pdf_fail_before_inference(tmp_path: Path) -> None:
    encrypted = _pdf_bytes([[(72, 740, "Secret")]], encrypted=True)
    _failure(extract_pdf_manifest(_validated(encrypted, tmp_path / "encrypted")), "PDF_ENCRYPTED")
    valid = _validated(_pdf_bytes([[(72, 740, "Valid")]]), tmp_path / "malformed")
    malformed = b"%PDF-1.7\nnot readable\n1 0 obj\nstartxref\n0\n%%EOF"
    digest = hashlib.sha256(malformed).hexdigest()
    artifact = SourceArtifactReference(
        artifact_id=f"readme-{digest}",
        sha256=digest,
        byte_size=len(malformed),
        media_type="application/pdf",
        storage_reference="test/document.pdf",
    )
    source = RetrievedREADME(
        record=valid.record.model_copy(update={"artifact": artifact}), content=malformed
    )
    _failure(extract_pdf_manifest(source), "PDF_MALFORMED")


def test_hash_page_character_block_and_timeout_limits(tmp_path: Path) -> None:
    source = _validated(_pdf_bytes([[(72, 740, "Some source text")]]), tmp_path)
    changed = RetrievedREADME(record=source.record, content=source.content + b"x")
    _failure(extract_pdf_manifest(changed), "PDF_HASH_MISMATCH")
    _failure(
        extract_pdf_manifest(
            source, settings=DatasetExtractionSettings(max_pages=0 + 1), pdf_opener=_two_page_fake
        ),
        "PDF_PAGE_LIMIT_EXCEEDED",
    )
    _failure(
        extract_pdf_manifest(source, settings=DatasetExtractionSettings(max_characters=2)),
        "PDF_CHARACTER_LIMIT_EXCEEDED",
    )
    _failure(
        extract_pdf_manifest(
            source, settings=DatasetExtractionSettings(max_blocks=1, max_block_characters=2)
        ),
        "PDF_BLOCK_LIMIT_EXCEEDED",
    )
    clock = iter((0.0, 61.0))
    _failure(extract_pdf_manifest(source, monotonic=lambda: next(clock)), "PDF_EXTRACTION_TIMEOUT")


class _FakePage:
    width = 600.0
    height = 800.0
    images = []
    lines = []
    rects = []

    def __init__(self, text: object = "text", *, words_error: bool = False) -> None:
        self.text = text
        self.words_error = words_error

    def extract_text(self, **kwargs):
        if isinstance(self.text, Exception):
            raise self.text
        return self.text

    def extract_words(self, **kwargs):
        if self.words_error:
            raise RuntimeError("coordinates unavailable")
        return []


@contextmanager
def _fake_document(pages):
    yield type("FakePDF", (), {"pages": pages})()


def _two_page_fake(stream):
    return _fake_document([_FakePage("one"), _FakePage("two")])


def test_page_failure_partial_failure_and_unavailable_coordinates(tmp_path: Path) -> None:
    source = _validated(_pdf_bytes([[(72, 740, "source")]]), tmp_path)
    _failure(
        extract_pdf_manifest(
            source, pdf_opener=lambda stream: _fake_document([_FakePage(RuntimeError("page"))])
        ),
        "PDF_PAGE_EXTRACTION_FAILED",
    )
    _failure(
        extract_pdf_manifest(source, pdf_opener=lambda stream: _fake_document([_FakePage(123)])),
        "PDF_PARTIAL_PAGE_FAILURE",
    )
    manifest = _success(
        extract_pdf_manifest(
            source, pdf_opener=lambda stream: _fake_document([_FakePage("text", words_error=True)])
        )
    )
    assert any(warning.code == "LAYOUT_COORDINATES_UNAVAILABLE" for warning in manifest.warnings)


def test_table_multicolumn_and_low_coverage_warnings(tmp_path: Path) -> None:
    table = _pdf_bytes([[(72, 740, "tiny")]])
    source = _validated(table, tmp_path)
    page = _FakePage("tiny")
    page.lines = [{}] * 8

    def words(**kwargs):
        return [{"x0": x, "x1": x + 5} for x in [10, 20, 30, 40, 50, 450, 460, 470, 480, 490]]

    page.extract_words = words
    manifest = _success(
        extract_pdf_manifest(source, pdf_opener=lambda stream: _fake_document([page]))
    )
    codes = {warning.code for warning in manifest.warnings}
    assert {
        "LOW_TEXT_COVERAGE",
        "TABLE_LIKE_LAYOUT",
        "LIKELY_MULTI_COLUMN",
        "SUSPICIOUS_READING_ORDER",
    } <= codes


def test_policy_version_change_is_explicit_but_content_ids_remain_source_based(
    tmp_path: Path,
) -> None:
    source = _validated(_pdf_bytes([[(72, 740, "Body text")]]), tmp_path)
    first = _success(extract_pdf_manifest(source))
    second = _success(
        extract_pdf_manifest(
            source,
            settings=DatasetExtractionSettings(heading_policy_version="deterministic-heading-v2"),
        )
    )
    assert first.document_id == second.document_id
    assert [page.page_id for page in first.pages] == [page.page_id for page in second.pages]
    assert first.heading_policy_version != second.heading_policy_version


def test_all_frozen_pdfs_extract_deterministically_and_validate_schema(tmp_path: Path) -> None:
    _, cases = load_fixture_manifest(Path("tests/fixtures/datasets/cases.json"))
    before = {
        path: path.read_bytes() for case in cases for path in (case.cmr_path, case.readme_path)
    }
    outcomes = {}
    for case in cases:
        meta = case.metadata
        cmr_bytes = case.cmr_path.read_bytes()
        pdf_bytes = case.readme_path.read_bytes()
        validated = validate_cmr_url(meta.cmr_url)
        cmr_source = CMRRetrievedSource(
            submitted_url=meta.cmr_url,
            final_url=meta.cmr_url,
            retrieved_at="2026-08-13T00:00:00Z",
            status_code=200,
            response_bytes=cmr_bytes,
            source_text=cmr_bytes.decode(),
            sha256=meta.cmr_sha256,
            parsed_json=json.loads(cmr_bytes),
        )
        resolved = resolve_cmr_collection(validated, cmr_source)
        discovery = discover_readme_candidates(
            resolved.blind_view, cmr_source_sha256=meta.cmr_sha256
        )
        candidate = next(
            item for item in discovery.candidates if item.url == meta.selected_readme_url
        )
        selection = select_readme_candidate(
            discovery, [(candidate.source_index, candidate.url)], utc_now=lambda: NOW
        )

        def handler(request, expected=meta.selected_readme_url, content=pdf_bytes):
            assert str(request.url) == expected
            return httpx.Response(
                200, headers={"content-type": "application/pdf"}, content=content, request=request
            )

        with httpx.Client(transport=httpx.MockTransport(handler), follow_redirects=False) as client:
            retrieved = retrieve_selected_readme(
                selection,
                client=client,
                resolver=lambda host: [PUBLIC],
                artifact_directory=tmp_path / meta.case_id,
                utc_now=lambda: NOW,
            )
        first = extract_pdf_manifest(retrieved)
        second = extract_pdf_manifest(retrieved)
        manifest = _success(first)
        assert manifest.model_dump(mode="json") == _success(second).model_dump(mode="json")
        assert manifest.document.sha256 == meta.readme_sha256
        assert manifest.lineage.cmr_source_sha256 == meta.cmr_sha256
        assert manifest.lineage.identity.concept_id == meta.expected_concept_id
        assert "ScienceKeywords" not in manifest.model_dump_json()
        DatasetExtractionManifest.model_validate(manifest.model_dump(mode="json"))
        assert schema_for(DatasetExtractionManifest)["title"] == "DatasetExtractionManifest"
        outcomes[meta.case_id] = (manifest.page_count, manifest.block_count, len(manifest.warnings))
    assert {path: path.read_bytes() for path in before} == before
    assert len(outcomes) == 4


def test_multiple_headings_on_page_and_heading_at_page_end(tmp_path: Path) -> None:
    source = _validated(
        _pdf_bytes(
            [
                [
                    (72, 740, "1. FIRST HEADING"),
                    (72, 720, "First body"),
                    (72, 700, "2. END HEADING"),
                ],
                [(72, 740, "Body continuing from end heading")],
            ]
        ),
        tmp_path,
    )
    manifest = _success(extract_pdf_manifest(source))
    literal = [section for section in manifest.sections if section.label_source == "literal"]
    assert [section.heading for section in literal] == ["1. FIRST HEADING", "2. END HEADING"]
    assert literal[-1].page_start == 1 and literal[-1].page_end == 2
    assert len(set(section.section_id for section in literal)) == 2