#!/usr/bin/env bash

WARNING_THRESHOLD=15
CRITICAL_THRESHOLD=25
AVERAGE_SECONDS=5
SHOW_TOP_PROCESSES=false
PROXMOX_CALCULATION=false

# Parse command line arguments
while getopts "w:c:n:tph" opt; do
  case $opt in
  w)
    WARNING_THRESHOLD="$OPTARG"
    ;;
  c)
    CRITICAL_THRESHOLD="$OPTARG"
    ;;
  n)
    AVERAGE_SECONDS="$OPTARG"
    ;;
  t)
    SHOW_TOP_PROCESSES=true
    ;;
  p)
    PROXMOX_CALCULATION=true
    ;;
  h)
    echo "Usage: check_iowait.sh [-w warning_threshold] [-c critical_threshold] [-n average_seconds] [-t] [-p] [-h]"
    echo "Options:"
    echo "  -w    Set warning threshold"
    echo "  -c    Set critical threshold"
    echo "  -n    Set average seconds"
    echo "  -t    Show top processes"
    echo "  -p    Enable iowait calculation similar to Proxmox"
    echo "  -h    Print this help message"
    exit 0
    ;;
  \?)
    echo "Usage: check_iowait.sh [-w warning_threshold] [-c critical_threshold] [-n average_seconds] [-t] [-p] [-h]"
    exit 1
    ;;
  esac
done

if ! command -v iostat &>/dev/null; then
  echo "UNKNOWN - iostat not found! Please install sysstat"
  exit -1
fi

# Get iowait value and top 3 processes causing iowait
if uname | grep -q "BSD"; then
  iowait=$(iostat -c $(($AVERAGE_SECONDS + 1)) -w 1 | awk 'NR==4 {print $5}')
  if $SHOW_TOP_PROCESSES; then
    top_processes=$(ps -axo pid,comm,%cpu,%mem | sort -k 3 -nr | head -n 4 | awk 'NR>1 {printf "%s%s", sep, $2; sep=", "} END {print ""}')
  fi
else
  iowait=$(iostat -c $AVERAGE_SECONDS 2 | awk 'NR==4 {print $4}')
  if $PROXMOX_CALCULATION; then
    idle=$(iostat -c $AVERAGE_SECONDS 2 | awk 'NR==4 {print $6}')
    non_idle=$(echo "100 - $idle" | bc -l)
    iowait=$(echo "$iowait / $non_idle * 100" | bc -l)
  fi
  if $SHOW_TOP_PROCESSES; then
    top_processes=$(pidstat -d -l -u -r 1 1 | awk 'NR>4 {print $1, $NF, $8}' | sort -k3 -nr | head -n 3 | awk '{printf "%s%s", sep, $2; sep=", "} END {print ""}')
  fi
fi

if $SHOW_TOP_PROCESSES; then
  iowait_str="- iowait: $iowait%"$'\n'"Top 3 processes causing iowait: $top_processes"
else
  iowait_str="- iowait: $iowait%"
fi
perfdata="iowait=$iowait%;$WARNING_THRESHOLD;$CRITICAL_THRESHOLD;0;100"

# Check if iowait is above thresholds
if (($(echo "$iowait > $CRITICAL_THRESHOLD" | bc -l))); then
  echo "CRITICAL $iowait_str | $perfdata"
  exit 2
elif (($(echo "$iowait > $WARNING_THRESHOLD" | bc -l))); then
  echo "WARNING $iowait_str | $perfdata"
  exit 1
else
  echo -e "OK $iowait_str | $perfdata"
  exit 0
fi