#!/usr/bin/env python3
import argparse, json, os, re, resource, subprocess, sys, threading, time

PAGE = resource.getpagesize()

def evict(path):
    fd = os.open(path, os.O_RDONLY)
    try:
        sz = os.fstat(fd).st_size
        step = 1 << 30
        off = 0
        while off < sz:
            n = min(step, sz - off)
            os.posix_fadvise(fd, off, n, os.POSIX_FADV_DONTNEED)
            off += n
    finally:
        os.close(fd)

def read_status(pid):
    d = {}
    try:
        with open(f"/proc/{pid}/status") as f:
            for line in f:
                k, _, v = line.partition(":")
                d[k.strip()] = v.strip()
    except OSError:
        return None
    return d

def kb(s):
    if s is None:
        return 0
    return int(s.split()[0]) * 1024

def swap_used():
    with open("/proc/meminfo") as f:
        for line in f:
            if line.startswith("SwapTotal:"):
                tot = int(line.split()[1])
            elif line.startswith("SwapFree:"):
                free = int(line.split()[1])
    return (tot - free) * 1024

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--bin", required=True)
    ap.add_argument("--levels-mb", type=int, nargs="+", required=True)
    ap.add_argument("--threads", type=int, default=8)
    ap.add_argument("--ctx", type=int, default=2048)
    ap.add_argument("--ubatch", type=int, default=128)
    ap.add_argument("--n-predict", type=int, default=48)
    ap.add_argument("--prompt", default="Write a detailed paragraph about mixture-of-experts routing.")
    ap.add_argument("--out-json", required=True)
    ap.add_argument("--evidence-dir", required=True)
    args = ap.parse_args()

    os.makedirs(args.evidence_dir, exist_ok=True)
    records = []
    for mb in args.levels_mb:
        lim = mb * 1024 * 1024
        evict(args.model)
        swap0 = swap_used()
        cmd = ["prlimit", f"--data={lim}", "--nofile=65536", args.bin,
               "-m", args.model, "-c", str(args.ctx), "-b", str(args.ubatch),
               "-t", str(args.threads), "-n", str(args.n_predict), "--temp", "0",
               "-no-cnv", "--no-jinja", "-p", args.prompt]
        t0 = time.time()
        p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        hwm_anon = 0
        hwm_rss = 0
        hwm_data = 0
        peak_read = 0
        infinity = 0
        while p.poll() is None:
            st = read_status(p.pid)
            if st:
                hwm_anon = max(hwm_anon, kb(st.get("RssAnon")))
                hwm_rss = max(hwm_rss, kb(st.get("VmRSS")))
                hwm_data = max(hwm_data, kb(st.get("VmData")))
            try:
                with open(f"/proc/{p.pid}/io") as f:
                    for line in f:
                        if line.startswith("read_bytes:"):
                            peak_read = max(peak_read, int(line.split()[1]))
            except OSError:
                pass
            time.sleep(0.2)
        out, err = p.communicate()
        dt = time.time() - t0
        text = (out + b"\n" + err).decode("utf-8", "replace")
        log = os.path.join(args.evidence_dir, f"upstream_{mb}m.log")
        open(log, "w").write(text)

        decode = prefill = None
        for ln in text.replace("\r", "\n").splitlines():
            m = re.search(r"(prompt eval|eval) time =.*?([\d.]+)\s*tokens per second", ln)
            if m:
                if m.group(1) == "prompt eval":
                    prefill = float(m.group(2))
                else:
                    decode = float(m.group(2))
        fail = p.returncode != 0 or decode is None
        rec = {
            "level_mb": mb,
            "limit_bytes": lim,
            "returncode": p.returncode,
            "seconds": round(dt, 2),
            "anon_hwm_gib": round(hwm_anon / (1 << 30), 3),
            "rss_hwm_gib": round(hwm_rss / (1 << 30), 3),
            "vmdata_hwm_gib": round(hwm_data / (1 << 30), 3),
            "read_bytes_gib": round(peak_read / (1 << 30), 3),
            "swap_delta_gib": round(max(0, swap_used() - swap0) / (1 << 30), 3),
            "decode_tps": decode,
            "prefill_tps": prefill,
            "result": "FAIL" if fail else "PASS",
            "log": log,
        }
        records.append(rec)
        print(json.dumps(rec), flush=True)

    json.dump(records, open(args.out_json, "w"), indent=2)
    print("wrote", args.out_json)

if __name__ == "__main__":
    main()