#!/usr/bin/env python3
"""Pair before/after urus-bench results and print an A/B comparison.

Layout produced by spin_ab_urus.sh:

    <outdir>/before/<server>-<scenario>-c<cores>/result.json
    <outdir>/after/<server>-<scenario>-c<cores>/result.json

We pair by the (server, scenario, cores) key and report sustained RPS
(after/before speedup) and p50/p99/p99.9 latency deltas. The spin gate is
active only at 2-4 schedulers, so rows at cores 2 and 4 are where an effect is
expected; cores 1 and 8 are controls (gate inert) and should stay flat.
"""

import argparse
import json
import os
import re
import sys

GATED = {2, 4}  # core counts where the RFC 004 spin policy is active


def parse_latency_ms(s):
    """wrk2 prints '0.50ms' / '1.83ms' / '24.50ms' / '1.20s' / '850.00us'."""
    if s is None:
        return None
    s = str(s).strip()
    m = re.match(r"^([0-9.]+)\s*(us|ms|s)?$", s)
    if not m:
        return None
    val = float(m.group(1))
    unit = m.group(2) or "ms"
    return {"us": val / 1000.0, "ms": val, "s": val * 1000.0}[unit]


def load(outdir, sub):
    """Return {(server, scenario, cores): result_dict} for one side."""
    base = os.path.join(outdir, sub)
    out = {}
    if not os.path.isdir(base):
        return out
    for name in os.listdir(base):
        rj = os.path.join(base, name, "result.json")
        if not os.path.isfile(rj):
            continue
        m = re.match(r"^(.*)-(s\d)-c(\d+)$", name)
        if not m:
            continue
        server, scenario, cores = m.group(1), m.group(2), int(m.group(3))
        try:
            with open(rj) as f:
                out[(server, scenario, cores)] = json.load(f)
        except (OSError, json.JSONDecodeError) as e:
            print(f"warn: could not read {rj}: {e}", file=sys.stderr)
    return out


def fnum(x):
    try:
        return float(x)
    except (TypeError, ValueError):
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("outdir")
    ap.add_argument("--before", default="before")
    ap.add_argument("--after", default="after")
    args = ap.parse_args()

    before = load(args.outdir, args.before)
    after = load(args.outdir, args.after)

    keys = sorted(set(before) | set(after))
    if not keys:
        print("no results found under", args.outdir, file=sys.stderr)
        return 1

    # --- RPS table ---------------------------------------------------------
    hdr = f"{'server':<10} {'scen':<4} {'cores':>5} {'before rps':>12} {'after rps':>12} {'speedup':>8}  gate"
    print("\n== Sustained RPS (higher is better) ==")
    print(hdr)
    print("-" * len(hdr))
    for k in keys:
        srv, scn, cores = k
        b = fnum(before.get(k, {}).get("sustained_rps"))
        a = fnum(after.get(k, {}).get("sustained_rps"))
        spd = f"{a / b:.2f}x" if (a and b) else "—"
        gate = "ACTIVE" if cores in GATED else "inert"
        bs = f"{b:,.0f}" if b is not None else "—"
        as_ = f"{a:,.0f}" if a is not None else "—"
        print(f"{srv:<10} {scn:<4} {cores:>5} {bs:>12} {as_:>12} {spd:>8}  {gate}")

    # --- latency tables ----------------------------------------------------
    for pct in ("p50", "p99", "p999"):
        print(f"\n== {pct} latency, ms (lower is better) ==")
        h = f"{'server':<10} {'scen':<4} {'cores':>5} {'before':>10} {'after':>10} {'delta':>9}  gate"
        print(h)
        print("-" * len(h))
        for k in keys:
            srv, scn, cores = k
            b = parse_latency_ms(before.get(k, {}).get("latency", {}).get(pct))
            a = parse_latency_ms(after.get(k, {}).get("latency", {}).get(pct))
            if a is not None and b is not None:
                delta = f"{(a - b):+.2f}"
            else:
                delta = "—"
            gate = "ACTIVE" if cores in GATED else "inert"
            bs = f"{b:.2f}" if b is not None else "—"
            as_ = f"{a:.2f}" if a is not None else "—"
            print(f"{srv:<10} {scn:<4} {cores:>5} {bs:>10} {as_:>10} {delta:>9}  {gate}")

    # --- error sanity (non-2xx should be ~0 on both sides) -----------------
    noisy = []
    for k in keys:
        for side, d in (("before", before), ("after", after)):
            n = d.get(k, {}).get("non_2xx", 0)
            if isinstance(n, (int, float)) and n > 0:
                noisy.append((side, k, n))
    if noisy:
        print("\n== WARNING: non-2xx responses (results may be unreliable) ==")
        for side, (srv, scn, cores), n in noisy:
            print(f"  {side} {srv} {scn} c{cores}: {n} non-2xx")

    print(
        "\nReading it: spin is 'worth it' if s3 (and/or s2) RPS rises and/or "
        "tail latency drops at cores 2/4 (gate ACTIVE), with cores 1/8 (inert) "
        "staying flat. Flat-or-worse at 2/4 = not worth it on real traffic."
    )
    if "axum-mem" in {k[0] for k in keys}:
        print(
            "axum-mem is the control: it doesn't depend on smarm, so its "
            "before/after should be ~identical. Any drift there means the run "
            "environment moved, not the patch."
        )
    return 0


if __name__ == "__main__":
    sys.exit(main())