arborist/tests/test_warrant_resolver.py

"""Tests for the warrant resolver (#000031 Phase 2).

Pure unit tests over the citation parser; the FTS5 resolver +
Merkle proof writer are exercised end-to-end via the
``arborist warrant-resolve`` CLI on real shards. The CLI smoke test
is documented in the ticket; this file stays offline / DB-free.
"""

from __future__ import annotations

import pytest

from arborist.qa.warrant_resolver import (
    Citation,
    parse_citation,
)


# --- "Title by Author" pattern (most common) -------------------------


def test_simple_title_by_author():
    out = parse_citation("Introduction to Mathematical Logic by Elliott Mendelson")
    assert len(out) == 1
    c = out[0]
    assert c.title == "Introduction to Mathematical Logic"
    assert c.authors == ("Elliott Mendelson",)


def test_title_with_punctuation():
    out = parse_citation("The Lambda Calculus: Its Syntax and Semantics by H.P. Barendregt")
    assert len(out) == 1
    c = out[0]
    assert c.title == "The Lambda Calculus: Its Syntax and Semantics"
    assert c.authors == ("H.P. Barendregt",)


# --- multi-author Oxford comma --------------------------------------


def test_multi_author_oxford_comma():
    out = parse_citation(
        "Classical Mechanics by Herbert Goldstein, Charles P. Poole, and John L. Safko"
    )
    assert len(out) == 1
    c = out[0]
    assert c.title == "Classical Mechanics"
    assert c.authors == ("Herbert Goldstein", "Charles P. Poole", "John L. Safko")


def test_multi_author_et_al():
    out = parse_citation("Classical Mechanics by Herbert Goldstein et al.")
    assert len(out) == 1
    c = out[0]
    assert c.title == "Classical Mechanics"
    assert c.authors == ("Herbert Goldstein",)


# --- semicolon-separated multi-citation -----------------------------


def test_semicolon_split():
    out = parse_citation(
        "Knuth TAOCP Volume 1 §1.2.6; Stanley §1.2; Brualdi §3.5"
    )
    assert len(out) == 3
    assert out[0].authors == ("Knuth",)
    assert out[1].authors == ("Stanley",)
    assert out[2].authors == ("Brualdi",)


def test_semicolon_with_year_and_section():
    out = parse_citation(
        "Knuth TAOCP Volume 1 §1.2.6 equation (13); Vandermonde 1772"
    )
    assert len(out) == 2
    assert out[0].section
    assert "§1.2.6" in out[0].section
    assert out[1].year == "1772"


# --- compact form ---------------------------------------------------


def test_compact_year_form():
    out = parse_citation("Pascal 1654")
    assert len(out) == 1
    c = out[0]
    assert c.year == "1654"
    assert "Pascal" in c.authors


def test_compact_section_only():
    out = parse_citation("Brualdi §3.5")
    assert len(out) == 1
    c = out[0]
    assert c.authors == ("Brualdi",)
    assert "§3.5" in c.section


# --- edge cases ------------------------------------------------------


def test_empty_input():
    assert parse_citation("") == []
    assert parse_citation("   ") == []


def test_year_extraction():
    out = parse_citation("Foundations of Probability by Andrey Kolmogorov 1933")
    assert len(out) == 1
    assert out[0].year == "1933"


def test_section_extraction():
    out = parse_citation("Classical Mechanics by Goldstein Volume 1 §3.2")
    assert len(out) == 1
    assert "§3.2" in out[0].section
    assert "Volume 1" in out[0].section


def test_oeis_identifier():
    out = parse_citation("OEIS A000108")
    assert len(out) == 1
    # OEIS identifiers parse as compact form — first token is "author".
    assert out[0].authors == ("OEIS",)


# --- Citation dataclass invariants ----------------------------------


def test_citation_is_empty():
    assert Citation().is_empty()
    assert not Citation(title="Foo").is_empty()
    assert not Citation(authors=("Bar",)).is_empty()


def test_raw_field_preserved():
    raw = "Some weird citation by Some Author"
    out = parse_citation(raw)
    assert out[0].raw == raw


# --- Phase 5: phrase + content-token cascade (#000039) --------------


def test_phrase_for_axiom_strips_categorical_prefix():
    from arborist.qa.warrant_resolver import _phrase_for_axiom

    assert _phrase_for_axiom("Axiom of Line Incidence") == '"line incidence"'
    assert _phrase_for_axiom("Axiom of Plane Incidence") == '"plane incidence"'
    assert _phrase_for_axiom("Theorem of Pythagoras") == '"pythagoras"' or _phrase_for_axiom("Theorem of Pythagoras") == ""


def test_phrase_for_axiom_drops_parenthetical():
    from arborist.qa.warrant_resolver import _phrase_for_axiom

    assert _phrase_for_axiom("Axiom of Side-Angle-Side (SAS)") == '"side angle side"'


def test_phrase_for_axiom_returns_empty_for_single_token():
    from arborist.qa.warrant_resolver import _phrase_for_axiom

    # Single-token axioms have nothing to phrase-match — fall through
    # to AND-join strategy.
    assert _phrase_for_axiom("Pasch's Axiom") == ""
    assert _phrase_for_axiom("Pythagorean Theorem") == ""


def test_content_tokens_filters_common_words():
    from arborist.qa.warrant_resolver import _content_tokens

    content = (
        "this axiom states that for every triangle there exists "
        "a unique line through any two points; the system follows. "
        "Triangle triangle triangle vertex vertex vertex."
    )
    toks = _content_tokens(content)
    # "triangle" appears 4 times → discriminating; "this", "that",
    # "every", "system" → common, filtered.
    assert "triangle" in toks
    assert "this" not in toks
    assert "system" not in toks
    assert "every" not in toks


def test_content_tokens_requires_count_at_least_2():
    from arborist.qa.warrant_resolver import _content_tokens

    # Singleton tokens ditched (likely typo / LaTeX residue).
    content = "betweenness betweenness consider unique helpfully"
    toks = _content_tokens(content)
    assert "betweenness" in toks
    # singletons dropped
    assert "consider" not in toks
    assert "unique" not in toks


def test_build_record_query_cascade_orders_correctly():
    from arborist.qa.warrant_resolver import (
        Citation,
        _build_record_query_cascade,
    )

    citation = Citation(
        title="The Foundations of Geometry",
        authors=("David Hilbert",),
        raw="The Foundations of Geometry by David Hilbert",
    )
    queries = _build_record_query_cascade(
        citation,
        theorem_name="Axiom of Line Incidence",
        record_content="The line line line connection between point point points axiom incidence relation."
    )
    # Phrase is first; legacy AND-join is last.
    assert queries[0] == '"line incidence"'
    # The content-token AND-join should appear in the cascade.
    assert any(" AND " in q for q in queries[1:])


# --- B-1: via_citation_alias attribution -----------------------


def test_citation_via_citation_alias_default_false():
    """Citations parsed from the original source_reference start
    with via_citation_alias=False — backward-compatible."""
    from arborist.qa.warrant_resolver import Citation, parse_citation

    cs = parse_citation("Foundations of Geometry by David Hilbert")
    assert all(c.via_citation_alias is False for c in cs)


def test_citation_explicit_via_citation_alias_flag():
    """Construct a substitute Citation as if from #000041 lookup —
    via_citation_alias=True. Used by warrant_resolve to flag chains
    that came through a citation-alias substitution rather than the
    parsed source_reference."""
    from arborist.qa.warrant_resolver import Citation

    sub = Citation(
        title="Russell IMP",
        authors=("Bertrand Russell",),
        raw="Russell IMP by Bertrand Russell",
        via_citation_alias=True,
    )
    assert sub.via_citation_alias is True


def test_resolve_chunks_propagates_via_citation_alias_to_match(tmp_path, monkeypatch):
    """When resolve_chunks is called with a Citation whose
    via_citation_alias=True is set, every ResolutionMatch it produces
    inherits via_alias=True so process_id correctly attributes the
    derivation as `warrant-resolver-v1+alias`."""
    import sqlite3
    from pathlib import Path

    from arborist.qa.warrant_resolver import (
        Citation,
        ResolutionMatch,
        resolve_chunks,
    )

    # Build a tiny fake shard cluster: shards/000.db (empty) +
    # crawl/textbook_test.db with one document + one chunk that the
    # _shard_matches_citation heuristic can find by author surname.
    shards_dir = tmp_path / "shards"
    crawl_dir = tmp_path / "crawl"
    shards_dir.mkdir()
    crawl_dir.mkdir()

    from arborist.store import SCHEMA_SQL

    main_db = shards_dir / "000.db"
    sqlite3.connect(str(main_db)).executescript(SCHEMA_SQL).close()

    sub_db = crawl_dir / "textbook_test.db"
    conn = sqlite3.connect(str(sub_db))
    conn.executescript(SCHEMA_SQL)
    conn.execute(
        "INSERT INTO documents "
        "(document_root, document_uri, source_type, kind, "
        " compression_depth, title, chunking_version, "
        " canonicalization_version, schema_version, ingest_ts) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
        ("a" * 64, "https://test/", "html", "surface", 0,
         "Test Book by Bertrand Russell", "tok-512-v1",
         "norm-v1", "v9.8.0", 0),
    )
    conn.execute(
        "INSERT INTO chunks (chunk_id, document_root, idx, leaf_hash, content) "
        "VALUES (1, ?, 0, ?, ?)",
        ("a" * 64, "h" * 64, "philosophy mathematics test sample"),
    )
    conn.execute(
        "INSERT INTO chunks_fts (rowid, content) VALUES (1, ?)",
        ("philosophy mathematics test sample",),
    )
    conn.commit()
    conn.close()

    # Citation with via_citation_alias=True (mimics what #000041
    # citation-alias lookup appends).
    sub_cit = Citation(
        title="Test Book",
        authors=("Bertrand Russell",),
        raw="Test Book by Bertrand Russell",
        via_citation_alias=True,
    )
    matches = resolve_chunks(
        sub_cit,
        shards_dir,
        theorem_name="Axiom of Test",
        record_content="philosophy mathematics test sample axiom",
        limit=3,
    )
    assert matches, "expected a match for the substitute citation"
    assert all(
        isinstance(m, ResolutionMatch) and m.via_alias is True
        for m in matches
    ), "every match from a via_citation_alias=True Citation must be flagged via_alias"

    # Sanity check the inverse — Citation without the flag → matches
    # are via_alias=False (the existing direct-cascade path).
    parsed_cit = Citation(
        title="Test Book",
        authors=("Bertrand Russell",),
        raw="Test Book by Bertrand Russell",
    )
    plain_matches = resolve_chunks(
        parsed_cit,
        shards_dir,
        theorem_name="Axiom of Test",
        record_content="philosophy mathematics test sample axiom",
        limit=3,
    )
    assert plain_matches
    assert all(m.via_alias is False for m in plain_matches), (
        "Citation parsed from source_ref (no via_citation_alias) "
        "must NOT inherit via_alias=True from the citation_alias_floor"
    )