#!/usr/bin/env python3
"""Run the current solver over training files with an external timeout.

The evaluator invokes the solver on uncompressed CNF, so this script mirrors
that by decompressing each xz file to a temporary directory first.
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import tempfile
import time


def run_one(path: str, solver: str, timeout: float) -> dict:
    t0 = time.time()
    with tempfile.TemporaryDirectory() as td:
        cnf = os.path.join(td, "f.cnf")
        pbp = os.path.join(td, "out.pbp")
        out = os.path.join(td, "stdout")
        dec = subprocess.run(["xz", "-dc", path], stdout=open(cnf, "wb"))
        if dec.returncode != 0:
            return {"file": os.path.basename(path), "status": "decompress-error", "time": time.time() - t0}
        with open(out, "wb") as stdout:
            try:
                proc = subprocess.run([solver, cnf, pbp], stdout=stdout, stderr=subprocess.DEVNULL, timeout=timeout)
                code = proc.returncode
            except subprocess.TimeoutExpired:
                code = "timeout"
        first = ""
        try:
            with open(out, "rt", errors="replace") as f:
                first = f.readline().strip()
        except FileNotFoundError:
            pass
        if code == 10:
            status = "sat"
        elif code == 20:
            status = "unsat"
        elif code == 0:
            status = "unknown"
        elif code == "timeout":
            status = "timeout"
        else:
            status = "crash"
        proof_lines = None
        if os.path.exists(pbp):
            try:
                with open(pbp, "rb") as f:
                    proof_lines = sum(1 for _ in f)
            except OSError:
                proof_lines = None
        return {
            "file": os.path.basename(path),
            "status": status,
            "exit": code,
            "time": round(time.time() - t0, 3),
            "first": first,
            "proof_lines": proof_lines,
        }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--solver", default="/workspace/solver/solver")
    ap.add_argument("--timeout", type=float, default=10.0)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--size-order", action="store_true")
    ap.add_argument("paths", nargs="*")
    args = ap.parse_args()

    paths = args.paths
    if not paths:
        paths = [
            os.path.join("/data/train", p)
            for p in os.listdir("/data/train")
            if p.endswith(".cnf.xz")
        ]
    if args.size_order:
        paths.sort(key=lambda p: os.path.getsize(p))
    else:
        paths.sort()
    if args.limit:
        paths = paths[: args.limit]

    counts: dict[str, int] = {}
    for path in paths:
        res = run_one(path, args.solver, args.timeout)
        counts[res["status"]] = counts.get(res["status"], 0) + 1
        print(json.dumps(res, sort_keys=True), flush=True)
    print(json.dumps({"summary": counts, "count": len(paths)}, sort_keys=True), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
