#!/usr/bin/env python3
"""Measure local SSD read bandwidth at expert-granularity (random reads) and sequential reads,
bypassing the page cache with F_NOCACHE (macOS). Writes a JSON record to benchmarks/raw/<out>.
Usage: ssd_bandwidth.py <file> --chunk-mb 3.4 --n 200 --out benchmarks/raw/gateA/ssd/<host>.json
"""
import argparse, fcntl, json, os, platform, random, socket, subprocess, time, statistics, sys

F_NOCACHE = 48  # from <sys/fcntl.h> on macOS

def read_chunks(path, offsets, size, nocache=True):
    fd = os.open(path, os.O_RDONLY)
    if nocache:
        fcntl.fcntl(fd, F_NOCACHE, 1)
    lat = []
    nbytes = 0
    t0 = time.perf_counter()
    for off in offsets:
        t = time.perf_counter()
        b = os.pread(fd, size, off)
        lat.append(time.perf_counter() - t)
        nbytes += len(b)
    tot = time.perf_counter() - t0
    os.close(fd)
    return nbytes, tot, lat

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("file")
    ap.add_argument("--chunk-mb", type=float, default=3.4)
    ap.add_argument("--n", type=int, default=200)
    ap.add_argument("--seq-mb", type=int, default=2048)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    size = int(a.chunk_mb * 1024 * 1024)
    fsize = os.path.getsize(a.file)
    rng = random.Random(a.seed)
    offsets = [rng.randrange(0, fsize - size) // 4096 * 4096 for _ in range(a.n)]
    rec = {"host": socket.gethostname(), "file": a.file, "file_bytes": fsize, "chunk_bytes": size, "n": a.n,
           "seed": a.seed, "machine": platform.machine(), "macos": platform.mac_ver()[0],
           "chip": subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip(),
           "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), "cmd": " ".join(sys.argv)}
    # random expert-sized reads, uncached
    nb, tot, lat = read_chunks(a.file, offsets, size, nocache=True)
    lat_ms = sorted(x * 1e3 for x in lat)
    rec["random_nocache"] = {"GBps": nb / tot / 1e9, "lat_ms_p50": statistics.median(lat_ms),
                              "lat_ms_p95": lat_ms[int(0.95 * len(lat_ms)) - 1], "lat_ms_p99": lat_ms[int(0.99 * len(lat_ms)) - 1],
                              "lat_ms_mean": statistics.mean(lat_ms)}
    # same offsets, page-cache allowed (second pass is warm)
    read_chunks(a.file, offsets, size, nocache=False)
    nb, tot, lat = read_chunks(a.file, offsets, size, nocache=False)
    rec["random_cached"] = {"GBps": nb / tot / 1e9, "lat_ms_p50": statistics.median(x * 1e3 for x in lat)}
    # sequential uncached
    seq = min(a.seq_mb * 1024 * 1024, fsize)
    offs = list(range(0, seq, 8 * 1024 * 1024))
    nb, tot, _ = read_chunks(a.file, offs, 8 * 1024 * 1024, nocache=True)
    rec["sequential_nocache_8MB"] = {"GBps": nb / tot / 1e9, "bytes": nb}
    os.makedirs(os.path.dirname(a.out), exist_ok=True)
    with open(a.out, "w") as f:
        json.dump(rec, f, indent=2)
    print(json.dumps(rec, indent=2))

if __name__ == "__main__":
    main()