#!/usr/bin/env python3
"""Generate a side-by-side comparison of SGLang and vLLM results.

Usage:
    python3 compare.py --sglang <sglang_result_root> --vllm <vllm_result_root> \
        [--output comparison.md]
"""
import argparse
import json
from collections import defaultdict
from pathlib import Path


def load_result(result_root: Path) -> dict:
    path = result_root / "results.json"
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def slo_status(ttft_p95_ms: float, tpot_mean_ms: float) -> str:
    # S2 tier targets from docs/SLO_STANDARDS.md.
    ttft_ok = ttft_p95_ms < 3000.0
    tpot_ok = tpot_mean_ms < 50.0
    if ttft_ok and tpot_ok:
        return "✅"
    if ttft_ok or tpot_ok:
        return "⚠️"
    return "❌"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--sglang", type=Path, required=True)
    parser.add_argument("--vllm", type=Path, required=True)
    parser.add_argument("-o", "--output", type=Path, default=Path("comparison.md"))
    args = parser.parse_args()

    sglang_data = load_result(args.sglang)
    vllm_data = load_result(args.vllm)

    by_scenario = defaultdict(dict)
    for data in (sglang_data, vllm_data):
        backend = data["metadata"]["engine"]
        for s in data.get("scenarios", []):
            key = s["name"]
            by_scenario[key][backend] = s

    with open(args.output, "w", encoding="utf-8") as f:
        f.write("# SGLang vs vLLM on DeepSeek-V4-Flash (H200, TP=8)\n\n")
        f.write("## Summary\n\n")
        f.write("- Model: `/data/models/DeepSeek-V4-Flash`\n")
        f.write("- Hardware: 8x NVIDIA H200 143GB\n")
        f.write("- Tensor Parallelism: 8\n")
        f.write("- Benchmark client: `sglang.bench_serving`\n")
        f.write("- SLO reference: S2 tier (TTFT P95 < 3s, TPOT < 50ms)\n\n")

        f.write("## Side-by-side results\n\n")
        f.write("| Scenario | Backend | Conc | Input | Output | Req/s | OutTok/s | TTFT P95(ms) | TTFT P99(ms) | TPOT Mean(ms) | TPOT P95(ms) | TPOT P99(ms) | E2E P99(ms) | SLO |\n")
        f.write("|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|\n")

        for scenario_name in sorted(by_scenario.keys()):
            for backend in ("sglang", "vllm"):
                s = by_scenario[scenario_name].get(backend)
                if s is None:
                    continue
                cfg = s["config"]
                m = s["metrics"]
                status = slo_status(m["ttft_ms"]["p95"], m["tpot_ms"]["mean"])
                f.write(
                    f"| {scenario_name} | {backend} | {cfg['concurrency']} | {cfg['input_len']} | {cfg['output_len']} | "
                    f"{m['request_throughput']:.2f} | {m['output_token_throughput']:.2f} | "
                    f"{m['ttft_ms']['p95']:.2f} | {m['ttft_ms']['p99']:.2f} | "
                    f"{m['tpot_ms']['mean']:.2f} | {m['tpot_ms']['p95']:.2f} | {m['tpot_ms']['p99']:.2f} | "
                    f"{m['e2e_ms']['p99']:.2f} | {status} |\n"
                )

        f.write("\n## Notes\n\n")
        f.write("- SLO check uses TTFT P95 and TPOT mean (the same criteria as `docs/SLO_STANDARDS.md`).\n")
        f.write("- A ⚠️ indicates one of the two metrics is out of target; ❌ indicates both are out.\n")

    print(f"Wrote comparison to {args.output}")


if __name__ == "__main__":
    main()