#!/usr/bin/env python3
import argparse
import json
import subprocess
import sys


def run_fio(directory, nblocks, iodepth, force_async,
            size,
            runtime):
    bs = nblocks * 8 * 1024
    cmd = [
        "fio",
        f"--directory={directory}",
        f"--size={size}",
        "--name=read",
        "--invalidate=0",
        "--rw=read",
        "--direct=0",
        "--buffered=1",
        "--time_based=1",
        f"--runtime={runtime}",
        "--ioengine=io_uring",
        f"--iodepth={iodepth}",
        f"--force_async={force_async}",
        f"--bs={bs}",
        "--output-format=json",
    ]

    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
    return json.loads(result.stdout)


def extract_metrics(data):
    """
    Extract bandwidth (GiB/s) and average latency (µs) from fio JSON.
    fio JSON reports.:
    """
    read_stats = data["jobs"][0]["read"]

    bw_kibs = read_stats["bw"]              # KiB/s
    bw_gibs = bw_kibs / 1024**2            # KiB/s → GiB/s

    lat_ns = read_stats["lat_ns"]["mean"]  # nanoseconds
    lat_usec = lat_ns / 1000.0             # → µs

    return bw_gibs, lat_usec


def main():
    parser = argparse.ArgumentParser(
        description="Run fio sequential read benchmarks across parameter combos."
    )
    parser.add_argument(
        "--directory", default="/srv/fio",
        help="fio test directory (default: /srv/fio)",
    )
    parser.add_argument(
        "--size", default="4GiB",
        help="fio file size (default: 4GiB)",
    )
    parser.add_argument(
        "--runtime", type=int, default=1,
        help="Seconds per test (default: 1)",
    )
    parser.add_argument(
        "--nblocks", type=int, nargs="+",
        default=[1, 2, 4, 8, 16, 32, 64, 128],
        help="Block-count values to test (bs = nblocks * 8 KiB)",
    )
    parser.add_argument(
        "--iodepths", type=int, nargs="+",
        default=[1, 2, 4, 8, 16, 32],
        help="iodepth values to test",
    )
    args = parser.parse_args()

    print("nblocks\tiod\tasync\tbw_gib_s\tlat_usec")

    for nblocks in args.nblocks:
        for iodepth in args.iodepths:
            for force_async in [0, 1]:
                try:
                    data = run_fio(
                        directory=args.directory,
                        nblocks=nblocks,
                        iodepth=iodepth,
                        force_async=force_async,
                        size=args.size,
                        runtime=args.runtime,
                    )
                    bw_gibs, lat_usec = extract_metrics(data)
                    print(f"{nblocks}\t{iodepth}\t{force_async}\t{bw_gibs:.4f}\t{lat_usec:.4f}")
                    sys.stdout.flush()
                except subprocess.CalledProcessError as exc:
                    print(f"# ERROR nblocks={nblocks} iod={iodepth} async={force_async}: {exc}",
                          file=sys.stderr)
                    sys.exit(1)


if __name__ == "__main__":
    main()