#!/usr/bin/env python3
"""Generate report.md from results.json for the max context length experiment.

Usage:
    python3 parse_results.py <result_root>
"""
import json
import sys
from pathlib import Path


def main():
    result_root = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("results")
    results_json = result_root / "results.json"
    report_path = result_root / "report.md"

    with open(results_json, "r", encoding="utf-8") as f:
        data = json.load(f)

    meta = data["metadata"]
    cfg = data["config"]

    with open(report_path, "w", encoding="utf-8") as f:
        f.write("# H200 Max Context Length Exploration Report\n\n")
        f.write(f"- Result root: `{result_root}`\n")
        f.write(f"- Model: `{meta['model']}`\n")
        f.write(f"- Hardware: {meta['hardware']}\n")
        f.write(f"- TP: {cfg['tp']}\n")
        f.write(f"- Output len: {cfg['output_len']}, num prompts: {cfg['num_prompts']}, concurrency: {cfg['max_concurrency']}\n\n")

        f.write("## Per-attempt results\n\n")
        f.write("| Backend | Target len | Max model/context len | Status | TTFT mean(ms) | TTFT P95(ms) | E2E mean(ms) | E2E P99(ms) | Error |\n")
        f.write("|---|---:|---:|---:|---:|---:|---:|---:|---:|\n")

        max_by_backend = {}
        for s in data.get("scenarios", []):
            backend = s["backend"]
            target = s["target_len"]
            status = "✅ pass" if s["success"] else "❌ fail"
            metrics = s.get("metrics") or {}
            ttft = metrics.get("ttft_ms", {})
            e2e = metrics.get("e2e_ms", {})
            f.write(
                f"| {backend} | {target} | {s['max_model_len']} | {status} | "
                f"{ttft.get('mean', 0):.2f} | {ttft.get('p95', 0):.2f} | "
                f"{e2e.get('mean', 0):.2f} | {e2e.get('p99', 0):.2f} | "
                f"{s.get('error', '') or ''} |\n"
            )
            if s["success"] and (backend not in max_by_backend or target > max_by_backend[backend]):
                max_by_backend[backend] = target

        f.write("\n## Maximum supported input length\n\n")
        for backend in ("sglang", "vllm"):
            max_len = max_by_backend.get(backend, 0)
            f.write(f"- **{backend}**: {max_len} tokens\n")
        f.write("\n")

    print(f"Wrote report to {report_path}")


if __name__ == "__main__":
    main()