#!/bin/bash
# GPU Fan Control — Delta QFR1212GHE (Tesla V100 / motherboard fallback)
# Version: 3.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Description:
#   Controls the Delta QFR1212GHE cooling fan based on GPU temperatures.
#
#   Mode selection logic:
#     1 NVIDIA GPU detected  → monitor that single GPU, use its temp.
#     2 NVIDIA GPUs detected → monitor both GPUs, use max(temp1, temp2).
#     0 NVIDIA GPUs detected → fall back to motherboard thermal zones,
#                               use max(temp) across all zones.
#
#   When the fan is controlled, its mode applies to ALL GPUs simultaneously
#   (one physical fan blows over both cards).
#
#   Thresholds (raw, with 2°C hysteresis dead-band):
#     Mode 6 - FULL speed    → TEMP >= 45°C or rise >= 3°C/2s → 120s burst
#     Mode 5 - HIGH speed    → 33-39°C  (hysteresis: 35-2 to 40-2)
#     Mode 4 - MED speed     → 28-32°C  (hysteresis: 30-2 to 35-2)
#     Mode 3 - LOW speed     → 23-27°C  (hysteresis: 25-2 to 30-2)
#     Mode 2 - STANDBY       → < 23°C   (hysteresis: 25-2)
#     Mode 1 - STOP          → not used
#
#   Hysteresis: 2°C dead-band prevents rapid mode switching.
#   Burst: if any GPU temperature rises by >=3°C within 2 seconds → FULL for 2 min.

LOG=/var/log/gpu_fan.log
DEV=/dev/ttyUSB0

# Timestamp helper
ts() { date '+%Y-%m-%dT%H:%M:%S'; }

# Open the serial device for writing to the fan controller.
exec 3>"$DEV"
stty -F "$DEV" 115200

# Save original stderr (→ systemd journal) as fd 4
exec 4>&2

# stdout + stderr → log file
exec 1>>"$LOG" 2>&1

# Thresholds
TH_COLD=25
TH_LOW=30
TH_MED=35
TH_HIGH=40
TH_BURST=45
TH_BURST_RISE=3
BURST_WINDOW=2
BURST_DURATION=120

# Hysteresis (2°C dead-band)
HYST=2

# Initialize variables
prevTemp=0
prevTime=0
burstTimer=0
lastMode=0
noDataCount=0

# ---- GPU detection -------------------------------------------------------

# Collect GPU indices into an array (temps re-read each loop iteration).
declare -a gpu_idx=()

if command -v nvidia-smi &>/dev/null; then
    while IFS=',' read -r gpu_idx_val gpu_temp_val; do
        # Trim whitespace from temp
        gpu_temp_val="${gpu_temp_val// /}"
        [[ -z "$gpu_temp_val" ]] && continue
        gpu_idx+=("${gpu_idx_val// /}")
    done < <(nvidia-smi --query-gpu=index,temperature.gpu --format=csv,noheader 2>&1 | grep -E '^[0-9]+,')
fi

gpuCount=${#gpu_idx[@]}

if [ "$gpuCount" -eq 0 ]; then
    echo "[$(ts)] [INFO] No NVIDIA GPUs detected. Falling back to motherboard thermal zones."
fi

# ---- Motherboard thermal zones (fallback) --------------------------------

# Store motherboard zone paths and types; temps re-read each iteration.
declare -a board_paths=()
declare -a board_types=()

for zone in /sys/class/thermal/thermal_zone*; do
    [ -f "$zone/temp" ] || continue
    zoneType=$(cat "$zone/type" 2>/dev/null)
    zoneTemp=$(cat "$zone/temp" 2>/dev/null)

    # Convert from millidegrees if needed
    if [[ "$zoneTemp" =~ ^[0-9]+$ ]] && [ "$zoneTemp" -gt 1000 ]; then
        zoneTemp=$((zoneTemp / 1000))
    fi

    # Sanity check — skip zones with empty type or invalid temp
    if [[ -n "$zoneType" ]] && [[ "$zoneTemp" =~ ^[0-9]+$ ]] && [ "$zoneTemp" -ge 0 ] && [ "$zoneTemp" -le 255 ]; then
        board_paths+=("$zone")
        board_types+=("$zoneType")
    fi
done

boardCount=${#board_paths[@]}

# ---- Helper: compute max temperature and build a display string ----------

# Read all GPU temperatures and find max.
# Sets: maxTemp, maxIdx, allTempsStr
read_max_gpu_temps() {
    maxTemp=0
    maxIdx=""
    allTempsStr=""
    local gpu_idx_val gpu_temp_val
    local sep=""
    while IFS=',' read -r gpu_idx_val gpu_temp_val; do
        # Trim whitespace from temp
        gpu_temp_val="${gpu_temp_val// /}"
        [[ -z "$gpu_temp_val" ]] && continue
        allTempsStr+="${sep}GPU-${gpu_idx_val// /}=$gpu_temp_val°C"
        sep=", "
        if [ "$gpu_temp_val" -gt "$maxTemp" ]; then
            maxTemp=$gpu_temp_val
            maxIdx="${gpu_idx_val// /}"
        fi
    done < <(nvidia-smi --query-gpu=index,temperature.gpu --format=csv,noheader 2>&1 | grep -E '^[0-9]+,')
}

# Read all motherboard thermal zone temperatures and find max.
# Sets: maxTemp, maxIdx, allTempsStr
read_max_board_temps() {
    maxTemp=0
    maxIdx=""
    allTempsStr=""
    local zoneTemp
    local sep=""
    for ((i=0; i<boardCount; i++)); do
        zoneTemp=$(cat "${board_paths[$i]}/temp" 2>/dev/null)
        # Convert from millidegrees if needed
        if [[ "$zoneTemp" =~ ^[0-9]+$ ]] && [ "$zoneTemp" -gt 1000 ]; then
            zoneTemp=$((zoneTemp / 1000))
        fi
        if [[ "$zoneTemp" =~ ^[0-9]+$ ]] && [ "$zoneTemp" -ge 0 ] && [ "$zoneTemp" -le 255 ]; then
            allTempsStr+="${sep}${board_types[$i]}=$zoneTemp°C"
            sep=", "
            if [ "$zoneTemp" -gt "$maxTemp" ]; then
                maxTemp=$zoneTemp
                maxIdx=${board_types[$i]}
            fi
        fi
    done
}

# ---- Initial read --------------------------------------------------------

if [ "$gpuCount" -gt 0 ]; then
    read_max_gpu_temps
    TEMP=$maxTemp
    SOURCE="gpu-$maxIdx"
else
    read_max_board_temps
    if [ -n "$maxIdx" ]; then
        TEMP=$maxTemp
        SOURCE="board-$maxIdx"
    else
        # No sensors at all — full speed to be safe
        TEMP=50
        SOURCE="none (full speed)"
    fi
fi

# Initialize tracking variables with current values
prevTemp=$TEMP
prevTime=$(date +%s)
burstTimer=0
lastMode=0

# Print detected GPUs (if any)
if [ "$gpuCount" -gt 0 ]; then
    echo "[$(ts)] [INFO] Detected $gpuCount NVIDIA GPU(s):"
    for ((i=0; i<gpuCount; i++)); do
        echo "[$(ts)]   GPU $i: index=${gpu_idx[$i]}"
    done
fi

# Print all motherboard zones
if [ "$boardCount" -gt 0 ]; then
    echo "[$(ts)] [INFO] Detected $boardCount motherboard thermal zone(s):"
    for ((i=0; i<boardCount; i++)); do
        echo "[$(ts)]   Zone $i: type=${board_types[$i]}"
    done
fi

echo "[$(ts)] [INFO] Starting fan control. Monitoring $SOURCE. Temperature: ${TEMP}°C"

# ---- Main loop -----------------------------------------------------------

while true; do
    now=$(date +%s)

    # Re-read current temperatures
    if [ "$gpuCount" -gt 0 ]; then
        read_max_gpu_temps
        SOURCE="gpu-$maxIdx"
    else
        read_max_board_temps
        if [ -n "$maxIdx" ]; then
            SOURCE="board-$maxIdx"
        fi
    fi
    TEMP=$maxTemp

    echo "[$(ts)] [STATUS] All temps: ${allTempsStr:-none}, max=$TEMP°C from $SOURCE" >&4

    # Check if the temperature source is still alive
    if [ "$TEMP" -le 0 ] 2>/dev/null; then
        noDataCount=$((noDataCount + 1))
        if [ $noDataCount -ge 10 ]; then
            echo "[$(ts)] [ERROR] No temperature data for 10+ seconds. Running fan at full speed."
            echo "6" >&3
            exit 1
        fi
        sleep 1; continue
    fi

    noDataCount=0

    # --- 1. Если burst-режим активен — держим режим 6 ---
    if [ $burstTimer -gt 0 ]; then
        burstTimer=$((burstTimer - 1))
        if [ $lastMode -ne 6 ]; then
            echo "[$(ts)] [MODE 6] Full speed (burst, ${burstTimer}s remaining)"
            lastMode=6
            echo "6" >&3
        fi
        prevTemp=$TEMP; prevTime=$now
        sleep 1; continue
    fi

    # --- 2. Если temp >= TH_BURST (45) и burst не активен — новый burst ---
    if [ "$TEMP" -ge $TH_BURST ]; then
        burstTimer=$BURST_DURATION
        echo "[$(ts)] [MODE 6] Full speed (re-burst: $TEMP°C)"
        lastMode=6; prevTemp=$TEMP; prevTime=$now
        echo "6" >&3
        sleep 1; continue
    fi

    # --- 3. Burst detection: быстрый рост температуры ---
    if [ $((now - prevTime)) -le $BURST_WINDOW ] && [ $prevTime -gt 0 ]; then
        rise=$((TEMP - prevTemp))
        if [ $rise -ge $TH_BURST_RISE ]; then
            burstTimer=$BURST_DURATION
            echo "[$(ts)] [BURST] $TEMP°C (up from $prevTemp°C in $((now - prevTime))s, window=$BURST_WINDOW)"
            lastMode=6
            prevTemp=$TEMP; prevTime=$now
            echo "6" >&3
            sleep 1; continue
        fi
    fi

    # --- 4. Обычная логика с гистерезисом ---
    newMode=5  # default: HIGH (35-39)
    if [ $TEMP -lt $((TH_COLD - HYST)) ]; then
        newMode=2  # STANDBY (< 23)
    elif [ $TEMP -lt $((TH_LOW - HYST)) ]; then
        newMode=3  # LOW (< 28)
    elif [ $TEMP -lt $((TH_MED - HYST)) ]; then
        newMode=4  # MED (< 33)
    fi

    if [ $newMode -ne $lastMode ]; then
        echo "[$(ts)] [MODE $newMode] Temperature: $TEMP°C, source: $SOURCE"
        lastMode=$newMode
        echo "$newMode" >&3
    fi

    prevTemp=$TEMP
    prevTime=$now
    sleep 1
done