"""3b. SAT collision search on r-step SHA-256 compression (question C2/C4).

Instance: two message blocks M, M' (fresh variables each) with M != M' (OR of the XOR differences over the words
W_0..W_{min(r,16)-1} that an r-step compression reads; unused words are pinned to 0), equal r-step output. Variants: feedforward in {False, True} x iv in {standard, free} (semi-free-start: IV is 8 free
words shared by both messages). Note that with a *shared* IV the feed-forward adds the same value to both sides,
so a collision with feed-forward is mathematically the same event as without; the two CNFs nevertheless differ
(extra adders), which is what the SAT comparison measures. Found collisions are verified with
wp34.sha256.compress and saved in hex. Solvers: CaDiCaL (pysat) and Kissat; timeout 600 s per cell.
Modes as in preimage_sat.py (--list-cells / --cell / --aggregate).
"""
from __future__ import annotations
import argparse, json, pathlib, random, time
from common import S, provenance, peak_rss_mb, MASK
import cnf_sha256 as C
import solvers as SV

R_LIST = [4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24]
TIMEOUT = 600.0
VARIANTS = [("fixed", False), ("fixed", True), ("free", False), ("free", True)]  # (iv, feedforward)


def list_cells(seeds, solvers):
    cells = []
    for iv, ff in VARIANTS:
        for r in R_LIST:
            for seed in seeds:
                for sol in solvers:
                    cells.append({"r": r, "seed": seed, "solver": sol, "iv": iv, "feedforward": ff})
    return cells


def build(r: int, free_iv: bool, ff: bool):
    cnf = C.CNF()
    W1 = [cnf.word() for _ in range(16)]
    W2 = [cnf.word() for _ in range(16)]
    IV = [cnf.word() for _ in range(8)] if free_iv else [C.CNF.const_word(h) for h in S.H0]
    O1 = cnf.compress(IV, W1, r, ff)
    O2 = cnf.compress(IV, W2, r, ff)
    for x, y in zip(O1, O2):
        cnf.eq_w(x, y)
    # M != M' restricted to the words an r-step compression actually reads (W_0..W_{min(r,16)-1}); otherwise a
    # difference in an unused word is a trivial "collision".
    used = min(r, 16)
    cnf.neq_w_any(list(zip(W1[:used], W2[:used])))
    for j in range(used, 16):  # unused words are forced equal (and to zero) to remove the useless freedom
        cnf.fix_word(W1[j], 0); cnf.fix_word(W2[j], 0)
    return cnf, W1, W2, IV


def run_cell(cell: dict, timeout: float) -> dict:
    r, seed, sol = cell["r"], cell["seed"], cell["solver"]
    free_iv, ff = cell["iv"] == "free", bool(cell["feedforward"])
    t0 = time.time()
    cnf, W1, W2, IV = build(r, free_iv, ff)
    # seed-dependent symmetry breaking is deliberately NOT added; the seed only changes the solver's random seed
    # (Kissat --seed) and is recorded for repeated runs of CaDiCaL (deterministic => identical runs).
    res = {"cell": cell, "timeout_s": timeout, **cnf.stats(), "encode_s": round(time.time() - t0, 3)}
    if sol == "cadical":
        out = SV.run_cadical(cnf.clauses, timeout, n_vars=cnf.nv, seed=seed)
    else:
        out = SV.run_kissat(cnf.clauses, cnf.nv, timeout, seed=seed)
    model = out.pop("model", None)
    res.update(out)
    if out["status"] == "sat":
        m1 = [C.decode_word(model, w) for w in W1]
        m2 = [C.decode_word(model, w) for w in W2]
        iv = [C.decode_word(model, w) for w in IV]
        b1, b2 = S.words_to_bytes(m1), S.words_to_bytes(m2)
        h1 = S.compress(iv, b1, r, feedforward=ff)
        h2 = S.compress(iv, b2, r, feedforward=ff)
        res.update({"message1_hex": b1.hex(), "message2_hex": b2.hex(), "iv_hex": "".join("%08x" % v for v in iv),
                    "output_hex": "".join("%08x" % v for v in h1),
                    "verified": (h1 == h2 and b1 != b2),
                    "message_xor_difference_hex": (int.from_bytes(b1, "big") ^ int.from_bytes(b2, "big")).to_bytes(64, "big").hex(),
                    "differing_words": [j for j in range(16) if m1[j] != m2[j]]})
    res["total_wall_s"] = round(time.time() - t0, 3)
    res["peak_rss_mb"] = max(peak_rss_mb(), peak_rss_mb(children=True))
    res["hostname"] = provenance.host_info()["hostname"]
    return res


def aggregate(d: str, out: str, started: float):
    cells = [c for c in (json.loads(p.read_text()) for p in sorted(pathlib.Path(d).glob("*.json"))) if isinstance(c, dict) and "cell" in c]
    summary = {}
    for iv, ff in VARIANTS:
        for sol in ("cadical", "kissat"):
            for r in R_LIST:
                cs = [c for c in cells if c["cell"]["solver"] == sol and c["cell"]["r"] == r
                      and c["cell"]["iv"] == iv and bool(c["cell"]["feedforward"]) == ff]
                if not cs:
                    continue
                solved = [c for c in cs if c["status"] == "sat"]
                times = sorted(c["wall_s"] for c in cs)
                conf = sorted(c["conflicts"] for c in solved if c.get("conflicts") is not None)
                summary[f"{sol}:iv_{iv}:ff{int(ff)}:r{r}"] = {
                    "solver": sol, "iv": iv, "feedforward": ff, "r": r, "n_runs": len(cs), "n_sat": len(solved),
                    "n_timeout": sum(c["status"] == "timeout" for c in cs), "solved_fraction": len(solved) / len(cs),
                    "n_unsat": sum(c["status"] == "unsat" for c in cs),
                    "median_wall_s_unsat": sorted(c["wall_s"] for c in cs if c["status"] == "unsat")[sum(c["status"] == "unsat" for c in cs) // 2] if any(c["status"] == "unsat" for c in cs) else None,
                    "median_wall_s": times[len(times) // 2],
                    "median_wall_s_solved": sorted(c["wall_s"] for c in solved)[len(solved) // 2] if solved else None,
                    "median_conflicts_solved": conf[len(conf) // 2] if conf else None,
                    "all_verified": all(c.get("verified") for c in solved) if solved else None,
                    "n_vars": cs[0]["n_vars"], "n_clauses": cs[0]["n_clauses"],
                    "max_peak_rss_mb": max(c["peak_rss_mb"] for c in cs), "hosts": sorted({c["hostname"] for c in cs}),
                }
    provenance.record("sha256_collision_sat", {"r_list": R_LIST, "timeout_s": TIMEOUT, "variants": VARIANTS,
                                               "solvers": ["cadical (pysat Cadical195)", "kissat 4.0.4"]},
                      {"summary": summary, "cells": cells}, seed="per-cell", extra={"n_cells": len(cells)},
                      out=out, started=started)
    print("aggregated", len(cells), "cells ->", out)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--list-cells", action="store_true")
    ap.add_argument("--seeds", default="0,1,2")
    ap.add_argument("--solvers", default="cadical,kissat")
    ap.add_argument("--cell")
    ap.add_argument("--timeout", type=float, default=TIMEOUT)
    ap.add_argument("--aggregate")
    ap.add_argument("--out")
    ap.add_argument("--seed", type=int, default=0)
    a = ap.parse_args()
    t0 = time.time()
    if a.list_cells:
        for c in list_cells([int(x) for x in a.seeds.split(",")], a.solvers.split(",")):
            print(json.dumps(c))
    elif a.cell:
        cell = json.loads(a.cell); cell.setdefault("seed", a.seed)
        res = run_cell(cell, a.timeout)
        if a.out:
            pathlib.Path(a.out).parent.mkdir(parents=True, exist_ok=True)
            pathlib.Path(a.out).write_text(json.dumps(res, indent=1, default=str))
        print(json.dumps(res, default=str)[:1500])
    elif a.aggregate:
        aggregate(a.aggregate, a.out, t0)