#!/usr/bin/env python3
import argparse
import json
import lzma
import os
import subprocess
import sys
import tempfile
import time


def load_labels(path):
    if not path:
        return {}
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    return {r["name"]: r["status"] for r in data.get("results", [])}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--solver", default="/workspace/solver/solver")
    ap.add_argument("--timeout", type=float, default=30.0)
    ap.add_argument("--labels")
    ap.add_argument("instances", nargs="+")
    args = ap.parse_args()
    labels = load_labels(args.labels)
    for src in args.instances:
        name = os.path.basename(src)
        with tempfile.TemporaryDirectory(prefix="diag-local-") as td:
            cnf = os.path.join(td, "formula.cnf")
            proof = os.path.join(td, "out.pbp")
            with lzma.open(src, "rb") as inp, open(cnf, "wb") as out:
                out.write(inp.read())
            t0 = time.time()
            try:
                p = subprocess.run(
                    [args.solver, cnf, proof],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True,
                    timeout=args.timeout,
                )
                elapsed = time.time() - t0
                if p.returncode == 10:
                    status = "sat"
                elif p.returncode == 20:
                    status = "unsat"
                elif p.returncode == 0:
                    status = "unknown"
                else:
                    status = f"crash({p.returncode})"
            except subprocess.TimeoutExpired as e:
                elapsed = time.time() - t0
                status = "timeout"
                p = e
            tail = ""
            err = getattr(p, "stderr", "") or ""
            if isinstance(err, bytes):
                err = err.decode("utf-8", "replace")
            if err:
                tail = err.strip().splitlines()[-1]
            print(json.dumps({
                "file": name,
                "status": status,
                "expected": labels.get(name),
                "time": round(elapsed, 3),
                "stderr_tail": tail,
            }, separators=(",", ":")))
            sys.stdout.flush()


if __name__ == "__main__":
    main()
