arborist/tests/test_loss_report.py

"""Adapter LossReport tests (ticket #000022).

Three layers of coverage:

1. ``test_to_base_*`` — wikitext.to_base() with collector emits the
   expected LossEvent set; without collector remains bit-identical to
   the pre-#000022 implementation.
2. ``test_record_losses_*`` — persistence layer: idempotent INSERT OR
   REPLACE on the (chunk_id, stage, canonicalization_version,
   loss_kind) primary key, sample_hash populated in PII-safe mode,
   policy hash decoupling.
3. ``test_ingest_*`` — end-to-end: ingesting a fake-wikipedia document
   populates adapter_loss_reports; toggling loss_report_enabled does
   NOT change cache_key shape (governance_policy_hash split per §3.7).

Loss-mode-aware byte-conservation property test (§3.8) lives at the
bottom of the file.
"""

from __future__ import annotations

from typing import Iterator

import pytest

from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.source import Source
from arborist.sources.loss_report import (
    LOSS_REPORT_VERSION,
    LossCollector,
    LossEvent,
    compute_loss_report_policy_hash,
    record_losses,
)
from arborist.store import connect


# ---------------------------------------------------------------------
# wikitext.to_base()
# ---------------------------------------------------------------------


def test_to_base_no_collector_is_bit_identical():
    """Regression guard: passing no collector must NOT change the prose
    output. Every prior caller (verifier, runner, query) relies on this.
    """
    from arborist.wikitext import to_base

    src = (
        "Cloud is great.<ref>Smith 2010, p. 5.</ref> "
        "He wields a [[Buster Sword|sword]]. "
        "[[File:cloud.jpg|thumb|250px|Cloud Strife]]"
        "[[Category:Final Fantasy VII characters]]"
    )
    no_collector = to_base(src)
    with_collector = to_base(src, loss_collector=LossCollector())
    assert no_collector == with_collector


def test_to_base_emits_ref_tag_event():
    from arborist.wikitext import BASE_VERSION, to_base

    collector = LossCollector()
    out = to_base(
        "Hello.<ref>citation</ref> World.",
        loss_collector=collector,
    )
    assert "citation" not in out
    events = {e.loss_kind: e for e in collector.events()}
    assert "ref_tag" in events
    assert events["ref_tag"].loss_mode == "pure_drop"
    assert events["ref_tag"].canonicalization_version == BASE_VERSION
    assert events["ref_tag"].occurrence_count == 1
    assert events["ref_tag"].bytes_dropped > 0


def test_to_base_emits_self_closing_ref_event():
    from arborist.wikitext import to_base

    collector = LossCollector()
    to_base("Hello.<ref name='x'/> World.", loss_collector=collector)
    kinds = {e.loss_kind for e in collector.events()}
    assert "self_closing_ref_tag" in kinds


def test_to_base_emits_namespace_link_events():
    from arborist.wikitext import to_base

    collector = LossCollector()
    to_base(
        "Body. [[File:a.jpg|thumb]] [[Image:b.png|left]] "
        "[[Category:Foo]]",
        loss_collector=collector,
    )
    kinds = {e.loss_kind: e for e in collector.events()}
    assert "file_link" in kinds
    assert "image_link" in kinds
    assert "category_link" in kinds
    assert all(kinds[k].loss_mode == "pure_drop" for k in (
        "file_link", "image_link", "category_link"
    ))


def test_to_base_lengths_recorded_when_collected():
    from arborist.wikitext import to_base

    collector = LossCollector()
    raw = "Hello.<ref>cite</ref> World."
    to_base(raw, loss_collector=collector)
    expected_input = len(raw.encode("utf-8"))
    events = collector.events()
    assert events
    assert all(e.input_length_bytes == expected_input for e in events)
    assert all(e.output_length_bytes is not None for e in events)
    assert all(
        e.output_length_bytes <= e.input_length_bytes for e in events
    )


def test_to_base_strip_code_emits_transform_event():
    """Surviving wikilinks rewrite to display text; that's a transform,
    not a deletion. Verifies the 'transform' loss_mode is emitted with
    a positive bytes_dropped delta."""
    from arborist.wikitext import to_base

    collector = LossCollector()
    to_base("See [[Cloud Strife|Cloud]].", loss_collector=collector)
    kinds = {e.loss_kind: e for e in collector.events()}
    assert "strip_code_transform" in kinds
    assert kinds["strip_code_transform"].loss_mode == "transform"
    assert kinds["strip_code_transform"].bytes_dropped > 0


def test_to_base_excerpt_disabled_keeps_hash():
    """PII-safe mode: excerpt is None but sample_hash stays populated."""
    from arborist.wikitext import to_base

    collector = LossCollector(excerpts_enabled=False)
    to_base("Hello.<ref>secret citation text</ref>", loss_collector=collector)
    events = [e for e in collector.events() if e.loss_kind == "ref_tag"]
    assert events
    ev = events[0]
    assert ev.sample_excerpt is None
    assert ev.sample_hash is not None
    assert len(ev.sample_hash) == 32  # sha256


def test_to_base_excerpt_truncates_at_byte_limit():
    from arborist.wikitext import to_base

    long_ref = "x" * 1000
    collector = LossCollector(max_excerpt_bytes=50)
    to_base(f"Body.<ref>{long_ref}</ref>", loss_collector=collector)
    events = [e for e in collector.events() if e.loss_kind == "ref_tag"]
    assert events
    ev = events[0]
    assert ev.sample_excerpt is not None
    assert len(ev.sample_excerpt.encode("utf-8")) <= 50


# ---------------------------------------------------------------------
# Persistence layer
# ---------------------------------------------------------------------


def test_record_losses_idempotent(tmp_path):
    """Re-running the same normalizer over the same chunk_id is a no-op.
    PK is (chunk_id, stage, canonicalization_version, loss_kind)."""
    db_path = tmp_path / "loss.db"
    conn = connect(db_path)
    try:
        # Seed a chunk row so document_root FK doesn't trip us up.
        conn.execute("BEGIN")
        conn.execute(
            "INSERT INTO documents (document_root, document_uri, source_type, "
            " title, chunking_version, canonicalization_version, schema_version, "
            " ingest_ts) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            ("aa" * 32, "test://x", "fake", "x", "tok-512-v1", "norm-v1", "v9.8.0", 0),
        )
        conn.execute(
            "INSERT INTO chunks (chunk_id, document_root, idx, leaf_hash) "
            "VALUES (?, ?, ?, ?)",
            (1, "aa" * 32, 0, "bb" * 32),
        )
        conn.execute("COMMIT")

        ev = LossEvent(
            stage="wikitext_base",
            canonicalization_version="wikitext-base-v1",
            loss_kind="ref_tag",
            loss_mode="pure_drop",
            bytes_dropped=42,
            occurrence_count=3,
        )
        policy_hash = compute_loss_report_policy_hash(
            enabled=True, excerpts=True, max_excerpt_bytes=200
        )
        conn.execute("BEGIN")
        record_losses(
            conn,
            chunk_id=1,
            document_root="aa" * 32,
            events=[ev],
            loss_report_policy_hash=policy_hash,
        )
        record_losses(
            conn,
            chunk_id=1,
            document_root="aa" * 32,
            events=[ev],
            loss_report_policy_hash=policy_hash,
        )
        conn.execute("COMMIT")

        n = conn.execute(
            "SELECT COUNT(*) FROM adapter_loss_reports WHERE chunk_id = ?",
            (1,),
        ).fetchone()[0]
        assert n == 1  # INSERT OR REPLACE — exactly one row
    finally:
        conn.close()


def test_loss_report_policy_hash_is_pure():
    """Same inputs → same hash; one input flips → hash flips."""
    h1 = compute_loss_report_policy_hash(
        enabled=True, excerpts=True, max_excerpt_bytes=200
    )
    h2 = compute_loss_report_policy_hash(
        enabled=True, excerpts=True, max_excerpt_bytes=200
    )
    h3 = compute_loss_report_policy_hash(
        enabled=False, excerpts=True, max_excerpt_bytes=200
    )
    h4 = compute_loss_report_policy_hash(
        enabled=True, excerpts=False, max_excerpt_bytes=200
    )
    h5 = compute_loss_report_policy_hash(
        enabled=True, excerpts=True, max_excerpt_bytes=100
    )
    assert h1 == h2
    assert h1 != h3
    assert h1 != h4
    assert h1 != h5
    assert len(h1) == 64


def test_loss_report_policy_version_pinned():
    """Tripwire: bumping LOSS_REPORT_VERSION moves all loss policy
    hashes — that's the intended semantics for a versioned canonicalizer
    of policy. Fail loudly so an accidental rename is caught."""
    assert LOSS_REPORT_VERSION == "loss-report-v1"


# ---------------------------------------------------------------------
# Ingest path integration
# ---------------------------------------------------------------------


class _FakeWikiSource(Source):
    """A wikipedia-shaped fake source so the ingest path runs to_base()
    with a collector. source_type starts with 'wikipedia_' which is the
    routing key in arborist.ingest."""

    source_type = "wikipedia_test"

    def __init__(self, docs: list[Document]):
        self.docs = docs

    def iter_documents(self) -> Iterator[Document]:
        yield from self.docs


def _wikipedia_doc(uri: str, content: str) -> Document:
    return Document(
        uri=uri,
        content=content,
        source_type="wikipedia_test",
        title=uri.rsplit("/", 1)[-1],
    )


def test_ingest_populates_adapter_loss_reports(tmp_path):
    src = _FakeWikiSource([
        _wikipedia_doc(
            "wiki://Cloud",
            "Cloud is great.<ref>Smith 2010</ref> "
            "He wields a [[Buster Sword|sword]]. "
            "[[File:cloud.jpg|thumb]] "
            "[[Category:Final Fantasy VII characters]] "
            "additional body prose to make sure the chunk has substance",
        ),
    ])
    db_path = tmp_path / "wiki.db"
    conn = connect(db_path)
    try:
        ingest_source(conn, src, loss_report_enabled=True)
        rows = conn.execute(
            "SELECT loss_kind, loss_mode, bytes_dropped, occurrence_count "
            "FROM adapter_loss_reports ORDER BY loss_kind"
        ).fetchall()
        kinds = {r["loss_kind"] for r in rows}
        # Must capture the four pure_drop classes from the fixture.
        assert "ref_tag" in kinds
        assert "file_link" in kinds
        assert "category_link" in kinds
        # Plus a transform delta from strip_code (rewriting wikilinks).
        assert "strip_code_transform" in kinds
        # All recorded rows have positive bytes_dropped.
        for r in rows:
            assert r["bytes_dropped"] > 0
    finally:
        conn.close()


def test_ingest_disable_loss_report_writes_no_rows(tmp_path):
    """With loss_report_enabled=False, adapter_loss_reports stays empty.
    Sidecar — toggling MUST NOT alter ingest semantics elsewhere."""
    src = _FakeWikiSource([
        _wikipedia_doc(
            "wiki://Tifa",
            "Body.<ref>cite</ref> "
            "[[File:tifa.jpg|thumb]] more prose words here for the chunk",
        ),
    ])
    db_path = tmp_path / "wiki.db"
    conn = connect(db_path)
    try:
        ingest_source(conn, src, loss_report_enabled=False)
        n = conn.execute(
            "SELECT COUNT(*) FROM adapter_loss_reports"
        ).fetchone()[0]
        assert n == 0
    finally:
        conn.close()


def test_ingest_loss_toggle_does_not_change_chunks_content(tmp_path):
    """document_root MUST be byte-identical across loss-report toggle —
    chunks.content stays raw wikitext, no chunking_version bump."""
    src_a = _FakeWikiSource([
        _wikipedia_doc(
            "wiki://Aerith",
            "Body.<ref>cite</ref> [[File:a.jpg]] additional words",
        ),
    ])
    src_b = _FakeWikiSource([
        _wikipedia_doc(
            "wiki://Aerith",
            "Body.<ref>cite</ref> [[File:a.jpg]] additional words",
        ),
    ])
    db_a = tmp_path / "a.db"
    db_b = tmp_path / "b.db"
    conn_a = connect(db_a)
    conn_b = connect(db_b)
    try:
        ingest_source(conn_a, src_a, loss_report_enabled=True)
        ingest_source(conn_b, src_b, loss_report_enabled=False)
        root_a = conn_a.execute(
            "SELECT document_root FROM documents"
        ).fetchone()["document_root"]
        root_b = conn_b.execute(
            "SELECT document_root FROM documents"
        ).fetchone()["document_root"]
        assert root_a == root_b
    finally:
        conn_a.close()
        conn_b.close()


# ---------------------------------------------------------------------
# Loss-mode-aware byte conservation property test (ticket §3.8)
# ---------------------------------------------------------------------


def test_byte_conservation_pure_drop_only(tmp_path):
    """For each (chunk_id, stage), sum(bytes_dropped) over loss_mode
    pure_drop / quarantine MUST be <= input_length_bytes. transform /
    normalize rows are excluded — they record byte deltas, not
    deletions (e.g. [[Velociraptor]] -> Velociraptor)."""
    src = _FakeWikiSource([
        _wikipedia_doc(
            "wiki://Property",
            (
                "Body prose with markup. "
                "<ref>cite one</ref> "
                "<ref>cite two longer text</ref> "
                "[[File:img.jpg|thumb|250px|caption]] "
                "[[Image:other.png|left]] "
                "[[Category:Things]] "
                "See [[Cloud Strife|Cloud]]. "
                "More body prose to give the chunk weight."
            ),
        ),
    ])
    db_path = tmp_path / "property.db"
    conn = connect(db_path)
    try:
        ingest_source(conn, src, loss_report_enabled=True)
        rows = conn.execute(
            "SELECT chunk_id, stage, "
            "       SUM(CASE WHEN loss_mode IN ('pure_drop','quarantine') "
            "                THEN bytes_dropped ELSE 0 END) AS pure_bytes, "
            "       MAX(input_length_bytes) AS input_bytes "
            "FROM adapter_loss_reports "
            "GROUP BY chunk_id, stage"
        ).fetchall()
        assert rows
        for r in rows:
            assert r["pure_bytes"] <= r["input_bytes"], (
                f"chunk={r['chunk_id']} stage={r['stage']} "
                f"pure_bytes={r['pure_bytes']} > input_bytes={r['input_bytes']}"
            )
    finally:
        conn.close()