"""Extract targeted H3 block metrics from an Nsight Compute raw CSV export."""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path


ROLES = ("qkv", "sage2", "attention_output", "fc1", "fc2")
EXPECTED_GRIDS = (
    "(296, 168, 1)",
    "(296, 56, 1)",
    "(296, 42, 1)",
    "(296, 224, 1)",
    "(296, 42, 1)",
)


def number(row: dict[str, str], name: str) -> float | None:
    value = row.get(name, "")
    if value in {"", "no data", "n/a"}:
        return None
    return float(value)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--traffic", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    with args.input.open(newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        units = next(reader)
        rows = list(reader)
    if len(rows) != len(ROLES):
        raise RuntimeError(f"expected five targeted launches, found {len(rows)}")
    with args.traffic.open(newline="", encoding="utf-8") as handle:
        traffic_reader = csv.DictReader(handle)
        next(traffic_reader)
        traffic_rows = list(traffic_reader)
    if len(traffic_rows) != len(ROLES):
        raise RuntimeError(f"expected five traffic launches, found {len(traffic_rows)}")
    for index, (role, row, traffic_row, expected_grid) in enumerate(
        zip(ROLES, rows, traffic_rows, EXPECTED_GRIDS, strict=True)
    ):
        if int(row["ID"]) != index or int(traffic_row["ID"]) != index:
            raise RuntimeError(f"{role} launch ID/order contract failed")
        if row["Kernel Name"] != traffic_row["Kernel Name"]:
            raise RuntimeError(f"{role} kernel differs between metric and traffic passes")
        if row["Grid Size"] != expected_grid or traffic_row["Grid Size"] != expected_grid:
            raise RuntimeError(f"{role} grid/order contract failed")
        if role == "sage2" and "qk_int_sv_f8_attn_kernel" not in row["Kernel Name"]:
            raise RuntimeError("Sage2 launch contract failed")
        if role != "sage2" and "block_scaled" not in row["Kernel Name"]:
            raise RuntimeError(f"{role} NVFP4 GEMM launch contract failed")

    stall_prefix = "smsp__average_warps_issue_stalled_"
    stall_suffix = "_per_issue_active.ratio"
    output = {}
    for role, row, traffic_row in zip(ROLES, rows, traffic_rows, strict=True):
        stalls = []
        for name in row:
            if name.startswith(stall_prefix) and name.endswith(stall_suffix):
                value = number(row, name)
                if value is not None:
                    stalls.append({
                        "reason": name[len(stall_prefix):-len(stall_suffix)],
                        "warps_per_issue_active": value,
                    })
        stalls.sort(key=lambda item: item["warps_per_issue_active"], reverse=True)
        output[role] = {
            "launch_id": int(row["ID"]),
            "kernel_name": row["Kernel Name"],
            "grid_size": row["Grid Size"],
            "block_size": row["Block Size"],
            "duration_ns": number(row, "gpu__time_duration.sum"),
            "registers_per_thread": number(row, "launch__registers_per_thread"),
            "achieved_occupancy_percent": number(
                row, "sm__warps_active.avg.pct_of_peak_sustained_active"
            ),
            "eligible_warps_per_scheduler": number(
                row, "smsp__warps_eligible.avg.per_cycle_active"
            ),
            "issue_active_percent": number(
                row, "smsp__issue_active.avg.pct_of_peak_sustained_active"
            ),
            "sm_throughput_percent": number(
                row, "sm__throughput.avg.pct_of_peak_sustained_elapsed"
            ),
            "tensor_pipe_active_percent": number(
                row, "sm__pipe_tensor_cycles_active.avg.pct_of_peak_sustained_elapsed"
            ),
            "l2_requested_bytes": number(traffic_row, "lts__t_bytes.sum"),
            "l2_hit_rate_percent": number(traffic_row, "lts__t_sector_hit_rate.pct"),
            "l2_throughput_percent": number(
                row, "lts__throughput.avg.pct_of_peak_sustained_elapsed"
            ),
            "memory_throughput_percent": number(
                traffic_row, "gpu__compute_memory_throughput.avg.pct_of_peak_sustained_elapsed"
            ),
            "local_spilling_requests": number(row, "derived__local_spilling_requests"),
            "top_scheduler_stalls": stalls[:5],
        }

    report = {
        "source_csv": str(args.input),
        "source_traffic_csv": str(args.traffic),
        "launch_order_contract": list(ROLES),
        "cache_control": "none (warmed/uncontrolled cache, as reported by NCU)",
        "metrics": output,
        "units": {
            "duration_ns": "ns",
            "l2_requested_bytes": "lts__t_bytes.sum",
            "throughput_and_hit_rate": "%",
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()