#!/usr/bin/env bash
# ============================================================
# collect_gpu_stats.sh — Continuous GPU monitoring during benchmark
#
# Usage:
#   bash scripts/collect_gpu_stats.sh [output_file] [interval_seconds]
#
# Logs nvidia-smi metrics every N seconds to a CSV file.
# Run in background during benchmark:
#   bash scripts/collect_gpu_stats.sh /tmp/gpu_stats.csv 1 &
#   GPU_PID=$!
#   ... run benchmark ...
#   kill $GPU_PID
#
# Default output: reports/gpu_stats_<timestamp>.csv
# ============================================================
set -euo pipefail

OUTPUT="${1:-}"
INTERVAL="${2:-1}"

# Default output path
if [[ -z "$OUTPUT" ]]; then
    SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
    PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
    TIMESTAMP=$(date +%Y%m%d_%H%M%S)
    OUTPUT="${PROJECT_DIR}/reports/gpu_stats_${TIMESTAMP}.csv"
fi

# Ensure output directory exists
mkdir -p "$(dirname "$OUTPUT")"

echo "============================================"
echo " GPU Stats Collector"
echo "============================================"
echo "  Output:   $OUTPUT"
echo "  Interval: ${INTERVAL}s"
echo "  PID:      $$"
echo "============================================"
echo ""
echo "Recording... Press Ctrl+C to stop."
echo ""

# Write CSV header
echo "timestamp,elapsed_s,gpu_index,gpu_name,pci_bus,driver_version,"\
"pstate,temperature_gpu,utilization_gpu,utilization_memory,"\
"memory_total_mib,memory_used_mib,memory_free_mib,"\
"power_draw_w,power_limit_w,"\
"clocks_sm_mhz,clocks_memory_mhz,"\
"fan_speed_pct,performance_state" > "$OUTPUT"

START_TIME=$(date +%s)

trap 'echo ""; echo "Stopped. Stats saved to: $OUTPUT"; exit 0' INT TERM

while true; do
    NOW=$(date +%s)
    ELAPSED=$((NOW - START_TIME))
    TS=$(date -u +%Y-%m-%dT%H:%M:%SZ)

    # Query nvidia-smi in CSV format (no header)
    nvidia-smi --query-gpu=index,name,pci.bus_id,driver_version,\
pstate,temperature.gpu,utilization.gpu,utilization.memory,\
memory.total,memory.used,memory.free,\
power.draw,power.limit,\
clocks.sm,clocks.memory,\
fan.speed,pstate \
        --format=csv,noheader,nounits 2>/dev/null | while IFS=',' read -r line; do
        echo "${TS},${ELAPSED},${line}" | tr -d ' ' >> "$OUTPUT"
    done

    sleep "$INTERVAL"
done