#!/usr/bin/env python3
"""Measure backup-compression codecs on the Cold Spare corpus.

One compress run and one decompress run per (corpus, codec). Output goes to
stdout as JSON Lines so a partial run is still usable. Timings are wall clock
around the codec process only; the corpus tars are read from and written to the
same volume, so I/O is included exactly as it would be for a real backup job.

Usage: ./compress-bench.py corpus/ > results.jsonl
"""
import json, os, subprocess, sys, time, platform, shutil

# (label, compress argv, decompress argv, file suffix)
CODECS = [
    ("gzip -6",            ["gzip", "-6", "-c"],                 ["gzip", "-dc"],  ".gz"),
    ("gzip -9",            ["gzip", "-9", "-c"],                 ["gzip", "-dc"],  ".gz"),
    ("bzip2 -9",           ["bzip2", "-9", "-c"],                ["bzip2", "-dc"], ".bz2"),
    ("xz -6",              ["xz", "-6", "-T1", "-c"],            ["xz", "-dc"],    ".xz"),
    ("xz -9",              ["xz", "-9", "-T1", "-c"],            ["xz", "-dc"],    ".xz"),
    ("xz -6 -T0",          ["xz", "-6", "-T0", "-c"],            ["xz", "-dc"],    ".xz"),
    ("zstd -1",            ["zstd", "-1", "-T1", "-c"],          ["zstd", "-dc"],  ".zst"),
    ("zstd -3",            ["zstd", "-3", "-T1", "-c"],          ["zstd", "-dc"],  ".zst"),
    ("zstd -9",            ["zstd", "-9", "-T1", "-c"],          ["zstd", "-dc"],  ".zst"),
    ("zstd -19",           ["zstd", "-19", "-T1", "-c"],         ["zstd", "-dc"],  ".zst"),
    ("zstd -19 --long=27", ["zstd", "-19", "--long=27", "-T1", "-c"],
                           ["zstd", "-dc", "--long=27"],                           ".zst"),
    ("zstd -3 -T0",        ["zstd", "-3", "-T0", "-c"],          ["zstd", "-dc"],  ".zst"),
    ("zstd -19 -T0",       ["zstd", "-19", "-T0", "-c"],         ["zstd", "-dc"],  ".zst"),
]


def run(argv, src, dst):
    """Run argv with src on stdin and dst on stdout. Returns wall seconds."""
    with open(src, "rb") as fin, open(dst, "wb") as fout:
        t = time.monotonic()
        subprocess.run(argv, stdin=fin, stdout=fout, check=True)
        return time.monotonic() - t


def machine():
    def sysctl(k):
        try:
            return subprocess.run(["sysctl", "-n", k], capture_output=True,
                                  text=True, check=True).stdout.strip()
        except Exception:
            return None
    return {
        "cpu": sysctl("machdep.cpu.brand_string") or platform.processor(),
        "cores": os.cpu_count(),
        "memory_bytes": int(sysctl("hw.memsize") or 0) or None,
        "os": f"{platform.system()} {platform.release()}",
        "python": platform.python_version(),
    }


def main():
    corpus = sys.argv[1] if len(sys.argv) > 1 else "corpus"
    scratch = os.path.join(corpus, ".bench-scratch")
    os.makedirs(scratch, exist_ok=True)
    print(json.dumps({"record": "machine", **machine()}), flush=True)

    tars = sorted(f for f in os.listdir(corpus) if f.endswith(".tar"))
    if not tars:
        sys.exit(f"no .tar files in {corpus} - run build-corpus.sh first")

    for tar in tars:
        src = os.path.join(corpus, tar)
        raw = os.path.getsize(src)
        for label, comp, decomp, suffix in CODECS:
            if shutil.which(comp[0]) is None:
                continue
            packed = os.path.join(scratch, tar + suffix)
            back = os.path.join(scratch, tar + ".restored")
            try:
                ct = run(comp, src, packed)
                size = os.path.getsize(packed)
                dt = run(decomp, packed, back)
                assert os.path.getsize(back) == raw, f"{label}/{tar} round trip changed size"
                print(json.dumps({
                    "record": "measure",
                    "corpus": tar[:-4],
                    "codec": label,
                    "raw_bytes": raw,
                    "packed_bytes": size,
                    "ratio": round(raw / size, 3),
                    "compress_s": round(ct, 3),
                    "decompress_s": round(dt, 3),
                    "compress_mbs": round(raw / ct / 1e6, 1),
                    "decompress_mbs": round(raw / dt / 1e6, 1),
                }), flush=True)
            except (subprocess.CalledProcessError, AssertionError) as e:
                print(json.dumps({"record": "error", "corpus": tar[:-4],
                                  "codec": label, "error": str(e)}), flush=True)
            finally:
                for f in (packed, back):
                    if os.path.exists(f):
                        os.remove(f)
    os.rmdir(scratch)


if __name__ == "__main__":
    main()