"""Uniform solver front-ends: CaDiCaL (pysat Cadical195, in-process), Kissat (subprocess), z3 (bit-vector SMT).

Each runner returns a dict with: status in {sat, unsat, timeout, error}, wall_s, conflicts, decisions, propagations
(when available), solver-reported extras, and `model` = set of true literals (SAT) / dict of word values (z3).
Peak RSS is measured by the caller (one fresh process per cell; see cluster_driver.py).
"""
from __future__ import annotations
import os, re, shutil, subprocess, sys, tempfile, threading, time
from common import S, MASK

KISSAT_CANDIDATES = [os.environ.get("KISSAT_BIN", ""), os.path.expanduser("~/wp34/bin/kissat"), "/opt/homebrew/bin/kissat"]
CADICAL_CANDIDATES = [os.environ.get("CADICAL_BIN", ""), os.path.expanduser("~/wp34/bin/cadical"), "/opt/homebrew/bin/cadical"]


def _find(cands, name):
    for c in cands:
        if c and os.path.exists(c):
            return c
    return shutil.which(name)


def kissat_bin() -> str | None:
    return _find(KISSAT_CANDIDATES, "kissat")


def cadical_bin() -> str | None:
    return _find(CADICAL_CANDIDATES, "cadical")


def _write_dimacs(clauses, n_vars):
    with tempfile.NamedTemporaryFile("w", suffix=".cnf", delete=False) as f:
        f.write("p cnf %d %d\n" % (n_vars, len(clauses)))
        for c in clauses:
            f.write(" ".join(map(str, c)) + " 0\n")
        return f.name


_STAT_RE = re.compile(r"^c (conflicts|decisions|propagations|restarts|maximum-resident-set-size|process-time|maximum resident set size of process|total real time since initialization):\s+([0-9.]+)")
_STAT_KEYS = {"maximum-resident-set-size": "kissat_rss_bytes", "process-time": "solver_process_s",
              "maximum resident set size of process": "cadical_rss_mb", "total real time since initialization": "solver_real_s"}


def _run_binary(cmd, timeout):
    """Run a DIMACS solver binary; returns (status, model, stats, wall)."""
    t0 = time.time()
    try:
        p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 60)
        txt = p.stdout
    except subprocess.TimeoutExpired as e:
        txt = e.stdout if isinstance(e.stdout, str) else (e.stdout or b"").decode(errors="replace")
        txt += "\ns UNKNOWN\n"
    wall = time.time() - t0
    status, model, stats = "timeout", set(), {}
    for line in txt.splitlines():
        if line.startswith("s "):
            status = {"s SATISFIABLE": "sat", "s UNSATISFIABLE": "unsat"}.get(line.strip(), "timeout")
        elif line.startswith("v "):
            model.update(int(x) for x in line[2:].split() if int(x) > 0)
        else:
            m = _STAT_RE.match(line)
            if m:
                key = _STAT_KEYS.get(m.group(1), m.group(1))
                val = float(m.group(2))
                stats[key] = int(val) if val.is_integer() and key not in ("cadical_rss_mb",) else val
    return status, model, stats, wall


def run_cadical(clauses, timeout: float, n_vars: int | None = None, seed: int = 0) -> dict:
    """CaDiCaL as the standalone binary (`cadical -t <s> --seed=<seed>`): pysat's in-process `interrupt()` does
    not stop CaDiCaL 1.9.5 (tested: a 5 s limit ran for minutes), so wall-clock limits are enforced by the
    binary's real-time limit and a subprocess kill guard."""
    binp = cadical_bin()
    if binp is None:
        return {"status": "error", "error": "cadical binary not found", "wall_s": 0.0}
    if n_vars is None:
        n_vars = max(abs(l) for c in clauses for l in c)
    path = _write_dimacs(clauses, n_vars)
    try:
        status, model, st, wall = _run_binary([binp, "-t", str(int(timeout)), "--seed=%d" % seed, path], timeout)
    finally:
        os.unlink(path)
    # the CaDiCaL binary omits a statistics line when its value is 0 (e.g. an instance solved by propagation)
    dflt = 0 if status in ("sat", "unsat") else None
    out = {"status": status, "wall_s": round(wall, 3), "conflicts": st.get("conflicts", dflt), "decisions": st.get("decisions", dflt),
           "propagations": st.get("propagations", dflt), "restarts": st.get("restarts", dflt), "model": model if status == "sat" else None}
    out.update({k: v for k, v in st.items() if k.startswith("cadical_") or k.startswith("solver_")})
    return out


def run_cadical_pysat(clauses, timeout: float) -> dict:
    """In-process pysat Cadical195 (used only where the instance is known to be fast; no reliable timeout)."""
    from pysat.solvers import Cadical195
    t0 = time.time()
    with Cadical195(bootstrap_with=clauses) as s:
        timer = threading.Timer(timeout, s.interrupt)
        timer.start()
        try:
            res = s.solve_limited(expect_interrupt=True)
        finally:
            timer.cancel()
        wall = time.time() - t0
        st = s.accum_stats() or {}
        out = {"status": "sat" if res is True else ("unsat" if res is False else "timeout"),
               "wall_s": round(wall, 3), "conflicts": st.get("conflicts"), "decisions": st.get("decisions"),
               "propagations": st.get("propagations"), "restarts": st.get("restarts"), "model": None}
        if res is True:
            out["model"] = {l for l in s.get_model() if l > 0}
    return out


def run_kissat(clauses, n_vars: int, timeout: float, seed: int = 0) -> dict:
    binp = kissat_bin()
    if binp is None:
        return {"status": "error", "error": "kissat binary not found", "wall_s": 0.0}
    path = _write_dimacs(clauses, n_vars)
    try:
        status, model, st, wall = _run_binary([binp, "--time=%d" % int(timeout), "--seed=%d" % seed, "-s", path], timeout)
    finally:
        os.unlink(path)
    out = {"status": status, "wall_s": round(wall, 3), "conflicts": st.get("conflicts"), "decisions": st.get("decisions"),
           "propagations": st.get("propagations"), "restarts": st.get("restarts"), "model": model if status == "sat" else None}
    out.update({k: v for k, v in st.items() if k.startswith("kissat_") or k.startswith("solver_")})
    return out


# ---------------------------------------------------------------------------------------------------------- z3 ---
def z3_compress(z3, IV, W16, rounds: int, feedforward: bool):
    """Bit-vector model of the reference compression (W16/IV are z3 BitVec(32) or python ints)."""
    def rotr(x, k): return z3.RotateRight(x, k)
    def S0(x): return rotr(x, 2) ^ rotr(x, 13) ^ rotr(x, 22)
    def S1(x): return rotr(x, 6) ^ rotr(x, 11) ^ rotr(x, 25)
    def s0(x): return rotr(x, 7) ^ rotr(x, 18) ^ z3.LShR(x, 3)
    def s1(x): return rotr(x, 17) ^ rotr(x, 19) ^ z3.LShR(x, 10)
    bv = lambda v: z3.BitVecVal(v, 32) if isinstance(v, int) else v
    W = [bv(w) for w in W16]
    for t in range(16, rounds):
        W.append(s1(W[t - 2]) + W[t - 7] + s0(W[t - 15]) + W[t - 16])
    a, b, c, d, e, f, g, h = [bv(x) for x in IV]
    for t in range(rounds):
        T1 = h + S1(e) + ((e & f) ^ (~e & g)) + z3.BitVecVal(S.K[t], 32) + W[t]
        T2 = S0(a) + ((a & b) ^ (a & c) ^ (b & c))
        a, b, c, d, e, f, g, h = T1 + T2, a, b, c, d + T1, e, f, g
    out = [a, b, c, d, e, f, g, h]
    if feedforward:
        out = [x + bv(y) for x, y in zip(out, IV)]
    return out


def run_z3(build, timeout: float, seed: int = 0) -> dict:
    """build(z3, solver) -> dict name -> BitVec to read back on SAT."""
    import z3
    z3.set_param("sat.random_seed", seed)
    z3.set_param("smt.random_seed", seed)
    s = z3.Solver()
    s.set("timeout", int(timeout * 1000))
    names = build(z3, s)
    t0 = time.time()
    r = s.check()
    wall = time.time() - t0
    st = {}
    try:
        stats = s.statistics()
        for k in stats.keys():
            try:
                st[k] = stats.get_key_value(k)
            except Exception:
                pass
    except Exception:
        pass
    status = "sat" if r == z3.sat else ("unsat" if r == z3.unsat else "timeout")
    out = {"status": status, "wall_s": round(wall, 3),
           "conflicts": st.get("conflicts", st.get("sat conflicts")), "decisions": st.get("decisions", st.get("sat decisions")),
           "propagations": st.get("propagations", st.get("sat propagations")), "z3_stats": st, "model": None}
    if r == z3.sat:
        m = s.model()
        out["model"] = {k: m.eval(v, model_completion=True).as_long() for k, v in names.items()}
    return out