"""Unit tests for XOR PIR helpers (Chor-style two-server PIR)."""

from zkac_cli import pir


def test_pad_roundtrip():
    d = {
        "grant_id": "abc",
        "claimed": False,
        "eph_pk_b64": "eA==",
        "ciphertext_b64": "cA==",
    }
    p = pir.pad_grant_record(d)
    assert len(p) == pir.PIR_RECORD_BYTES
    assert pir.unpad_grant_record(p) == d


def test_xor_pir_row_recovery():
    rows = [
        {
            "grant_id": f"id{i}",
            "claimed": False,
            "eph_pk_b64": f"e{i}==",
            "ciphertext_b64": f"c{i}==",
        }
        for i in range(5)
    ]
    want_i = 3
    sa, sb = pir.pir_query_indices(len(rows), want_i)

    def fold(idxs):
        pads = [pir.pad_grant_record(rows[j]) for j in sorted(set(idxs))]
        return pir.xor_bytes_many(pads)

    xa = fold(sa)
    xb = fold(sb)
    got = pir.pir_recover(xa, xb)
    assert got == pir.pad_grant_record(rows[want_i])


def test_pir_fold_empty_is_zero_block():
    assert pir.xor_bytes_many([]) == b"\x00" * pir.PIR_RECORD_BYTES