#!/usr/bin/env python3
"""Summarize prompt-free vLLM scheduler/operator-profiler step records."""

from __future__ import annotations

import argparse
import json
import math
import statistics
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--cell-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def percentile(values: list[float], fraction: float) -> float:
    ordered = sorted(values)
    index = min(len(ordered) - 1, max(0, math.ceil(fraction * len(ordered)) - 1))
    return ordered[index]


def stats(values: list[float]) -> dict[str, float]:
    return {
        "min": min(values),
        "mean": statistics.fmean(values),
        "p50": percentile(values, 0.50),
        "p95": percentile(values, 0.95),
        "p99": percentile(values, 0.99),
        "max": max(values),
    }


def summarize_cell(path: Path) -> dict[str, Any]:
    streams = sorted(path.glob("opprof/*.jsonl"))
    if len(streams) != 1:
        raise ValueError(f"expected one opprof stream in {path}, got {len(streams)}")
    groups: dict[tuple[str, str], dict[str, list[float]]] = defaultdict(
        lambda: defaultdict(list)
    )
    phase_counts: Counter[str] = Counter()
    graph_counts: Counter[str] = Counter()
    prefix_queries = 0
    prefix_hits = 0
    total_records = 0
    dropped_records_max = 0
    with streams[0].open() as handle:
        for line in handle:
            record = json.loads(line)
            if not record.get("model_executed", False):
                continue
            prefill_tokens = int(record.get("prefill_tokens", 0))
            decode_tokens = int(record.get("decode_tokens", 0))
            if prefill_tokens and decode_tokens:
                phase = "true_mixed"
            elif prefill_tokens:
                phase = "pure_prefill"
            elif decode_tokens:
                phase = "pure_decode"
            else:
                phase = "empty"
            graph = str((record.get("cudagraph") or {}).get("runtime_mode", "UNKNOWN"))
            duration_ms = (
                int(record["complete_mono_ns"]) - int(record["submit_mono_ns"])
            ) / 1e6
            if duration_ms < 0:
                raise ValueError(f"negative duration in {streams[0]}")
            group = groups[(phase, graph)]
            group["submit_to_complete_ms"].append(duration_ms)
            group["scheduled_requests"].append(int(record["scheduled_requests"]))
            group["prefill_tokens"].append(prefill_tokens)
            group["decode_tokens"].append(decode_tokens)
            group["total_tokens"].append(prefill_tokens + decode_tokens)
            group["queue_waiting"].append(int(record["queues"]["waiting"]))
            group["kv_usage"].append(float(record["kv"]["usage"]))
            local_prefix = (record.get("prefix") or {}).get("local") or {}
            prefix_queries += int(local_prefix.get("queries", 0))
            prefix_hits += int(local_prefix.get("hits", 0))
            dropped_records_max = max(
                dropped_records_max, int(record.get("dropped_records_before", 0))
            )
            phase_counts[phase] += 1
            graph_counts[graph] += 1
            total_records += 1
    if total_records == 0:
        raise ValueError(f"empty opprof stream: {streams[0]}")
    summarized_groups = []
    for (phase, graph), metrics in sorted(groups.items()):
        summarized_groups.append(
            {
                "phase": phase,
                "cudagraph_runtime_mode": graph,
                "steps": len(metrics["submit_to_complete_ms"]),
                **{name: stats(values) for name, values in metrics.items()},
            }
        )
    cell = path.name
    tp_text, mns_text = cell.split("_")
    return {
        "cell": cell,
        "tensor_parallel_size": int(tp_text.removeprefix("tp")),
        "max_num_seqs": int(mns_text.removeprefix("mns")),
        "stream": str(streams[0]),
        "records": total_records,
        "phase_counts": dict(phase_counts),
        "cudagraph_runtime_mode_counts": dict(graph_counts),
        "prefix": {
            "queries": prefix_queries,
            "hits": prefix_hits,
            "hit_rate": prefix_hits / prefix_queries if prefix_queries else 0.0,
        },
        "dropped_records_before_max": dropped_records_max,
        "groups": summarized_groups,
    }


def main() -> None:
    args = parse_args()
    cells = [summarize_cell(path) for path in sorted(args.cell_root.glob("tp*_mns*"))]
    if not cells:
        raise SystemExit(f"no cells found in {args.cell_root}")
    payload = {
        "schema_version": "qwen30_vllm020_opprof_summary.v1",
        "measurement_semantics": (
            "complete_mono_ns - submit_mono_ns from the existing vLLM opprof record; "
            "this includes the instrumented submit-to-completion interval and is not "
            "claimed to be a single-kernel CUDA-event duration"
        ),
        "contains_prompt_text": False,
        "cell_root": str(args.cell_root),
        "cells": cells,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps({"cells": len(cells), "records": sum(c["records"] for c in cells)}))


if __name__ == "__main__":
    main()