#!/usr/bin/env python3
"""
Post-process AIPerf suite results into management-ready deliverables.
Generates: summary.csv, capacity_curve.png, latency_curve.png, sales_comparison.png
"""
import json
import pathlib
import sys
from collections import defaultdict

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def get_metric(data, key, stat="avg"):
    m = data.get(key, {})
    return m.get(stat) if isinstance(m, dict) else m


def load_results(results_dir):
    results = defaultdict(lambda: defaultdict(list))
    for json_file in sorted(results_dir.rglob("profile_export_aiperf.json")):
        rel = json_file.relative_to(results_dir)
        parts = list(rel.parts)
        model = parts[0]
        scenario = parts[1]
        conc_label = parts[2] if len(parts) > 2 else "steady"
        with open(json_file) as f:
            data = json.load(f)
        results[model][scenario].append({
            "conc_label": conc_label,
            "requests": get_metric(data, "request_count", "avg"),
            "ttft_avg_ms": get_metric(data, "time_to_first_token", "avg"),
            "ttft_p99_ms": get_metric(data, "time_to_first_token", "p99"),
            "throughput_tps": get_metric(data, "output_token_throughput", "avg"),
            "itl_p99_ms": get_metric(data, "inter_token_latency", "p99"),
            "prefill_tps": get_metric(data, "prefill_throughput_per_user", "avg"),
            "duration_s": get_metric(data, "benchmark_duration", "avg"),
        })
    return results


def write_csv(results, output_path):
    rows = ["model,scenario,concurrency,ttft_avg_ms,ttft_p99_ms,throughput_tok_per_s,requests,duration_s,prefill_tok_per_s,itl_p99_ms"]
    for model in sorted(results):
        for scenario in ["sales", "capacity", "steady"]:
            for r in sorted(results[model][scenario], key=lambda x: int(x["conc_label"].replace("conc","")) if x["conc_label"].startswith("conc") else 0):
                rows.append(f"{model},{scenario},{r['conc_label']},{r['ttft_avg_ms'] or ''},{r['ttft_p99_ms'] or ''},{r['throughput_tps'] or ''},{r['requests'] or ''},{r['duration_s'] or ''},{r['prefill_tps'] or ''},{r['itl_p99_ms'] or ''}")
    output_path.parent.mkdir(exist_ok=True)
    output_path.write_text("\n".join(rows))
    print(f"CSV: {output_path}")


def plot_scenario(results, scenario, y_key, y_label, title, output_path, logscale=True):
    fig, ax = plt.subplots(figsize=(10, 6))
    for model in sorted(results):
        runs = sorted(results[model][scenario], key=lambda x: int(x["conc_label"].replace("conc","")) if x["conc_label"].startswith("conc") else 0)
        if not runs:
            continue
        concs = []
        vals = []
        for r in runs:
            try:
                c = int(r["conc_label"].replace("conc",""))
            except ValueError:
                c = r["conc_label"]
            concs.append(c)
            vals.append(r[y_key] or 0)
        ax.plot(concs, vals, marker="o", linewidth=2, label=model)
    ax.set_xlabel("Concurrency")
    ax.set_ylabel(y_label)
    ax.set_title(title, fontweight="bold")
    ax.legend()
    ax.grid(True, alpha=0.3)
    if logscale:
        ax.set_xscale("log")
    fig.tight_layout()
    fig.savefig(output_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"Plot: {output_path}")


def main():
    if len(sys.argv) < 2:
        print("Usage: plot-suite-results.py <results-dir>")
        sys.exit(1)
    results_dir = pathlib.Path(sys.argv[1])
    results = load_results(results_dir)
    out = results_dir / "plots"
    out.mkdir(exist_ok=True)
    write_csv(results, out / "summary.csv")
    plot_scenario(results, "capacity", "throughput_tps", "Throughput (tok/s)", "Capacity: Throughput vs Concurrency", out / "capacity_curve.png")
    plot_scenario(results, "capacity", "ttft_p99_ms", "TTFT p99 (ms)", "Capacity: TTFT p99 vs Concurrency", out / "latency_curve.png")
    plot_scenario(results, "sales", "ttft_p99_ms", "TTFT p99 (ms)", "Sales: Latency at Low Concurrency", out / "sales_comparison.png", logscale=False)
    print(f"\nDone. Outputs in {out}")


if __name__ == "__main__":
    main()