"""Gate for the Tseitin SHA-256 encoder: fixed message + IV, solved with CaDiCaL, must reproduce the reference."""
import pathlib, random, sys
ROOT = pathlib.Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
sys.path.insert(0, str(ROOT / "experiments" / "reduced_round_sha256"))
from wp34 import sha256 as S
import cnf_sha256 as C
from pysat.solvers import Cadical195


def _solve(cnf):
    with Cadical195(bootstrap_with=cnf.clauses) as s:
        assert s.solve()
        return C.model_to_set(s.get_model())


def test_encoder_matches_reference():
    rng = random.Random(2026)
    for r in (1, 4, 16, 64):
        for ff in (False, True):
            for _ in range(3):
                block = bytes(rng.getrandbits(8) for _ in range(64))
                cnf, W, IV, OUT = C.build_compression(r, feedforward=ff)
                for w, v in zip(W, S.bytes_to_words(block)):
                    cnf.fix_word(w, v)
                m = _solve(cnf)
                assert [C.decode_word(m, w) for w in OUT] == S.compress(S.H0, block, r, feedforward=ff)


def test_encoder_free_iv_matches_reference():
    rng = random.Random(7)
    for r in (3, 20):
        block = bytes(rng.getrandbits(8) for _ in range(64))
        iv = [rng.getrandbits(32) for _ in range(8)]
        cnf, W, IV, OUT = C.build_compression(r, feedforward=True, free_iv=True)
        for w, v in zip(W, S.bytes_to_words(block)):
            cnf.fix_word(w, v)
        for w, v in zip(IV, iv):
            cnf.fix_word(w, v)
        m = _solve(cnf)
        assert [C.decode_word(m, w) for w in OUT] == S.compress(iv, block, r, feedforward=True)


def test_preimage_small_instance_is_sat_and_verifies():
    rng = random.Random(3)
    block = bytes(rng.getrandbits(8) for _ in range(64))
    target = S.compress(S.H0, block, 8, feedforward=False)
    cnf, W, IV, OUT = C.build_compression(8)
    words = S.bytes_to_words(block)
    for j, w in enumerate(W):
        if j == 0:
            for i in range(16, 32):  # 16 free bits, rest fixed
                cnf.assert_eq(w[i], bool((words[0] >> i) & 1))
        else:
            cnf.fix_word(w, words[j])
    for o, v in zip(OUT, target):
        cnf.fix_word(o, v)
    m = _solve(cnf)
    found = [C.decode_word(m, w) for w in W]
    assert S.compress(S.H0, S.words_to_bytes(found), 8, feedforward=False) == target