#!/usr/bin/env python3
"""Parse dspark benchmark grid logs and generate a markdown report."""

import os
import re
import sys
from pathlib import Path


def parse_log(log_path: Path) -> dict:
    text = log_path.read_text(errors="ignore")
    lines = text.splitlines()

    # Find benchmark header line
    header_idx = None
    for i, line in enumerate(lines):
        if "============ Serving Benchmark Result ============" in line:
            header_idx = i
            break
    if header_idx is None:
        return {}

    section = "\n".join(lines[header_idx:])

    def get_float(pattern):
        m = re.search(pattern + r"\s+([\d.]+)", section)
        return float(m.group(1)) if m else None

    def get_int(pattern):
        m = re.search(pattern + r"\s+(\d+)", section)
        return int(m.group(1)) if m else None

    return {
        "duration_s": get_float(r"Benchmark duration \(s\):"),
        "successful_requests": get_int(r"Successful requests:"),
        "req_throughput": get_float(r"Request throughput \(req/s\):"),
        "input_tok_throughput": get_float(r"Input token throughput \(tok/s\):"),
        "output_tok_throughput": get_float(r"Output token throughput \(tok/s\):"),
        "total_tok_throughput": get_float(r"Total token throughput \(tok/s\):"),
        "mean_e2e_ms": get_float(r"Mean E2E Latency \(ms\):"),
        "p95_e2e_ms": get_float(r"P95 E2E Latency \(ms\):"),
        "p99_e2e_ms": get_float(r"P99 E2E Latency \(ms\):"),
        "mean_ttft_ms": get_float(r"Mean TTFT \(ms\):"),
        "p95_ttft_ms": get_float(r"P95 TTFT \(ms\):"),
        "p99_ttft_ms": get_float(r"P99 TTFT \(ms\):"),
        "mean_tpot_ms": get_float(r"Mean TPOT \(ms\):"),
        "p95_tpot_ms": get_float(r"P95 TPOT \(ms\):"),
        "p99_tpot_ms": get_float(r"P99 TPOT \(ms\):"),
        "mean_itl_ms": get_float(r"Mean ITL \(ms\):"),
    }


def main():
    result_root = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("/data/user1/yy/experiments/legacy_bench_results/dspark_grid_20260707-132641")
    if not result_root.exists():
        print(f"Result root not found: {result_root}", file=sys.stderr)
        sys.exit(1)

    phases = ["p1_quick", "p2_core", "p3_extension"]
    report_lines = [
        "# vllm-dspark Benchmark Grid Report",
        "",
        f"- Result root: `{result_root}`",
        f"- Model: `/data/models/DeepSeek-V4-Flash-DSpark`",
        f"- Backend: vllm-dspark (TP=8, FP8 KV cache, spec-method=dspark, spec-tokens=5)",
        f"- Benchmark client: `sglang.bench_serving --backend vllm`",
        "",
    ]

    for phase in phases:
        phase_dir = result_root / phase
        if not phase_dir.exists():
            continue

        report_lines.append(f"## {phase}")
        report_lines.append("")

        scenarios = sorted([p.name for p in phase_dir.iterdir() if p.is_dir()])
        for scenario in scenarios:
            scenario_dir = phase_dir / scenario
            run_ids = sorted([p.name for p in scenario_dir.iterdir() if p.is_dir()])
            for run_id in run_ids:
                run_dir = scenario_dir / run_id
                logs = sorted(run_dir.glob("c*.log"), key=lambda p: int(p.stem[1:]))

                report_lines.append(f"### {scenario} ({run_id})")
                report_lines.append("")
                report_lines.append("| Concurrency | Duration(s) | Success | Req/s | In tok/s | Out tok/s | Total tok/s | Mean E2E(ms) | P95 E2E(ms) | P99 E2E(ms) | Mean TTFT(ms) | P95 TTFT(ms) | P99 TTFT(ms) | Mean TPOT(ms) | P95 TPOT(ms) | P99 TPOT(ms) | Mean ITL(ms) |")
                report_lines.append("|------------|------------|---------|------|---------|----------|------------|-------------|-------------|-------------|--------------|-------------|-------------|--------------|-------------|-------------|-------------|")

                for log in logs:
                    concurrency = int(log.stem[1:])
                    metrics = parse_log(log)
                    if not metrics:
                        continue

                    def fmt(v):
                        return f"{v:.2f}" if v is not None else "N/A"

                    report_lines.append(
                        f"| {concurrency} | {fmt(metrics['duration_s'])} | {metrics['successful_requests'] or 'N/A'} | "
                        f"{fmt(metrics['req_throughput'])} | {fmt(metrics['input_tok_throughput'])} | "
                        f"{fmt(metrics['output_tok_throughput'])} | {fmt(metrics['total_tok_throughput'])} | "
                        f"{fmt(metrics['mean_e2e_ms'])} | {fmt(metrics['p95_e2e_ms'])} | {fmt(metrics['p99_e2e_ms'])} | "
                        f"{fmt(metrics['mean_ttft_ms'])} | {fmt(metrics['p95_ttft_ms'])} | {fmt(metrics['p99_ttft_ms'])} | "
                        f"{fmt(metrics['mean_tpot_ms'])} | {fmt(metrics['p95_tpot_ms'])} | {fmt(metrics['p99_tpot_ms'])} | "
                        f"{fmt(metrics['mean_itl_ms'])} |"
                    )

                report_lines.append("")

    report_text = "\n".join(report_lines)
    report_path = result_root / "report.md"
    report_path.write_text(report_text)
    print(report_text)
    print(f"\nReport saved to: {report_path}", file=sys.stderr)


if __name__ == "__main__":
    main()