#!/usr/bin/env python3
"""Aggregate H3 profile outputs and produce machine-readable bottleneck hints."""

from __future__ import annotations

import argparse
import csv
import json
import math
import re
import statistics
from collections import defaultdict
from pathlib import Path
from typing import Any


def pct(xs: list[float], q: float) -> float:
    if not xs:
        return 0.0
    xs = sorted(xs)
    pos = (len(xs) - 1) * q
    lo, hi = math.floor(pos), math.ceil(pos)
    return xs[lo] if lo == hi else xs[lo] * (hi - pos) + xs[hi] * (pos - lo)


def load_rows(root: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for path in root.rglob("results.jsonl"):
        for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
            try:
                row = json.loads(line)
                row["result_file"] = str(path)
                # Older/light clients may carry the resident-layout label even
                # when the phase activated all four instances. The phase path
                # is authoritative for this run topology.
                if "light_RVA_4active" in str(path) and row.get("deployment") == "LIGHT_tp2x4_one_active":
                    row["deployment"] = "LIGHT_tp2x4_four_active"
                rows.append(row)
            except json.JSONDecodeError:
                pass
    return rows


def parse_stage_logs(root: Path) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    timing_re = re.compile(r"(?:stage|name)[=: ]+([A-Za-z0-9_./-]+).*?(?:time|duration)[_ ]?(?:ms|s)?[=: ]+([0-9.]+)")
    for path in root.rglob("server.log"):
        for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
            if "stage" not in line.lower() and "timing" not in line.lower():
                continue
            match = timing_re.search(line)
            if match:
                records.append({"server_log": str(path), "stage": match.group(1), "value": float(match.group(2)), "raw": line[-2000:]})
    return records


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("root", type=Path)
    args = parser.parse_args()
    rows = load_rows(args.root)
    successful = [r for r in rows if r.get("success")]
    grouped: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
    for row in successful:
        grouped[(row["deployment"], row["scenario"])].append(row)

    summaries = []
    for (deployment, scenario), items in sorted(grouped.items()):
        lat = [float(x["latency_s"]) for x in items]
        inf = [float(x["inference_time_s"]) for x in items if x.get("inference_time_s") is not None]
        wall = max(float(x["finished_at_epoch"]) for x in items) - min(float(x["started_at_epoch"]) for x in items)
        summaries.append({
            "deployment": deployment,
            "scenario": scenario,
            "completed": len(items),
            "latency_mean_s": statistics.fmean(lat),
            "latency_p50_s": pct(lat, 0.5),
            "latency_p95_s": pct(lat, 0.95),
            "inference_mean_s": statistics.fmean(inf) if inf else None,
            "serving_gap_mean_s": statistics.fmean(lat) - statistics.fmean(inf) if inf else None,
            "machine_wall_s": wall,
            "machine_qps": len(items) / wall if wall > 0 else None,
        })

    args.root.mkdir(parents=True, exist_ok=True)
    (args.root / "summary.json").write_text(json.dumps({
        "requests_recorded": len(rows),
        "completed": len(successful),
        "failed": len(rows) - len(successful),
        "groups": summaries,
    }, ensure_ascii=False, indent=2), encoding="utf-8")
    fields = list(summaries[0]) if summaries else ["deployment", "scenario"]
    with (args.root / "summary.csv").open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields)
        writer.writeheader()
        writer.writerows(summaries)

    stage_records = parse_stage_logs(args.root)
    with (args.root / "stage_log_extract.jsonl").open("w", encoding="utf-8") as f:
        for record in stage_records:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    by_key = {(x["deployment"], x["scenario"]): x for x in summaries}
    comparisons = []
    for scenario in sorted({x["scenario"] for x in summaries}):
        base = by_key.get(("P1_tp2_single", scenario))
        for deployment in ("P2_tp2x4_one_active", "P3_tp2x4_fl_active", "P4_tp2x4_ref_active", "P5_tp2x4_mixed"):
            other = by_key.get((deployment, scenario))
            if base and other:
                ratio = other["latency_mean_s"] / base["latency_mean_s"]
                comparisons.append({
                    "scenario": scenario,
                    "baseline": base["deployment"],
                    "deployment": deployment,
                    "latency_ratio": ratio,
                    "contention_pct": (ratio - 1.0) * 100.0,
                    "classification": "cross_instance_primary" if ratio >= 1.20 else "cross_instance_secondary" if ratio >= 1.10 else "not_significant",
                })
    (args.root / "bottleneck_comparisons.json").write_text(json.dumps(comparisons, ensure_ascii=False, indent=2), encoding="utf-8")

    torch_traces = [str(p) for p in args.root.rglob("*.pt.trace.json")]
    nsys_reports = [str(p) for p in args.root.rglob("*.nsys-rep")]
    inventory = {
        "stage_records": len(stage_records),
        "torch_trace_count": len(torch_traces),
        "torch_traces": torch_traces,
        "nsys_report_count": len(nsys_reports),
        "nsys_reports": nsys_reports,
        "note": "Final primary/secondary bottleneck classification requires stage shares plus Torch/Nsight kernel and uncovered-NCCL review.",
    }
    (args.root / "profile_inventory.json").write_text(json.dumps(inventory, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"wrote {args.root / 'summary.csv'}; completed={len(successful)} failed={len(rows)-len(successful)}")
    return int(any(not r.get("success") for r in rows))


if __name__ == "__main__":
    raise SystemExit(main())