#!/usr/bin/env python3
"""Extract a compact metrics dict from a sglang.bench_serving JSONL output.

Usage:
    python3 extract_metrics.py <raw_jsonl_path>
"""
import json
import sys
from pathlib import Path


def parse_jsonl(path: Path) -> dict | None:
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                return json.loads(line)
            except json.JSONDecodeError:
                continue
    return None


def compute(data: dict) -> dict:
    return {
        "success_count": data.get("completed", 0),
        "total_count": len(data.get("input_lens", [])),
        "duration_s": data.get("duration", 0.0),
        "request_throughput": data.get("request_throughput", 0.0),
        "input_token_throughput": data.get("input_throughput", 0.0),
        "output_token_throughput": data.get("output_throughput", 0.0),
        "total_input_tokens": data.get("total_input_tokens", 0),
        "total_output_tokens": data.get("total_output_tokens", 0),
        "e2e_ms": {
            "mean": data.get("mean_e2e_latency_ms", 0.0),
            "p50": data.get("median_e2e_latency_ms", 0.0),
            "p90": data.get("p90_e2e_latency_ms", 0.0),
            "p95": data.get("p95_e2e_latency_ms", 0.0),
            "p99": data.get("p99_e2e_latency_ms", 0.0),
        },
        "ttft_ms": {
            "mean": data.get("mean_ttft_ms", 0.0),
            "p50": data.get("median_ttft_ms", 0.0),
            "p90": data.get("p90_ttft_ms", 0.0),
            "p95": data.get("p95_ttft_ms", 0.0),
            "p99": data.get("p99_ttft_ms", 0.0),
        },
        "tpot_ms": {
            "mean": data.get("mean_tpot_ms", 0.0),
            "p50": data.get("median_tpot_ms", 0.0),
            "p90": data.get("p90_tpot_ms", 0.0),
            "p95": data.get("p95_tpot_ms", 0.0),
            "p99": data.get("p99_tpot_ms", 0.0),
        },
        "itl_ms": {
            "mean": data.get("mean_itl_ms", 0.0),
            "p50": data.get("median_itl_ms", 0.0),
            "p90": data.get("p90_itl_ms", 0.0),
            "p95": data.get("p95_itl_ms", 0.0),
            "p99": data.get("p99_itl_ms", 0.0),
        },
    }


def main():
    path = Path(sys.argv[1])
    data = parse_jsonl(path)
    if data is None:
        raise SystemExit("No valid JSON found")
    print(json.dumps(compute(data), indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()