#!/usr/bin/env python3
"""Per-size pure_transfer aggregate from A's send_blocks events alone.

Used when B's receive_kv events are missing (e.g. EngineCore env-var
propagation failed on the consumer host). Pure transfer time still
recoverable from the producer side.
"""
from __future__ import annotations

import argparse
import json
import statistics
from pathlib import Path


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--a-log", type=Path, required=True)
    p.add_argument("--out", type=Path, required=True)
    p.add_argument("--min-tokens", type=int, default=32,
                   help="Skip events smaller than this (drop the spurious "
                        "init-time tiny sends)")
    args = p.parse_args()

    events = []
    with args.a_log.open() as f:
        for line in f:
            try:
                events.append(json.loads(line))
            except json.JSONDecodeError:
                continue

    by_size: dict[int, list[float]] = {}
    for e in events:
        if e.get("event") != "send_blocks":
            continue
        sz_tokens = e["total_bytes"] // 98304
        if sz_tokens < args.min_tokens:
            continue
        by_size.setdefault(sz_tokens, []).append(e["duration_s"])

    summary = []
    for sz in sorted(by_size):
        durs = by_size[sz]
        sz_bytes = sz * 98304
        sz_mib = sz_bytes / (1024 * 1024)
        bw = [sz_bytes / d / 1e9 for d in durs]
        summary.append({
            "input_tokens": sz,
            "kv_mib": round(sz_mib, 1),
            "n": len(durs),
            "pure_transfer_ms_mean": round(statistics.mean(durs) * 1000, 2),
            "pure_transfer_ms_p50":  round(statistics.median(durs) * 1000, 2),
            "pure_transfer_ms_min":  round(min(durs) * 1000, 2),
            "pure_transfer_ms_max":  round(max(durs) * 1000, 2),
            "throughput_gbps_mean":  round(statistics.mean(bw), 2),
            "throughput_gbps_p50":   round(statistics.median(bw), 2),
            "throughput_gbps_max":   round(max(bw), 2),
        })

    print(f"loaded {len(events)} events; kept {sum(s['n'] for s in summary)} send_blocks")
    print()
    print(f"{'in_tok':>8} {'KV_MiB':>8} {'n':>4} "
          f"{'pure_p50':>10} {'pure_min':>10} {'pure_max':>10} "
          f"{'GB/s_p50':>10} {'GB/s_max':>10}")
    for s in summary:
        print(f"{s['input_tokens']:>8} {s['kv_mib']:>8.1f} {s['n']:>4} "
              f"{s['pure_transfer_ms_p50']:>10.1f} "
              f"{s['pure_transfer_ms_min']:>10.1f} "
              f"{s['pure_transfer_ms_max']:>10.1f} "
              f"{s['throughput_gbps_p50']:>10.2f} "
              f"{s['throughput_gbps_max']:>10.2f}")

    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps({"summary": summary}, indent=2))
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()