#!/usr/bin/env bash

set -e

# 默认风扇转速，单位百分比
FAN_SPEED="${1:-100}"

# X 显示号
DISPLAY_ID="${DISPLAY:-:0}"

export DISPLAY="$DISPLAY_ID"

echo "Using DISPLAY=$DISPLAY"
echo "Target fan speed: ${FAN_SPEED}%"

if ! command -v nvidia-settings >/dev/null 2>&1; then
    echo "Error: nvidia-settings not found."
    exit 1
fi

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

# 检查风扇转速范围
if ! [[ "$FAN_SPEED" =~ ^[0-9]+$ ]]; then
    echo "Error: fan speed must be an integer, e.g. 80 or 100."
    exit 1
fi

if [ "$FAN_SPEED" -lt 0 ] || [ "$FAN_SPEED" -gt 100 ]; then
    echo "Error: fan speed must be between 0 and 100."
    exit 1
fi

echo
echo "Detected GPUs:"
nvidia-smi --query-gpu=index,name,temperature.gpu --format=csv,noheader

echo
echo "Enabling manual fan control for all GPUs..."

GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)

for ((gpu=0; gpu<GPU_COUNT; gpu++)); do
    echo "Setting GPU ${gpu} fan control to manual..."
    nvidia-settings -a "[gpu:${gpu}]/GPUFanControlState=1" || true
done

echo
echo "Detecting fans from nvidia-settings..."

FANS=$(nvidia-settings -q fans 2>/dev/null | grep -o "fan:[0-9]\+" | sort -u | cut -d: -f2)

if [ -z "$FANS" ]; then
    echo "Error: no controllable fans detected."
    echo
    echo "Possible reasons:"
    echo "1. Coolbits is not enabled."
    echo "2. Xorg is not running."
    echo "3. DISPLAY is wrong."
    echo "4. Your GPU fan is controlled by BMC/IPMI, not NVIDIA driver."
    exit 1
fi

echo "Detected fans:"
echo "$FANS" | sed 's/^/fan:/'

echo
echo "Setting all detected fans to ${FAN_SPEED}%..."

for fan in $FANS; do
    echo "Setting fan:${fan} to ${FAN_SPEED}%"
    nvidia-settings -a "[fan:${fan}]/GPUTargetFanSpeed=${FAN_SPEED}" || true
done

echo
echo "Done."

echo
echo "Current fan status:"
nvidia-settings -q fans 2>/dev/null || true