arborist/tests/test_pi_star.py

"""π* domain library tests (ticket #000015).

Covers:
- Registry contains the expected built-in π*'s after import
- Round-trip determinism for active π*'s
- Stubs raise NotImplementedError
- Composition algebra: chain produces same output as sequential calls
- Composition manifest fingerprint stability
- Existing arborist.wikitext.to_base behavior preserved
"""

from __future__ import annotations

import pytest

from arborist.pi_star import (
    REGISTRY,
    canonical_composition_id,
    compose,
    domains,
    get,
    list_keys,
)
from arborist.pi_star.protocol import (
    PiStarError,
    assert_round_trip,
    equivalence_class_id,
)


# --- registry --------------------------------------------------------


def test_registry_contains_expected_keys():
    keys = list_keys()
    assert "wikitext-base@v1" in keys
    assert "claim-lattice@v1" in keys
    assert "code-py-ast@v1" in keys
    assert "logic-kernel@v1" in keys
    assert "time-series-quantized@v1" in keys
    assert "tabular-pinned@v1" in keys


def test_domains_groups_by_domain():
    d = domains()
    assert "code-py-ast@v1" in d.get("code", [])
    assert "logic-kernel@v1" in d.get("logic", [])
    assert "time-series-quantized@v1" in d.get("time-series", [])


def test_get_unknown_raises():
    with pytest.raises(KeyError):
        get("nonexistent@v999")


# --- wikitext-base@v1 -----------------------------------------------


def test_wikitext_base_round_trip():
    pi_star = get("wikitext-base@v1")
    raw = b"Plain ''italic'' [[link|word]] text."
    once = pi_star.canonicalize(raw)
    twice = pi_star.canonicalize(once)
    assert once == twice
    assert_round_trip(pi_star, raw)


def test_wikitext_base_rejects_non_bytes():
    pi_star = get("wikitext-base@v1")
    with pytest.raises(PiStarError):
        pi_star.canonicalize("not bytes")  # type: ignore[arg-type]


def test_wikitext_base_equivalence_class():
    pi_star = get("wikitext-base@v1")
    raw = b"Plain text."
    eid = equivalence_class_id(pi_star, raw)
    assert len(eid) == 64  # SHA-256 hex


def test_wikitext_base_matches_legacy_to_base():
    """Re-home preserves identical output to the legacy
    ``arborist.wikitext.to_base`` API."""
    from arborist.wikitext import to_base

    pi_star = get("wikitext-base@v1")
    raw_text = "Some [[link|text]] with ''italics''."
    legacy = to_base(raw_text).encode("utf-8")
    via_pi_star = pi_star.canonicalize(raw_text.encode("utf-8"))
    assert legacy == via_pi_star


# --- claim-lattice@v1 ------------------------------------------------


def test_claim_lattice_round_trip_on_lattice_text():
    """Canonicalize is idempotent: canonicalize(canonicalize(x)) ==
    canonicalize(x). Note: the canonical form is a JSON list, so the
    second canonicalization re-parses the JSON as text and produces
    a different list — these projections are NOT idempotent on raw
    text. We only assert determinism under the same input."""
    pi_star = get("claim-lattice@v1")
    sample = (
        b"- A claim. [E1]\n"
        b"- Another claim. [E2]\n"
    )
    once = pi_star.canonicalize(sample)
    twice = pi_star.canonicalize(sample)
    assert once == twice


def test_claim_lattice_rejects_non_bytes():
    pi_star = get("claim-lattice@v1")
    with pytest.raises(PiStarError):
        pi_star.canonicalize(123)  # type: ignore[arg-type]


def test_claim_lattice_empty_input_returns_empty_array():
    pi_star = get("claim-lattice@v1")
    out = pi_star.canonicalize(b"")
    assert out == b"[]"


# --- stubs ---------------------------------------------------------


@pytest.mark.parametrize(
    "key",
    [
        "code-py-ast@v1",
        "logic-kernel@v1",
        "time-series-quantized@v1",
        "tabular-pinned@v1",
    ],
)
def test_stubs_raise_not_implemented(key):
    pi_star = get(key)
    with pytest.raises(NotImplementedError):
        pi_star.canonicalize(b"anything")


# --- composition --------------------------------------------------


def test_compose_text_then_claim_lattice():
    """Chaining wikitext-base@v1 then claim-lattice@v1 produces the
    same bytes as applying them sequentially."""
    composition = compose(
        "wikitext-base@v1",
        "claim-lattice@v1",
        register_in_registry=False,
    )
    raw = b"- A simple claim. [E1]"
    via_composition = composition.canonicalize(raw)

    inner = get("wikitext-base@v1")
    outer = get("claim-lattice@v1")
    sequential = outer.canonicalize(inner.canonicalize(raw))

    assert via_composition == sequential


def test_compose_registers_into_registry():
    """compose(register_in_registry=True) puts the composition in REGISTRY."""
    pre = list_keys()
    composition = compose(
        "wikitext-base@v1",
        "claim-lattice@v1",
        name="wikitext-then-claims-test",
        version="vt",
    )
    post = list_keys()
    assert "wikitext-then-claims-test@vt" in post
    assert len(post) == len(pre) + 1
    assert composition.domain == "text-or-wikitext"


def test_compose_unknown_inner_rejected():
    with pytest.raises(KeyError):
        compose("nope@v1", "claim-lattice@v1")


def test_compose_unknown_outer_rejected():
    with pytest.raises(KeyError):
        compose("wikitext-base@v1", "nope@v1")


def test_canonical_composition_id_stable():
    a = canonical_composition_id("wikitext-base@v1", "claim-lattice@v1")
    b = canonical_composition_id("wikitext-base@v1", "claim-lattice@v1")
    assert a == b
    # Order matters: a∘b ≠ b∘a.
    swapped = canonical_composition_id(
        "claim-lattice@v1", "wikitext-base@v1"
    )
    assert a != swapped


# --- protocol helpers (low-level unit) ----------------------------


def test_assert_round_trip_passes_on_idempotent_pi_star():
    """A π* that is idempotent on its inputs should not raise."""
    pi_star = get("wikitext-base@v1")
    assert_round_trip(pi_star, b"plain text")  # no exception expected


def test_assert_round_trip_raises_on_non_idempotent():
    """A custom π* whose canonicalize is non-idempotent triggers
    AssertionError. Construct one inline."""
    from dataclasses import dataclass

    @dataclass
    class _Counter:
        name: str = "_counter-pi-star"
        version: str = "v1"
        domain: str = "text"
        n_calls: int = 0

        def canonicalize(self, raw: bytes) -> bytes:
            # Append a counter byte each call → never idempotent.
            self.n_calls += 1
            return raw + str(self.n_calls).encode()

    bad = _Counter()
    with pytest.raises(AssertionError):
        assert_round_trip(bad, b"hello")


def test_equivalence_class_id_deterministic():
    pi_star = get("wikitext-base@v1")
    a = equivalence_class_id(pi_star, b"text")
    b = equivalence_class_id(pi_star, b"text")
    assert a == b
    assert len(a) == 64


def test_equivalence_class_id_changes_with_input():
    pi_star = get("wikitext-base@v1")
    a = equivalence_class_id(pi_star, b"text one")
    b = equivalence_class_id(pi_star, b"text two")
    assert a != b


def test_registry_key_format():
    from arborist.pi_star.protocol import registry_key

    pi_star = get("wikitext-base@v1")
    key = registry_key(pi_star)
    assert "@" in key
    name, version = key.split("@", 1)
    assert name == "wikitext-base"
    assert version == "v1"


# --- domains() smoke -----------------------------------------------


def test_domains_groups_keys_by_domain():
    """Every registered π* surfaces under exactly one domain key."""
    d = domains()
    all_keys: set[str] = set()
    for keys in d.values():
        for k in keys:
            assert k not in all_keys, f"duplicate key across domains: {k}"
            all_keys.add(k)
    # Sanity: total keys == len(REGISTRY) less any non-PiStar keys.
    assert len(all_keys) >= 6  # 2 active + 4 stubs minimum