#!/usr/bin/env python3
"""Parse raw sglang.bench_serving JSONL outputs for one backend.

Reads JSONL files like {backend}_{label}_MMDD_concurrency_inputlen_outputlen.jsonl
and updates results.json + report.md in the given result root.

Usage:
    python3 parse_backend.py <result_root> [--backend sglang|vllm]
"""
import argparse
import json
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_metrics(data: dict) -> dict:
    completed = data.get("completed", 0)
    total = len(data.get("input_lens", []))
    failed = total - completed if total > 0 else 0
    duration_s = data.get("duration", 0.0)

    return {
        "success": completed,
        "failed": failed,
        "duration_s": duration_s,
        "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_token_throughput": data.get("total_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 scenario_name(concurrency: int, input_len: int, output_len: int) -> str:
    return f"c{concurrency}_i{input_len}_o{output_len}"


def slo_status(metrics: dict, ttft_limit_ms: float = 3000.0, tpot_limit_ms: float = 50.0) -> dict:
    """Check SLO: TTFT P95 < limit, TPOT mean < limit. Defaults to S2 tier."""
    ttft_ok = metrics["ttft_ms"]["p95"] < ttft_limit_ms
    tpot_ok = metrics["tpot_ms"]["mean"] < tpot_limit_ms
    if ttft_ok and tpot_ok:
        mark = "✅"
    elif ttft_ok or tpot_ok:
        mark = "⚠️"
    else:
        mark = "❌"
    return {
        "ttft_p95_ok": ttft_ok,
        "tpot_mean_ok": tpot_ok,
        "overall": mark,
    }


def parse_gpu_memory_csv(jsonl_path: Path) -> dict | None:
    """Parse a paired nvidia-smi CSV for GPU memory/utilization statistics.

    The CSV is expected to live in a sibling `gpu_logs/` directory or in the same
    `raw_outputs/` directory, named `gpu_mem_c<conc>_i<isl>_o<dsl>.csv`.
    """
    result_root = jsonl_path.parent.parent
    scenario_id = jsonl_path.stem.split("_", 2)[2]  # e.g. c1_i1024_o128
    csv_name = f"gpu_mem_{scenario_id}.csv"

    csv_path = None
    for candidate in (
        result_root / "gpu_logs" / csv_name,
        result_root / "raw_outputs" / csv_name,
    ):
        if candidate.exists() and candidate.stat().st_size > 0:
            csv_path = candidate
            break

    if csv_path is None:
        return None

    per_gpu = {}
    total_mb = None
    try:
        with open(csv_path, "r", encoding="utf-8") as f:
            header = f.readline()
            if not header.strip():
                return None
            for line in f:
                line = line.strip()
                if not line:
                    continue
                parts = [p.strip() for p in line.split(",")]
                if len(parts) < 5:
                    continue
                idx = parts[1]
                try:
                    used = float(parts[2].split()[0])
                    total = float(parts[3].split()[0])
                    util = float(parts[4].split()[0])
                except (ValueError, IndexError):
                    continue
                per_gpu.setdefault(idx, {"used": [], "util": []})
                per_gpu[idx]["used"].append(used)
                per_gpu[idx]["util"].append(util)
                if total_mb is None:
                    total_mb = total
    except Exception:
        return None

    if not per_gpu or total_mb is None:
        return None

    peak_used = max(max(g["used"]) for g in per_gpu.values())
    avg_used = sum(sum(g["used"]) / len(g["used"]) for g in per_gpu.values()) / len(per_gpu)
    peak_util = max(max(g["util"]) for g in per_gpu.values())

    return {
        "peak_used_mb": peak_used,
        "avg_used_mb": avg_used,
        "peak_utilization_pct": peak_util,
        "memory_total_mb": total_mb,
    }


def generate_report(result_root: Path, backend: str, scenarios: list[dict], metadata: dict | None) -> None:
    report_path = result_root / "report.md"
    model = metadata.get("model", "unknown") if metadata else "unknown"
    hardware = metadata.get("hardware", "unknown") if metadata else "unknown"

    with open(report_path, "w", encoding="utf-8") as f:
        f.write(f"# {hardware} {backend.upper()} Benchmark Report\n\n")
        f.write(f"- Result root: `{result_root}`\n")
        f.write(f"- Model: `{model}`\n")
        f.write(f"- Backend: {backend.upper()}\n")
        f.write(f"- Benchmark client: `sglang.bench_serving --backend {backend}`\n\n")

        f.write("## Results\n\n")
        f.write("| Scenario | Phase | Concurrency | Input | Output | Duration(s) | Success | Req/s | In tok/s | Out tok/s | Total tok/s | Mean TTFT(ms) | P95 TTFT(ms) | P99 TTFT(ms) | Mean TPOT(ms) | P95 TPOT(ms) | P99 TPOT(ms) | Mean E2E(ms) | P95 E2E(ms) | P99 E2E(ms) | Peak GPU mem | SLO |\n")
        f.write("|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|\n")

        for s in scenarios:
            if "metrics" not in s:
                # Manually-recorded skipped/failed scenario without metrics;
                # listed separately below instead of crashing the report.
                continue
            cfg = s["config"]
            m = s["metrics"]
            slo = s.get("slo_status", {}).get("overall", "")
            gpu = m.get("gpu_memory")
            if gpu:
                gpu_str = f"{gpu['peak_used_mb']:.0f}/{gpu['memory_total_mb']:.0f} MiB ({100*gpu['peak_used_mb']/gpu['memory_total_mb']:.1f}%)"
            else:
                gpu_str = "-"
            f.write(
                f"| {s['name']} | {cfg['phase']} | {cfg['concurrency']} | {cfg['input_len']} | {cfg['output_len']} | "
                f"{m['duration_s']:.2f} | {m['success']} | {m['request_throughput']:.2f} | "
                f"{m['input_token_throughput']:.2f} | {m['output_token_throughput']:.2f} | "
                f"{m['total_token_throughput']:.2f} | "
                f"{m['ttft_ms']['mean']:.2f} | {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']['mean']:.2f} | {m['e2e_ms']['p95']:.2f} | {m['e2e_ms']['p99']:.2f} | {gpu_str} | {slo} |\n"
            )
        f.write("\n")
        f.write("SLO: S2 tier — TTFT P95 < 3000ms, TPOT mean < 50ms. ✅ pass, ⚠️ partial, ❌ fail.\n\n")

        skipped = [s for s in scenarios if "metrics" not in s]
        if skipped:
            f.write("## Skipped or failed scenarios\n\n")
            f.write("| Scenario | Status | Note |\n")
            f.write("|---|---|---|\n")
            for s in skipped:
                f.write(f"| {s['name']} | {s.get('status', '')} | {s.get('note', '')} |\n")
            f.write("\n")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("result_root", type=Path)
    parser.add_argument("--backend", default=None, choices=["sglang", "vllm"])
    args = parser.parse_args()

    result_root = args.result_root
    raw_dir = result_root / "raw_outputs"
    results_json = result_root / "results.json"

    if not raw_dir.exists():
        raise SystemExit(f"raw_outputs directory not found: {raw_dir}")

    backend = args.backend
    if backend is None:
        for p in raw_dir.iterdir():
            if p.name.startswith("sglang_"):
                backend = "sglang"
                break
            if p.name.startswith("vllm_"):
                backend = "vllm"
                break
    if backend is None:
        raise SystemExit("Could not infer backend from raw outputs")

    metadata = None
    old_scenarios = []
    if results_json.exists():
        with open(results_json, "r", encoding="utf-8") as f:
            try:
                existing = json.load(f)
                metadata = existing.get("metadata")
                old_scenarios = existing.get("scenarios", [])
            except json.JSONDecodeError:
                pass

    scenarios = []
    for jsonl_path in sorted(raw_dir.glob(f"{backend}_*.jsonl")):
        parts = jsonl_path.stem.split("_")
        if len(parts) < 5:
            continue

        label = parts[1]
        if label == "sharegpt":
            phase = "sharegpt"
            dataset = "sharegpt"
        else:
            phase = label
            dataset = "random"

        try:
            concurrency, input_len, output_len = int(parts[-3]), int(parts[-2]), int(parts[-1])
        except ValueError:
            continue

        data = parse_jsonl(jsonl_path)
        if data is None:
            continue

        metrics = compute_metrics(data)
        metrics["gpu_memory"] = parse_gpu_memory_csv(jsonl_path)
        scenario = {
            "name": scenario_name(concurrency, input_len, output_len),
            "config": {
                "phase": phase,
                "concurrency": concurrency,
                "input_len": input_len,
                "output_len": output_len,
                "dataset": dataset,
                "num_prompts": metrics["success"] + metrics["failed"],
            },
            "metrics": metrics,
            "slo_status": slo_status(metrics),
            "raw_file": str(jsonl_path),
        }
        scenarios.append(scenario)

    if not scenarios and not old_scenarios:
        print("No benchmark outputs found to parse")
        return

    # Preserve any manually-recorded skipped/failed scenarios (e.g. optional
    # combinations that OOMed) that do not have a fresh raw output.
    parsed_names = {s["name"] for s in scenarios}
    for old in old_scenarios:
        if old.get("status") and old["name"] not in parsed_names:
            scenarios.append(old)
    scenarios.sort(key=lambda s: (
        s["config"]["input_len"],
        s["config"]["output_len"],
        s["config"]["concurrency"],
    ))

    if results_json.exists():
        with open(results_json, "r", encoding="utf-8") as f:
            data = json.load(f)
        data["scenarios"] = scenarios
        with open(results_json, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    generate_report(result_root, backend, scenarios, metadata)
    print(f"Parsed {len(scenarios)} scenarios into {result_root}/report.md")


if __name__ == "__main__":
    main()