#!/bin/bash
#
# check_nvidia_smi_temperature - Nagios plugin to check NVIDIA GPU temperatures via nvidia-smi
#
# Copyright (C) 2026 Thomas Wagner <thomas.wagner@jku.at>
# SPDX-License-Identifier: GPL-2.0-or-later

PROGNAME=$(basename "$0")

STATE_OK=0
STATE_WARNING=1
STATE_CRITICAL=2
STATE_UNKNOWN=3

WARNING=80
CRITICAL=90
INDICES=""

print_usage() {
    echo "Usage: $PROGNAME [-w <warning>] [-c <critical>] [-i <index[,index...]>]"
    echo ""
    echo "  -w <warning>   Warning temperature threshold in Celsius (default: $WARNING)"
    echo "  -c <critical>  Critical temperature threshold in Celsius (default: $CRITICAL)"
    echo "  -i <indices>   Comma-separated list of GPU indices to check (default: all cards)"
    echo "  -h             Show this help message"
}

while getopts "w:c:i:h" opt; do
    case $opt in
        w) WARNING=$OPTARG ;;
        c) CRITICAL=$OPTARG ;;
        i) INDICES=$OPTARG ;;
        h) print_usage; exit $STATE_OK ;;
        *) print_usage; exit $STATE_UNKNOWN ;;
    esac
done

if ! [[ "$WARNING" =~ ^[0-9]+$ ]] || ! [[ "$CRITICAL" =~ ^[0-9]+$ ]]; then
    echo "UNKNOWN: Warning and critical thresholds must be non-negative integers"
    exit $STATE_UNKNOWN
fi

if [ -n "$INDICES" ] && ! [[ "$INDICES" =~ ^[0-9]+(,[0-9]+)*$ ]]; then
    echo "UNKNOWN: GPU indices must be a comma-separated list of non-negative integers"
    exit $STATE_UNKNOWN
fi

if [ "$WARNING" -ge "$CRITICAL" ]; then
    echo "UNKNOWN: Warning threshold must be less than critical threshold"
    exit $STATE_UNKNOWN
fi

if ! command -v nvidia-smi >/dev/null 2>&1; then
    echo "UNKNOWN: nvidia-smi command not found"
    exit $STATE_UNKNOWN
fi

NVIDIA_SMI_ARGS=(--query-gpu=index,name,temperature.gpu --format=csv,noheader,nounits)
if [ -n "$INDICES" ]; then
    NVIDIA_SMI_ARGS=(-i "$INDICES" "${NVIDIA_SMI_ARGS[@]}")
fi

OUTPUT=$(nvidia-smi "${NVIDIA_SMI_ARGS[@]}" 2>/dev/null)
if [ $? -ne 0 ] || [ -z "$OUTPUT" ]; then
    echo "UNKNOWN: Unable to query GPU temperatures via nvidia-smi"
    exit $STATE_UNKNOWN
fi

STATE=$STATE_OK
STATUS_MSG=""
PERFDATA=""

while IFS=',' read -r index name temp; do
    index=$(echo "$index" | xargs)
    name=$(echo "$name" | xargs)
    temp=$(echo "$temp" | xargs)

    if ! [[ "$temp" =~ ^[0-9]+$ ]]; then
        continue
    fi

    GPU_STATE=$STATE_OK
    if [ "$temp" -ge "$CRITICAL" ]; then
        GPU_STATE=$STATE_CRITICAL
    elif [ "$temp" -ge "$WARNING" ]; then
        GPU_STATE=$STATE_WARNING
    fi

    if [ "$GPU_STATE" -gt "$STATE" ]; then
        STATE=$GPU_STATE
    fi

    STATUS_MSG="${STATUS_MSG}GPU${index} (${name}): ${temp}C, "
    PERFDATA="${PERFDATA}gpu${index}_temp=${temp};${WARNING};${CRITICAL};0; "
done <<< "$OUTPUT"

STATUS_MSG="${STATUS_MSG%, }"
PERFDATA="${PERFDATA% }"

if [ -z "$STATUS_MSG" ]; then
    echo "UNKNOWN: No valid GPU temperature data returned by nvidia-smi"
    exit $STATE_UNKNOWN
fi

case $STATE in
    $STATE_OK) LABEL="OK" ;;
    $STATE_WARNING) LABEL="WARNING" ;;
    $STATE_CRITICAL) LABEL="CRITICAL" ;;
esac

echo "${LABEL}: ${STATUS_MSG} | ${PERFDATA}"
exit $STATE
