#!/usr/bin/env python3
"""Summarize the longest graph-on execute window in each TP rank."""

from __future__ import annotations

import argparse
import gzip
import json
from collections import Counter, defaultdict
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--trace-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def load_events(path: Path) -> list[dict]:
    opener = gzip.open if path.suffix == ".gz" else open
    with opener(path, "rt") as source:
        return json.load(source)["traceEvents"]


def classify(name: str) -> str:
    lower = name.lower()
    if any(token in lower for token in ("nccl", "allreduce", "all_reduce")):
        return "collective"
    if "topkgating" in lower or "fused_topk" in lower:
        return "router"
    if any(
        token in lower
        for token in (
            "fused_moe",
            "moefcgemm",
            "tensorrt_llm::kernels::cutlass_kernels",
            "groupproblemshape",
            "memcpy32_post",
        )
    ):
        return "moe"
    if any(
        token in lower
        for token in (
            "flashattn",
            "flashattnfwd",
            "reshape_and_cache",
            "prepare_varlen_num_blocks",
        )
    ):
        return "attention"
    if any(
        token in lower
        for token in ("nvjet", "cublaslt", "rms_norm", "rsqrt", "rope")
    ):
        return "linear_norm_rope"
    return "other"


def analyze_rank(path: Path) -> dict:
    events = load_events(path)
    kernels = [event for event in events if event.get("cat") == "kernel"]
    windows = [
        event
        for event in events
        if event.get("cat") == "gpu_user_annotation"
        and str(event.get("name", "")).startswith("execute_")
    ]
    if not windows:
        raise ValueError(f"{path}: no execute annotation")
    selected = max(windows, key=lambda event: float(event["dur"]))
    start = float(selected["ts"])
    end = start + float(selected["dur"])
    current = [
        kernel for kernel in kernels if start <= float(kernel["ts"]) < end
    ]
    components: dict[str, float] = defaultdict(float)
    kernel_totals: dict[str, float] = defaultdict(float)
    for kernel in current:
        duration_ms = float(kernel["dur"]) / 1000
        name = str(kernel["name"])
        components[classify(name)] += duration_ms
        kernel_totals[name] += duration_ms
    kernel_rows = [
        {"name": name, "duration_ms": duration}
        for name, duration in sorted(
            kernel_totals.items(), key=lambda item: -item[1]
        )
    ]
    wall_ms = float(selected["dur"]) / 1000
    busy_ms = sum(components.values())
    return {
        "trace": str(path),
        "selected_execute_annotation": str(selected["name"]),
        "execute_annotation_histogram": dict(
            sorted(Counter(str(window["name"]) for window in windows).items())
        ),
        "all_execute_windows": [
            {
                "name": str(window["name"]),
                "duration_ms": float(window["dur"]) / 1000,
            }
            for window in sorted(windows, key=lambda event: float(event["ts"]))
        ],
        "execute_wall_ms": wall_ms,
        "gpu_kernel_busy_ms": busy_ms,
        "non_kernel_gap_ms": wall_ms - busy_ms,
        "components_ms": dict(sorted(components.items())),
        "kernel_rows": kernel_rows,
    }


def main() -> None:
    args = parse_args()
    traces = sorted(args.trace_root.rglob("*.pt.trace.json*"))
    if not traces:
        raise ValueError(f"no traces below {args.trace_root}")
    ranks = [analyze_rank(path) for path in traces]
    critical = max(ranks, key=lambda rank: rank["execute_wall_ms"])
    payload = {
        "schema": "frontier-tp2-prefill-serving-smoke.v1",
        "contract": {
            "selection": "longest execute annotation per TP rank",
            "critical_path": "rank with largest selected execute wall",
            "component_time": "sum of CUDA kernel duration within selected window",
        },
        "ranks": ranks,
        "critical_rank": critical,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(
        json.dumps(
            {
                "ranks": len(ranks),
                "execute_wall_ms": critical["execute_wall_ms"],
                "components_ms": critical["components_ms"],
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()