#!/usr/bin/env python3
"""Compare an es0 (shadow) vs es1 (real-state) run: TTFT, decision split,
routing flips, load distribution. Optional 3rd arg = es1 freshness jsonl."""
import json, sys, statistics, os

def load(d):
    ms = {}
    for l in open(os.path.join(d, "metrics.jsonl")):
        m = json.loads(l); ms[m["request_id"]] = m
    bd = {x["request_id"]: x for x in json.load(open(os.path.join(d, "breakdown.json")))}
    return ms, bd

def pct(xs, q):
    xs = sorted(xs)
    return xs[min(len(xs) - 1, int(q * len(xs)))] if xs else 0

def chosen(x):
    return x.get("routed_to", x.get("chosen_idx"))

d0, d1 = sys.argv[1], sys.argv[2]
m0, b0 = load(d0); m1, b1 = load(d1)

def ttfts(ms):
    return [m["ttft_s"] for m in ms.values() if not m.get("error") and m.get("ttft_s") is not None]

print("=== overall TTFT ===")
for tag, ms in [("es0/shadow", m0), ("es1/real ", m1)]:
    t = ttfts(ms)
    print(f"{tag}: {len(t)}/{len(ms)} ok  p50={pct(t,.5):.2f} p90={pct(t,.9):.2f} "
          f"p99={pct(t,.99):.2f} max={max(t):.2f} mean={statistics.mean(t):.2f}")

def byclass(ms, bd):
    cls = {}
    for rid, m in ms.items():
        if m.get("error") or m.get("ttft_s") is None: continue
        dec = bd.get(rid, {}).get("decision", "?")
        cls.setdefault(dec, []).append(m["ttft_s"])
    return cls

print("\n=== decision split (n / p90 / p99) ===")
for tag, ms, bd in [("es0", m0, b0), ("es1", m1, b1)]:
    print(f"  [{tag}]")
    for dec, ts in sorted(byclass(ms, bd).items()):
        print(f"    {dec:18s} n={len(ts):4d} p90={pct(ts,.9):7.2f} p99={pct(ts,.99):7.2f}")

common = set(b0) & set(b1)
flips = [r for r in common if chosen(b0[r]) != chosen(b1[r])]
decflip = [r for r in common if b0[r].get("decision") != b1[r].get("decision")]
print(f"\n=== routing changes (common reqs={len(common)}) ===")
print(f"  instance flips    : {len(flips)} ({100*len(flips)/max(1,len(common)):.1f}%)")
print(f"  decision-type flips: {len(decflip)} ({100*len(decflip)/max(1,len(common)):.1f}%)")

# TTFT of flipped vs non-flipped (es1 side)
fl_t = [m1[r]["ttft_s"] for r in flips if not m1[r].get("error") and m1[r].get("ttft_s") is not None]
if fl_t:
    print(f"  flipped reqs es1 TTFT: p50={pct(fl_t,.5):.2f} p90={pct(fl_t,.9):.2f} mean={statistics.mean(fl_t):.2f}")

def dist(bd):
    d = {}
    for x in bd.values():
        d[chosen(x)] = d.get(chosen(x), 0) + 1
    return dict(sorted(d.items(), key=lambda kv: str(kv[0])))
print("\n=== per-instance request count ===")
print("  es0:", dist(b0))
print("  es1:", dist(b1))

if len(sys.argv) > 3 and os.path.exists(sys.argv[3]):
    rows = [json.loads(l) for l in open(sys.argv[3]) if l.strip()]
    ages = [r["age_s"] for r in rows]
    busy = [r["age_s"] for r in rows if (r.get("num_prefilling") or 0) > 0]
    print(f"\n=== es1 feed freshness (full run, n={len(ages)}) ===")
    if ages:
        print(f"  age_s med={statistics.median(ages):.3f} p90={pct(ages,.9):.3f} max={max(ages):.3f} "
              f"stale>2s={sum(1 for a in ages if a>2)}")
    if busy:
        print(f"  during-prefill n={len(busy)} med={statistics.median(busy):.3f} max={max(busy):.3f}")