#!/usr/bin/env python3
"""Analyze nsys traces for DSpark profiling.

Usage:
    python scripts/analyze_dspark_nsys.py <nsys-rep-file>
"""

import json
import subprocess
import sys
from pathlib import Path


def run_nsys_report(rep_file: Path, report_name: str) -> str:
    cmd = [
        "nsys", "stats",
        "--report", report_name,
        "--format", "json",
        str(rep_file),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"nsys {report_name} failed: {result.stderr}")
        return ""
    return result.stdout


def summarize_cuda_sum(rep_file: Path) -> list[dict]:
    """Return top CUDA kernels by total time."""
    text = run_nsys_report(rep_file, "cuda_kernel_sum")
    if not text:
        return []
    # nsys stats --format json outputs JSON after some header lines.
    # Find the first '[' character.
    start = text.find("[")
    if start < 0:
        return []
    data = json.loads(text[start:])
    # Sort by total time descending.
    rows = data[1:] if len(data) > 1 else data
    rows_sorted = sorted(rows, key=lambda x: x.get("Total Time (ns)", 0), reverse=True)
    return rows_sorted[:50]


def summarize_nvtx_sum(rep_file: Path) -> list[dict]:
    text = run_nsys_report(rep_file, "nvtx_sum")
    start = text.find("[")
    if start < 0:
        return []
    data = json.loads(text[start:])
    rows = data[1:] if len(data) > 1 else data
    rows_sorted = sorted(rows, key=lambda x: x.get("Total Time (ns)", 0), reverse=True)
    return rows_sorted[:30]


def main():
    if len(sys.argv) < 2:
        print("Usage: analyze_dspark_nsys.py <nsys-rep-file>")
        sys.exit(1)
    rep_file = Path(sys.argv[1])
    if not rep_file.exists():
        print(f"File not found: {rep_file}")
        sys.exit(1)

    print(f"Analyzing {rep_file}...\n")

    cuda_kernels = summarize_cuda_sum(rep_file)
    print("=== Top CUDA Kernels by Total Time ===")
    for row in cuda_kernels[:20]:
        name = row.get("Name", "N/A")
        total_ns = row.get("Total Time (ns)", 0)
        count = row.get("Instances", 0)
        avg_ns = row.get("Avg (ns)", 0)
        print(f"  {total_ns/1e6:8.2f} ms  {count:6d} calls  avg={avg_ns/1e6:6.3f} ms  {name[:80]}")

    nvtx = summarize_nvtx_sum(rep_file)
    if nvtx:
        print("\n=== Top NVTX Ranges ===")
        for row in nvtx[:20]:
            name = row.get("Range", "N/A")
            total_ns = row.get("Total Time (ns)", 0)
            count = row.get("Instances", 0)
            print(f"  {total_ns/1e6:8.2f} ms  {count:6d} calls  {name[:80]}")


if __name__ == "__main__":
    main()