#!/usr/bin/env python3
"""Parse spec-tokens 3 vs 5 comparison results and generate a 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()

    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 collect_results(result_root: Path) -> dict:
    """Collect results into {(scenario, concurrency, spec_tokens): metrics}."""
    results = {}
    focused_dir = result_root / "focused"
    if not focused_dir.exists():
        return results

    for scenario_dir in focused_dir.iterdir():
        if not scenario_dir.is_dir():
            continue
        scenario = scenario_dir.name
        for run_dir in scenario_dir.iterdir():
            if not run_dir.is_dir():
                continue
            run_id = run_dir.name
            # run_id format: <timestamp>_st{3,5}
            if "_st" not in run_id:
                continue
            spec_tokens = run_id.split("_st")[-1]
            for log in run_dir.glob("c*.log"):
                concurrency = int(log.stem[1:])
                metrics = parse_log(log)
                if metrics:
                    results[(scenario, concurrency, spec_tokens)] = metrics
    return results


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


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

    results = collect_results(result_root)
    if not results:
        print("No results found.", file=sys.stderr)
        sys.exit(1)

    # Group by scenario and concurrency
    keys = sorted({(s, c) for (s, c, _) in results.keys()})

    report_lines = [
        "# vllm-dspark `--spec-tokens` 对比报告",
        "",
        f"- 结果目录：`{result_root}`",
        f"- 模型：`/data/models/DeepSeek-V4-Flash-DSpark`",
        f"- 后端：vllm-dspark (TP=8, FP8 KV cache)",
        f"- 对比参数：`--spec-tokens 3` vs `--spec-tokens 5`",
        f"- Warmup：100 条",
        f"- 压测客户端：`sglang.bench_serving --backend vllm`",
        "",
        "## 核心指标对比",
        "",
        "| Scenario | Concurrency | Spec | Duration(s) | Req/s | Out tok/s | Total tok/s | Mean E2E(ms) | P95 E2E(ms) | P99 E2E(ms) | Mean TTFT(ms) | P99 TTFT(ms) | Mean TPOT(ms) | P99 TPOT(ms) |",
        "|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|",
    ]

    for scenario, concurrency in keys:
        for st in ["3", "5"]:
            m = results.get((scenario, concurrency, st), {})
            if not m:
                continue
            report_lines.append(
                f"| {scenario} | {concurrency} | {st} | {fmt(m['duration_s'])} | {fmt(m['req_throughput'])} | "
                f"{fmt(m['output_tok_throughput'])} | {fmt(m['total_tok_throughput'])} | {fmt(m['mean_e2e_ms'])} | "
                f"{fmt(m['p95_e2e_ms'])} | {fmt(m['p99_e2e_ms'])} | {fmt(m['mean_ttft_ms'])} | {fmt(m['p99_ttft_ms'])} | "
                f"{fmt(m['mean_tpot_ms'])} | {fmt(m['p99_tpot_ms'])} |"
            )

    # Summary table: best total tok/s per scenario/concurrency
    report_lines.append("")
    report_lines.append("## 吞吐 winner 统计")
    report_lines.append("")
    report_lines.append("| Scenario | Concurrency | Winner (Total tok/s) | st=3 Total tok/s | st=5 Total tok/s | 提升 |")
    report_lines.append("|---|---:|---:|---:|---:|---:|")

    for scenario, concurrency in keys:
        m3 = results.get((scenario, concurrency, "3"), {})
        m5 = results.get((scenario, concurrency, "5"), {})
        t3 = m3.get("total_tok_throughput", 0) or 0
        t5 = m5.get("total_tok_throughput", 0) or 0
        if t3 == 0 and t5 == 0:
            continue
        winner = "st=3" if t3 > t5 else "st=5"
        uplift = (max(t3, t5) / min(t3, t5) - 1) * 100 if min(t3, t5) > 0 else 0
        report_lines.append(
            f"| {scenario} | {concurrency} | {winner} | {fmt(t3)} | {fmt(t5)} | {fmt(uplift)}% |"
        )

    report_lines.append("")
    report_lines.append("## 观察与建议")
    report_lines.append("")

    # Compute overall wins
    st3_wins = 0
    st5_wins = 0
    for scenario, concurrency in keys:
        m3 = results.get((scenario, concurrency, "3"), {})
        m5 = results.get((scenario, concurrency, "5"), {})
        t3 = m3.get("total_tok_throughput", 0) or 0
        t5 = m5.get("total_tok_throughput", 0) or 0
        if t3 > t5:
            st3_wins += 1
        elif t5 > t3:
            st5_wins += 1

    report_lines.append(f"- **st=3 在 {st3_wins} 个配置中吞吐更高，st=5 在 {st5_wins} 个配置中吞吐更高。**")
    report_lines.append("- 具体结论请根据上表分析。")
    report_lines.append("")

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


if __name__ == "__main__":
    main()