#!/bin/bash

# ganglia_name | nvidia-smi query | type | units | aggregation (sum|avg|none) | title
metrics=(
    "gpu_driver_version|driver_version|string||none|GPU driver version"
    "gpu_name|gpu_name|string||none|GPU name"
    "gpu_vbios_version|vbios_version|string||none|GPU VBIOS version"
    "gpu_pcie_link_gen_current|pcie.link.gen.gpucurrent|uint8||none|GPU current PCIe link generation"
    "gpu_pcie_link_gen_max|pcie.link.gen.max|uint8||none|GPU maximum PCIe link generation"
    "gpu_pcie_link_width_current|pcie.link.width.current|uint8||none|GPU current PCIe link width"
    "gpu_pcie_link_width_max|pcie.link.width.max|uint8||none|GPU maximum PCIe link width"
    "gpu_memory_total|memory.total|uint32|MB|sum|GPU total memory"
    "gpu_memory_reserved|memory.reserved|uint32|MB|sum|GPU reserved memory"
    "gpu_memory_used|memory.used|uint32|MB|sum|GPU used memory"
    "gpu_memory_free|memory.free|uint32|MB|sum|GPU free memory"
    "gpu_fan_speed|fan.speed|float|%|avg|GPU fan speed"
    "gpu_utilization|utilization.gpu|float|%|avg|GPU utilization"
    "gpu_memory_utilization|utilization.memory|float|%|avg|GPU memory utilization"
    "gpu_temperature|temperature.gpu|float|C|avg|GPU temperature"
    "gpu_power_draw|power.draw|float|W|sum|GPU power draw"
    "gpu_power_limit|power.limit|float|W|sum|GPU power limit"
    "gpu_clocks_current_graphics|clocks.current.graphics|float|MHz|avg|GPU graphics clock"
    "gpu_clocks_current_sm|clocks.current.sm|float|MHz|avg|GPU streaming multiprocessor clock"
    "gpu_clocks_current_memory|clocks.current.memory|float|MHz|avg|GPU memory clock"
    "gpu_clocks_current_video|clocks.current.video|float|MHz|avg|GPU video clock"
)

# Build comma-separated query
smi_query=""
for m in "${metrics[@]}"; do
    IFS='|' read -r _ query _ _ _ _ <<< "$m"
    smi_query+="$query,"
done
smi_query=${smi_query%,}

raw=$(nvidia-smi --format=csv,noheader,nounits --query-gpu="$smi_query" 2>/dev/null)

# Get GPU count (nvidia-smi prints it once per GPU, so take first line)
gpu_count=$(nvidia-smi --format=csv,noheader,nounits --query-gpu=count 2>/dev/null | head -1)

# Accumulators
declare -a agg_sum
declare -a agg_count
for i in "${!metrics[@]}"; do
    agg_sum[$i]=0
    agg_count[$i]=0
done

metric() {
    local name="$1" value="$2" type="$3" units="$4" title="$5"
    [[ -z "$value" ]] && return
    local args=(--name="$name" --title="$title" --value="$value" --type="$type" --group=gpu)
    [[ -n "$units" ]] && args+=(--units="$units")
    gmetric "${args[@]}"
}

gpu_idx=0
while IFS= read -r line; do
    [[ -z "$line" ]] && continue
    IFS=',' read -ra values <<< "$line"

    mi=0
    for m in "${metrics[@]}"; do
        IFS='|' read -r ganglia_name _ type units agg title <<< "$m"
        value="${values[$mi]}"
        # nvidia-smi CSV puts a space after each comma; trim leading whitespace
        value="${value#"${value%%[![:space:]]*}"}"
        mi=$((mi + 1))
        [[ -z "$value" ]] && continue

        # Single GPU: emit non-indexed (preserves history)
        # Multiple GPUs: emit indexed per-GPU
        if [[ $gpu_count -eq 1 ]]; then
            metric "$ganglia_name" "$value" "$type" "$units" "$title"
        else
            base="${ganglia_name#gpu_}"
            # Replace 'GPU' with 'GPU N' in title for per-GPU metrics
            gpu_title="${title/GPU/GPU $gpu_idx}"
            metric "gpu${gpu_idx}_${base}" "$value" "$type" "$units" "$gpu_title"
        fi

        # Accumulate for aggregate
        if [[ "$agg" == "sum" || "$agg" == "avg" ]]; then
            idx=$((mi - 1))
            agg_sum[$idx]=$(awk "BEGIN {print ${agg_sum[$idx]} + $value}")
            agg_count[$idx]=$((agg_count[$idx] + 1))
        fi
    done

    gpu_idx=$((gpu_idx + 1))
done <<< "$raw"

# Emit aggregates (only when multiple GPUs; single GPU already emitted above)
mi=0
for m in "${metrics[@]}"; do
    IFS='|' read -r ganglia_name _ type units agg title <<< "$m"
    if [[ "$agg" == "none" || $gpu_count -eq 1 ]]; then
        mi=$((mi + 1))
        continue
    fi

    if [[ "$agg" == "sum" ]]; then
        metric "$ganglia_name" "${agg_sum[$mi]}" "$type" "$units" "Total $title"
    elif [[ "$agg" == "avg" && ${agg_count[$mi]} -gt 0 ]]; then
        avg=$(awk "BEGIN {printf \"%.2f\", ${agg_sum[$mi]} / ${agg_count[$mi]}}")
        metric "$ganglia_name" "$avg" "$type" "$units" "Average $title"
    fi
    mi=$((mi + 1))
done