#!/usr/bin/env python3
"""
Merge PD-sep event logs from P and D into per-request breakdown.

Reads JSONL event logs produced by the patched mooncake_connector +
scheduler, joins events by request_id / transfer_id, and emits a
per-request CSV with the lifecycle phase durations:

  prefill_compute    : P_prefill_done.t - <P start time>
                       (we don't have explicit P start; use min of D_get_num_matched - 0
                        or use prefill duration ≈ P_zmq_received - P_prefill_done? NO.
                        P_prefill_done = when prefill finished, blocks ready.
                        We approximate prefill_compute = P_prefill_done - D_get_num_matched
                        because D and P receive the request simultaneously from proxy.)

  zmq_handshake      : P_rdma_start - P_zmq_received
                       (time from D's pull request reaching P to RDMA write start)

  rdma_transfer      : P_rdma_end - P_rdma_start
                       (pure RDMA write duration on P side)

  completion_signal  : D_recv_complete - P_rdma_end
                       (RDMA completion event back to D side)

  D_promote_delay    : D_request_promoted - D_recv_complete
                       (D scheduler step delay to wake the blocked request)

  full_pdsep_overhead: D_request_promoted - D_get_num_matched
                       (total server-side overhead from request arrival to schedulable)

Usage:
  python analyze_events.py --events-dir LOGDIR --out breakdown.csv
"""

import argparse
import json
from collections import defaultdict
from pathlib import Path

import csv


def load_events(paths):
    """Yield event dicts from multiple JSONL files."""
    for p in paths:
        with open(p) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    yield json.loads(line)
                except json.JSONDecodeError:
                    continue


def group_events(events):
    """Group events by transfer_id (preferred) or req_id."""
    by_key = defaultdict(dict)  # key -> {event_name: event}

    # First pass: figure out req_id <-> transfer_id mapping
    req_to_transfer = {}
    for ev in events:
        tid = ev.get("transfer_id", "")
        rid = ev.get("req_id", "")
        if tid and rid:
            req_to_transfer[rid] = tid

    # Second pass: assign each event to its transfer_id key
    # Need to re-iterate; load again
    return req_to_transfer


def _merge_event(slot: dict, ev: dict) -> None:
    """Add event to per-request slot. For repeating events:
       - 'end'/'complete'/'promoted' → keep the LATEST
       - everything else → keep the EARLIEST."""
    name = ev["event"]
    existing = slot.get(name)
    if existing is None:
        slot[name] = ev
        return
    is_end_like = any(s in name for s in ("rdma_end", "recv_complete", "promoted", "prefill_done"))
    if is_end_like:
        if ev["t_ns"] > existing["t_ns"]:
            slot[name] = ev
    else:
        if ev["t_ns"] < existing["t_ns"]:
            slot[name] = ev


def build_per_request(event_files):
    """Walk all events, group by transfer_id, compute breakdown."""

    # Collect all events into memory (these logs are small enough)
    all_events = list(load_events(event_files))

    # Map req_id <-> transfer_id
    req_to_transfer = {}
    for ev in all_events:
        tid = ev.get("transfer_id", "")
        rid = ev.get("req_id", "")
        if tid and rid:
            req_to_transfer[rid] = tid

    # Pre-pass: assign transfer_ids to P_zmq_received events from their `data.transfer_ids` field
    # Also collect P_zmq_received timestamps with their transfer_ids so we can link
    # P_rdma_start/end events that happen nearby.
    zmq_records = []  # list of (t_ns, [transfer_ids...])
    for ev in all_events:
        if ev["event"] == "P_zmq_received":
            tids = ev.get("data", {}).get("transfer_ids", []) or []
            if tids:
                # Synthetically tag this event with the first transfer_id for primary key.
                ev["transfer_id"] = tids[0]
                ev["_tids_in_batch"] = tids
                zmq_records.append((ev["t_ns"], tids))

    # Sort zmq_records by time so we can binary-search later
    zmq_records.sort()

    def find_zmq_batch(t_ns):
        """Find the most recent ZMQ batch whose timestamp <= t_ns and within 5 seconds."""
        best = None
        for ts, tids in zmq_records:
            if ts <= t_ns and (t_ns - ts) < 5e9:
                best = tids  # take the latest qualifying
            elif ts > t_ns:
                break
        return best

    # Tag P_rdma_start / P_rdma_end with the transfer_ids from the nearest preceding ZMQ batch
    for ev in all_events:
        if ev["event"] in ("P_rdma_start", "P_rdma_end") and not ev.get("transfer_id"):
            tids = find_zmq_batch(ev["t_ns"])
            if tids:
                ev["transfer_id"] = tids[0]
                ev["_tids_in_batch"] = tids

    # Group events by transfer_id
    by_xfer = defaultdict(dict)
    orphans = []
    for ev in all_events:
        tid = ev.get("transfer_id", "")
        rid = ev.get("req_id", "")
        # find key
        if tid:
            key = tid
        elif rid in req_to_transfer:
            key = req_to_transfer[rid]
        else:
            orphans.append(ev)
            continue
        # Also handle events that belong to multiple transfers in a single batch:
        # we replicate the event under each transfer_id key for fan-out.
        tids_in_batch = ev.get("_tids_in_batch", [tid] if tid else [])
        if len(tids_in_batch) > 1:
            for t in tids_in_batch:
                _merge_event(by_xfer[t], ev)
        else:
            _merge_event(by_xfer[key], ev)

    print(f"Loaded {len(all_events)} events, grouped into {len(by_xfer)} requests, "
          f"{len(orphans)} orphans")

    # Build per-request rows
    rows = []
    for tid, evmap in by_xfer.items():
        def t(name):
            e = evmap.get(name)
            return e["t_ns"] if e else None

        def d(name, field, default=None):
            e = evmap.get(name)
            return e["data"].get(field, default) if e else default

        row = {
            "transfer_id": tid,
            "n_events": len(evmap),
            "events_seen": ",".join(sorted(evmap.keys())),
            # data fields
            "num_local_cached": d("D_get_num_matched", "num_local_cached"),
            "prompt_tokens":    d("D_get_num_matched", "prompt_tokens"),
            "remote_total":     d("D_get_num_matched", "remote_total"),
            "delta_to_pull":    d("D_get_num_matched", "delta_to_pull"),
            "num_send_blocks":  d("P_prefill_done", "num_send_blocks"),
            "num_prompt_tokens_P": d("P_prefill_done", "num_prompt_tokens"),
            "rdma_num_ops":     d("P_rdma_end", "num_ops"),
            "rdma_bytes":       d("P_rdma_end", "bytes_total"),
        }

        # timestamps (ns → ms)
        ts = {
            "t_D_get_num_matched":   t("D_get_num_matched"),
            "t_P_prefill_done":      t("P_prefill_done"),
            "t_P_zmq_received":      t("P_zmq_received"),
            "t_P_rdma_start":        t("P_rdma_start"),
            "t_P_rdma_end":          t("P_rdma_end"),
            "t_D_recv_complete":     t("D_recv_complete"),
            "t_D_request_promoted":  t("D_request_promoted"),
        }
        row.update(ts)

        # phase durations in ms
        def dur(a, b):
            ta, tb = ts.get(a), ts.get(b)
            if ta is None or tb is None:
                return None
            return (tb - ta) / 1e6

        row["d_to_p_dispatch_ms"] = dur("t_D_get_num_matched", "t_P_zmq_received")
        row["prefill_compute_ms"] = dur("t_P_zmq_received",    "t_P_prefill_done")
        row["build_params_ms"]    = dur("t_P_prefill_done",    "t_P_rdma_start")
        row["rdma_transfer_ms"]   = dur("t_P_rdma_start",      "t_P_rdma_end")
        row["completion_sig_ms"]  = dur("t_P_rdma_end",        "t_D_recv_complete")
        row["D_promote_ms"]       = dur("t_D_recv_complete",   "t_D_request_promoted")
        row["full_overhead_ms"]   = dur("t_D_get_num_matched", "t_D_request_promoted")

        # transfer bandwidth
        rt = row["rdma_transfer_ms"]
        bytes_ = row["rdma_bytes"]
        if rt and bytes_ and rt > 0:
            row["rdma_bandwidth_gbps"] = (bytes_ * 8 / (rt / 1000)) / 1e9
        else:
            row["rdma_bandwidth_gbps"] = None

        rows.append(row)

    return rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--events", nargs="+", required=True,
                    help="One or more JSONL event log files")
    ap.add_argument("--out", default="breakdown.csv")
    args = ap.parse_args()

    paths = [Path(p) for p in args.events]
    for p in paths:
        if not p.exists():
            raise SystemExit(f"Not found: {p}")

    rows = build_per_request(paths)
    if not rows:
        print("No grouped requests found.")
        return

    # Write CSV
    fieldnames = sorted({k for r in rows for k in r.keys()})
    with open(args.out, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        w.writerows(rows)
    print(f"Wrote {len(rows)} request breakdowns → {args.out}")

    # Quick summary
    complete = [r for r in rows if r.get("full_overhead_ms") is not None]
    print(f"\nComplete-event requests: {len(complete)}/{len(rows)}")
    if complete:
        import statistics as st
        for k in ("d_to_p_dispatch_ms", "prefill_compute_ms", "build_params_ms",
                  "rdma_transfer_ms", "completion_sig_ms", "D_promote_ms",
                  "full_overhead_ms", "rdma_bandwidth_gbps",
                  "delta_to_pull", "num_local_cached", "rdma_bytes"):
            vals = [r[k] for r in complete if r.get(k) is not None]
            if vals:
                med = st.median(vals)
                print(f"  {k:<24} median={med:.1f}  n={len(vals)}")


if __name__ == "__main__":
    main()