#!/bin/bash

function smart_disks() {
  if [[ $# != 1 ]]; then
    echo "Wrong number of arguments"
    return 3
  fi

  POOL=$1
  DISKS=$(sudo btrfs device stats -c / | grep -oP "^\[[^\]]+\]" | sed 's/[][]//g' | sort -u)
  ERR_OUTPUT=""
  HIGHEST_RES=0
  for disk in $DISKS; do
    OUTPUT=""
    RES=0
    if echo $disk | grep "nvme" > /dev/null; then
      OUTPUT=$(/usr/lib64/nagios/plugins/check_nvme $disk)
      RES=$?
    else
      OUTPUT=$(/usr/lib64/nagios/plugins/check_smart $disk)
      RES=$?
    fi
    if [[ $RES != 0 ]]; then
      ERR_OUTPUT="$ERR_OUTPUT - $OUTPUT"
    fi
    if [[ $RES -gt $HIGHEST_RES ]]; then
      HIGHEST_RES=$RES
    fi
  done

  echo -n "BTRFS POOL SMART $POOL "
  if [[ $HIGHEST_RES == 1 ]]; then
    echo "WARNING${ERR_OUTPUT}"
    return 1
  elif [[ $HIGHEST_RES == 2 ]]; then
    echo "CRITICAL${ERR_OUTPUT}"
    return 2
  elif [[ $HIGHEST_RES == 3 ]]; then
    echo "UNKNOWN${ERR_OUTPUT}"
    return 3
  fi

  echo "OK"
  return 0
}

function pool_errors() {
  if [[ $# != 1 ]]; then
    echo "Wrong number of arguments"
    return 3
  fi

  OUTPUT="$(sudo btrfs device stats -c $1)"
  STATUS=$?

  if [[ $STATUS != 0 ]]; then
    echo "BTRFS POOL ERRORS $1 - CRITICAL: $OUTPUT"
    return 2
  fi

  echo "BTRFS POOL ERRORS $1 - OK"
  return 0
}

if [[ $# -lt 2 ]]; then
  echo "Wrong number of arguments"
  exit 3
fi

ACTION=$1
POOL=$2

if [[ $ACTION == "smart" ]]; then
  smart_disks $POOL
  exit $?
elif [[ $ACTION == "errors" ]]; then
  pool_errors $POOL
  exit $?
fi
echo "Unknown command: $ACTION"
exit 3