#!/usr/bin/env bash
set -Eeuo pipefail

ROOT=${ROOT:-/data/wxy}
RUN_ID=${RUN_ID:-h3-sdpa-kernel-matrix-$(date +%Y%m%d-%H%M%S)}
OUT=${OUT:-$ROOT/profile_results/$RUN_ID}
PYTHON=${PYTHON:-/root/.miniconda3/envs/sglang/bin/python}
BENCH=${BENCH:-$ROOT/h3_profile/h3_sdpa_kernel_bench.py}

mkdir -p "$OUT"/{jobs,system,metadata}
cp "$0" "$BENCH" "$OUT/metadata/"
date -Ins > "$OUT/metadata/date.txt"
nvidia-smi -q > "$OUT/metadata/nvidia-smi-q.txt"
nvidia-smi topo -m > "$OUT/metadata/topology.txt"
"$PYTHON" -m pip freeze > "$OUT/metadata/pip-freeze.txt"

declare -a PIDS=()
cleanup(){
  local rc=$? p
  trap - EXIT INT TERM
  for p in "${PIDS[@]:-}"; do kill -TERM "$p" 2>/dev/null || true; done
  exit "$rc"
}
trap cleanup EXIT INT TERM

nvidia-smi dmon -s pucvmet -d 1 -o DT > "$OUT/system/nvidia-dmon.log" 2>&1 &
DMON_PID=$!

backends=(flash cudnn efficient auto)
for gpu in 0 1 2 3; do
  backend=${backends[$gpu]}
  CUDA_VISIBLE_DEVICES=$gpu PYTHONUNBUFFERED=1 "$PYTHON" "$BENCH" \
    --scenario F3 --backend "$backend" --sequence 41778 \
    --output "$OUT/jobs/gpu${gpu}-F3-${backend}.json" \
    > "$OUT/jobs/gpu${gpu}-F3-${backend}.log" 2>&1 &
  PIDS+=("$!")
done
for gpu in 4 5 6 7; do
  backend=${backends[$((gpu-4))]}
  CUDA_VISIBLE_DEVICES=$gpu PYTHONUNBUFFERED=1 "$PYTHON" "$BENCH" \
    --scenario RVA --backend "$backend" --sequence 81547 \
    --output "$OUT/jobs/gpu${gpu}-RVA-${backend}.json" \
    > "$OUT/jobs/gpu${gpu}-RVA-${backend}.log" 2>&1 &
  PIDS+=("$!")
done

failed=0
for p in "${PIDS[@]}"; do wait "$p" || failed=1; done
PIDS=()
kill -TERM "$DMON_PID" 2>/dev/null || true
wait "$DMON_PID" 2>/dev/null || true

"$PYTHON" - "$OUT" <<'PY'
import json, pathlib, sys
root = pathlib.Path(sys.argv[1])
rows = []
for path in sorted((root / "jobs").glob("*.json")):
    payload = json.loads(path.read_text())
    for result in payload.get("results", []):
        rows.append({
            "scenario": payload["scenario"],
            "backend": payload["backend"],
            "layout": result["layout"],
            "median_ms": result["warm"]["median_ms"],
            "cold_median_ms": result["cold"]["median_ms"],
            "cold_over_warm": result["cold_over_warm"],
            "effective_tflops": result["effective_tflops"],
            "cuda_graph_ms": result["cuda_graph"].get("median_ms"),
            "q_stride": result["q_stride"],
        })
(root / "summary.json").write_text(json.dumps(rows, indent=2) + "\n")
print(json.dumps(rows, indent=2))
PY

if ((failed)); then
  touch "$OUT/FAILED"
  exit 1
fi
touch "$OUT/DONE"