#!/usr/bin/env bash


###
#
# Options Section
#
###

set -e
set -o pipefail


###
#
# Variables Section
#
###

if [[ -f '/etc/ugreen-nas.conf' ]]; then
  source '/etc/ugreen-nas.conf'
fi

DISK_ZFS="${DISK_ZFS:='false'}"
DISK_SMART="${DISK_SMART:='false'}"
DISK_LED_BRIGHTNESS="${DISK_LED_BRIGHTNESS:='255'}"
DISK_LED_COLOR_HEALTH="${DISK_LED_COLOR_HEALTH:='0 255 0'}"
DISK_LED_COLOR_FAILED="${DISK_LED_COLOR_FAILED:='255 0 0'}"
DISK_LED_COLOR_STANDBY="${DISK_LED_COLOR_STANDBY:='0 0 255'}"

MODEL=$(/usr/sbin/dmidecode --string 'system-product-name')
SERIAL_MAP=(${DISK_SERIAL})

declare -A DEVICES
declare -A DEV_MAP
declare -A DEV_LED_MAP
declare -A ZFS_LED_MAP


###
#
# Function Section
#
###

function trap_exit () {
  if [[ -n "${PID_ZFS}" ]]; then
    /usr/bin/kill "${PID_ZFS}" &> '/dev/null'
  fi
  if [[ -n "${PID_SMART}" ]]; then
    /usr/bin/kill "${PID_SMART}" &> '/dev/null'
  fi
  if [[ -n "${PID_DISK}" ]]; then
    /usr/bin/kill "${PID_DISK}" &> '/dev/null'
  fi
  if [[ -n "${PID_DISK_STANDBY}" ]]; then
    /usr/bin/kill "${PID_DISK_STANDBY}" &> '/dev/null'
  fi
}

function load_module () {
  MODULE="${1}"
  if ! /usr/sbin/lsmod | /usr/bin/grep --quiet "${MODULE}"; then
    /usr/sbin/modprobe "${MODULE}"
  fi
}

function check_disk_health () {
  if [[ "${1}" == "${DISK_LED_COLOR_HEALTH}" || \
        "${1}" == "${DISK_LED_COLOR_STANDBY}" ]]; then
    return 0
  else
    return 1
  fi
}

function check_disk_zfs () {
  while true; do
    while read LINE; do
      LINE=(${LINE})
      DEV_NAME="${LINE[0]}"
      DEV_STATE="${LINE[1]}"
      if [[ -v "ZFS_LED_MAP[${DEV_NAME}]" ]]; then
        LED="${ZFS_LED_MAP[${DEV_NAME}]}"
        LED_COLOR=$(/usr/bin/cat "/sys/class/leds/${LED}/color")
        if ! check_disk_health "${LED_COLOR}"; then
          continue;
        fi
        if [[ "${DEV_STATE}" != "ONLINE" ]]; then
          /usr/bin/echo "${COLOR_ZPOOL_FAIL}" > "/sys/class/leds/${LED}/color"
        fi
      fi
    done <<< $(/usr/bin/zpool status -L | /usr/bin/grep --extended-regexp ^\\s*\(nvme\|dm\))
    /usr/bin/sleep '15s'
  done
}

function disk_parameters () {
  for LED in "${!DEVICES[@]}"; do
    /usr/bin/echo "${DEVICES[${LED}]} ${LED}"
  done
}


###
#
# Runtime Environment
#
###

if [[ -z "${DISK_SERIAL}" ]]; then
  /usr/bin/echo "The serial number of the hard drives could not be found in the configuration."
  exit 0
fi

trap 'trap_exit' EXIT

case "${MODEL}" in
  DXP2800*)
    LED_MAP=(disk1 disk2)
  ;;
  DH4300*)
    LED_MAP=(disk1 disk2 disk3 disk4)
  ;;
  DXP4800*)
    LED_MAP=(disk1 disk2 disk3 disk4)
  ;;
  DXP6800*)
    LED_MAP=(disk1 disk2 disk3 disk4 disk5 disk6)
  ;;
  DXP8800*)
    LED_MAP=(disk1 disk2 disk3 disk4 disk5 disk6 disk7 disk8)
  ;;
  *)
    /usr/bin/echo "The UGREEN NAS ${MODEL} is not supported."
    exit 1
  ;;
esac

load_module 'ledtrig_oneshot'

while read LINE; do
  LINE=(${LINE})
  SERIAL="${LINE[1]}"
  DISK="${LINE[0]}"
  DEV_MAP[${SERIAL}]="${DISK}"
done <<< $(/usr/bin/lsblk --all --output 'name,serial,tran' | /usr/bin/grep --extended-regexp '^nvme[0-9]+n[0-9]+ ')

for MAP in "${!LED_MAP[@]}"; do
  LED="${LED_MAP[MAP]}"
  if [[ -d "/sys/class/leds/${LED}" ]]; then
    /usr/bin/echo 'oneshot' > "/sys/class/leds/${LED}/trigger"
    /usr/bin/echo '1' > "/sys/class/leds/${LED}/invert"
    /usr/bin/echo '100' > "/sys/class/leds/${LED}/delay_on"
    /usr/bin/echo '100' > "/sys/class/leds/${LED}/delay_off"
    /usr/bin/echo "${DISK_LED_COLOR_HEALTH}" > "/sys/class/leds/${LED}/color"
    /usr/bin/echo "${DISK_LED_BRIGHTNESS}" > "/sys/class/leds/${LED}/brightness"
    TMP_STR="SERIAL_MAP[@]"
    TMP_ARR=("${!TMP_STR}")
    if [[ -v "DEV_MAP[${TMP_ARR[MAP]}]" ]]; then
      DEV="${DEV_MAP[${TMP_ARR[MAP]}]}"
      if [[ -f "/sys/class/block/${DEV}/stat" ]]; then
        DEVICES[${LED}]="${DEV}"
        DEV_LED_MAP[${DEV}]="${LED}"
      else
        /usr/bin/echo '0' > "/sys/class/leds/${LED}/brightness"
        /usr/bin/echo 'none' > "/sys/class/leds/${LED}/trigger"
      fi
    else
      /usr/bin/echo '0' > "/sys/class/leds/${LED}/brightness"
      /usr/bin/echo 'none' > "/sys/class/leds/${LED}/trigger"
    fi
  fi
done

if [[ "${DISK_ZFS}" == 'true' ]]; then
  while read LINE; do
    LINE=(${LINE})
    DEV_NAME="${LINE[0]}"
    SCSI_DEV_NAME="unknown"
    case "${DEV_NAME}" in
      nvme*)
        SCSI_DEV_NAME="${DEV_NAME}"
      ;;
      dm*)
        DM_SLAVES=($(/usr/bin/ls "/sys/block/${DEV_NAME}/slaves"))
        SCSI_DEV_NAME="${DM_SLAVES[0]}"
      ;;
      *)
        break
      ;;
    esac
    if [[ -v "DEV_LED_MAP[${SCSI_DEV_NAME}]" ]]; then
      ZFS_LED_MAP[${DEV_NAME}]="${DEV_LED_MAP[${SCSI_DEV_NAME}]}"
    fi
  done <<< $(/usr/bin/zpool status -L | /usr/bin/grep --extended-regexp ^\\s*\(nvme\|dm\))
  check_disk_zfs &
  PID_ZFS="${!}"
fi

if [[ "${DISK_SMART}" == 'true' ]]; then
(
  while true; do
    for LED in "${!DEVICES[@]}"; do
      LED_COLOR=$(/usr/bin/cat "/sys/class/leds/${LED}/color")
      if ! check_disk_health "${LED_COLOR}"; then
        continue;
      fi
      DEV="${DEVICES[${LED}]}"
      /usr/sbin/smartctl --health "/dev/${DEV}" --nocheck='standby,0' &> '/dev/null'
      RETURN="${?}"
      if (( "${RETURN}" & ~32 )); then
        /usr/bin/echo "${DISK_LED_COLOR_FAILED}" > "/sys/class/leds/${LED}/color"
        continue
      fi
    done
    /usr/bin/sleep '900s'
  done
) &
  PID_SMART="${!}"
fi

(
  while true; do
    for LED in "${!DEVICES[@]}"; do
      DEV="${DEVICES[${LED}]}"
      LED_COLOR=$(/usr/bin/cat "/sys/class/leds/${LED}/color")
      if ! check_disk_health "${LED_COLOR}"; then
        continue;
      fi
      if [[ ! -f "/sys/class/block/${DEV}/stat" ]]; then
        /usr/bin/echo "${DISK_LED_COLOR_FAILED}"  > "/sys/class/leds/${LED}/color"
        continue
      fi
    done
    /usr/bin/sleep '15s'
  done
) &
PID_DISK="${!}"

if [[ -x '/usr/sbin/ugreen-led-disk-standby' ]]; then
  /usr/sbin/ugreen-led-disk-standby '15' "${DISK_LED_COLOR_STANDBY}" "${DISK_LED_COLOR_HEALTH}" $(disk_parameters) &
  PID_DISK_STANDBY="${!}"
fi

if [[ -x '/usr/sbin/ugreen-led-disk-blink' ]]; then
  /usr/sbin/ugreen-led-disk-blink '0.1' $(disk_parameters)
else
  declare -A DISKIO_RW
  while true; do
    for LED in "${!DEVICES[@]}"; do
      DISKIO_RW_CMP=$(/usr/bin/cat "/sys/block/${DEVICES[${LED}]}/stat" 2> '/dev/null')
      if [[ "${DISKIO_RW[${LED}]}" != "${DISKIO_RW_COMP}" ]]; then
        /usr/bin/echo '1' > "/sys/class/leds/${LED}/shot"
      fi
      DISKIO_RW[${LED}]="${DISKIO_RW_CMP}"
    done
    /usr/bin/sleep '0.1s'
  done
fi