#!/bin/bash
# function for removing lockfile
exit-ugreen-leds() {
  if [[ -f "/var/run/ugreen-leds.lock" ]]; then
    rm -f "/var/run/ugreen-leds.lock"
  fi
  kill $smart_check_pid 2>/dev/null
  kill $netdev_check_pid 2>/dev/null
  kill $disk_online_check_pid 2>/dev/null
}

# trap exit and remove lockfile
trap 'exit-ugreen-leds' EXIT

# check if script is already running
if [[ -f "/var/run/ugreen-leds.lock" ]]; then
  echo "ugreen-leds: already running!"
  logger "ugreen-leds: Error: ugreen-leds already running!"
  exit 1
fi
touch /var/run/ugreen-leds.lock

# use variables from config file (unRAID specific)
source /boot/config/plugins/ugreenleds-driver/settings.cfg

# hctl, $> lsblk -S -x hctl -o hctl,serial,name 
# NOTE: It is reported that the order below should be adjusted for each model. 
#     Please check the disk mapping section in https://github.com/miskcoo/ugreen_dx4600_leds_controller/blob/master/README.md.
hctl_map=("0:0:0:0" "1:0:0:0" "2:0:0:0" "3:0:0:0" "4:0:0:0" "5:0:0:0" "6:0:0:0" "7:0:0:0")
# serial number, $> lsblk -S -x hctl -o hctl,serial,name 
serial_map=(${DISK_SERIAL})
# ata number, $> ls /sys/block | egrep ata\d
ata_map=("ata1" "ata2" "ata3" "ata4" "ata5" "ata6" "ata7" "ata8")
# create led map
led_map=(disk1 disk2 disk3 disk4 disk5 disk6 disk7 disk8)

# get default unRAID network interface
INTERFACE="$(cat /boot/config/network.cfg | grep "IFNAME\[0\]" | cut -d '=' -f2 | sed 's/\"//g' | head -1 | tr -d '[:space:]')"
if [ -z "${INTERFACE}" ]; then
  INTERFACES="$(ls /sys/class/net | grep -E '^eth[0-9]+|^bond[0-9]+|^br[0-9]+' | sed 's/@$//')"
  if [[ $INTERFACES =~ ^bond ]]; then
    INTERFACE="$(echo "$INTERFACES" | grep "bond" | head -1 | tr -d '[:space:]')"
  elif [[ $INTERFACES =~ ^br ]]; then
    INTERFACE="$(echo "$INTERFACES" | grep "br" | head -1 | tr -d '[:space:]')"
  elif [[ $INTERFACES =~ ^eth ]]; then
    INTERFACE="$(echo "$INTERFACES" | grep "eth" | head -1 | tr -d '[:space:]')"
  fi
fi

# define functions
default_modules() {
  /sbin/modprobe i2c-dev
  /sbin/modprobe led-ugreen
  sleep 1
  I2C_DEV=$(i2cdetect -l | grep "SMBus I801 adapter" | grep -Po "i2c-\d+")
  echo "led-ugreen 0x3a" > /sys/bus/i2c/devices/${I2C_DEV}/new_device 2>/dev/null
  /sbin/modprobe ledtrig-oneshot
  /sbin/modprobe ledtrig-netdev
  sleep 2
}

default_netdevmon() {
  echo netdev > /sys/class/leds/netdev/trigger
  echo ${INTERFACE} > /sys/class/leds/netdev/device_name
  echo 1 > /sys/class/leds/netdev/link
  echo 1 > /sys/class/leds/netdev/tx
  echo 1 > /sys/class/leds/netdev/rx
  echo 500 > /sys/class/leds/netdev/interval
}

dxp480t_static() {
  /sbin/modprobe i2c-dev
  sleep 1
  I2C_DEV=$(i2cdetect -l | grep "SMBus I801 adapter" | grep -Po "i2c-\d+" | cut -d '-' -f2)

  ## Set status led to static white
  i2cset -y ${I2C_DEV} 0x26 0xa0 1 b # red off
  i2cset -y ${I2C_DEV} 0x26 0xa0 2 b # white off
  i2cset -y ${I2C_DEV} 0x26 0xb1 2 b # white on
  i2cset -y ${I2C_DEV} 0x26 0x51 0 b # Slow Flash off
  i2cset -y ${I2C_DEV} 0x26 0x50 0 b # Permanently on
}

product_name=$(dmidecode --string system-product-name)
case "${product_name}" in 
  DXP6800*)   # tested on DXP6800 Pro
    logger "ugreen-leds: Info: Found UGREEN DXP6800 series" 
    hctl_map=("2:0:0:0" "3:0:0:0" "4:0:0:0" "5:0:0:0" "0:0:0:0" "1:0:0:0")
    ata_map=("ata3" "ata4" "ata5" "ata6" "ata1" "ata2")
    default_modules
    default_netdevmon
    ;;
  DX4600*)   # tested on DX4600 Pro
    logger "ugreen-leds: Info: Found UGREEN DX4600 series"
    default_modules
    default_netdevmon
    ;;
  DX4700*)
    logger "ugreen-leds: Info: Found UGREEN DX4700 series"
    default_modules
    default_netdevmon
    ;;
  DXP2800*) 
    logger "ugreen-leds: Info: UGREEN DXP2800 series"
    default_modules
    default_netdevmon
    ;;
  DXP4800*) 
    logger "ugreen-leds: Info: UGREEN DXP4800 series"
    default_modules
    default_netdevmon
    ;;
  DXP8800*)  # tested on DXP8800 Plus
    logger "ugreen-leds: Info: UGREEN DXP8800 series"
    default_modules
    default_netdevmon
    # using the default mapping
    ;;
  DXP480T*)  # tested on DXP480T Plus
    logger "ugreen-leds: Info: UGREEN DXP480T series"
    dxp480t_static
    # set led to static white and exit
    exit 0
    ;;
  *)
    if [[ "${MAPPING_METHOD}" == "hctl" || "${MAPPING_METHOD}" == "ata" ]]; then
      logger "ugreen-leds: Error: no supported UGREEN NAS found!"
    fi
    exit 1
    ;;
esac
declare -A devices

function disk_enumerating_string() {
  if [[ $MAPPING_METHOD == ata ]]; then
    ls -ahl /sys/block | sed 's/\/$//' | awk '{
      if (match($0, /ata[0-9]+/)) {
        ata = substr($0, RSTART, RLENGTH);
        if (match($0, /[^\/]+$/)) {
          basename = substr($0, RSTART, RLENGTH);
          print basename, ata;
        }
      }
    }'
  elif [[ $MAPPING_METHOD == hctl || $MAPPING_METHOD == serial ]]; then
    lsblk -S -o name,${MAPPING_METHOD},tran | grep sata
  else
    logger "ugreen-leds: Error: Unsupported mapping method: ${MAPPING_METHOD}"
    exit 1
  fi
}

function set_netdev_color() {
  case $(cat /sys/class/net/${1}/speed) in
    100) color="${COLOR_NETDEV_LINK_100}";;
    1000) color="${COLOR_NETDEV_LINK_1000}";;
    2500) color="${COLOR_NETDEV_LINK_2500}";;
    10000) color="${COLOR_NETDEV_LINK_10000}";;
  esac
  if [ -z "${color}" ]; then
    color="${COLOR_NETDEV_NORMAL}"
  fi
  echo "$color" > /sys/class/leds/netdev/color
}

logger "ugreen-leds: Info: Enumerating disks based on $MAPPING_METHOD..."
declare -A dev_map
while read line
do
  blk_line=($line)
  key=${blk_line[1]}
  val=${blk_line[0]}
  dev_map[${key}]=${val}
  logger "ugreen-leds: Info: mapping $MAPPING_METHOD ${key} to ${dev_map[${key}]}"
done <<< "$(disk_enumerating_string)"

# initialize LEDs
declare -A dev_to_led_map
for i in "${!led_map[@]}"; do
  led=${led_map[i]} 
  if [[ -d /sys/class/leds/$led ]]; then
    echo oneshot > /sys/class/leds/$led/trigger
    echo 1 > /sys/class/leds/$led/invert
    echo 100 > /sys/class/leds/$led/delay_on
    echo 100 > /sys/class/leds/$led/delay_off
    echo "$COLOR_DISK_HEALTH" > /sys/class/leds/$led/color
    echo "$BRIGHTNESS_DISK_LEDS" > /sys/class/leds/$led/brightness
    # find corresponding device
    _tmp_str=${MAPPING_METHOD}_map[@]
    _tmp_arr=(${!_tmp_str})
    
    if [[ -v "dev_map[${_tmp_arr[i]}]" ]]; then
      dev=${dev_map[${_tmp_arr[i]}]}

      if [[ -f /sys/class/block/${dev}/stat ]]; then
        devices[$led]=${dev}
        dev_to_led_map[$dev]=$led
      else
        # turn off the led if no disk installed on this slot
        echo 0 > /sys/class/leds/$led/brightness
        echo none > /sys/class/leds/$led/trigger
      fi
    else
      # turn off the led if no disk installed on this slot
      echo 0 > /sys/class/leds/$led/brightness
      echo none > /sys/class/leds/$led/trigger
    fi
  fi
done

# check gw connection
if [ "$CHECK_GATEWAY_CONNECTIVITY" == "true" ]; then
(
  while true; do
    gw=$(ip route | awk '/default/ { print $3 }' | head -1)
    if ping -q -c 1 -W 1 $gw >/dev/null; then
      set_netdev_color ${INTERFACE}
    else
      echo $COLOR_NETDEV_GATEWAY_UNREACHABLE > /sys/class/leds/netdev/color
    fi
    sleep ${CHECK_NETDEV_INTERVAL}s
  done
) &
netdev_check_pid=$!
fi

# check disk online status
(
  while true; do
    for led in "${!devices[@]}"; do
      dev=${devices[$led]}
      if [[ "$(cat /sys/class/leds/$led/color)" != "$COLOR_DISK_HEALTH" ]]; then
        continue;
      fi

      if [[ ! -f /sys/class/block/${dev}/stat ]]; then
        echo "$COLOR_DISK_UNAVAIL" > /sys/class/leds/$led/color 2>/dev/null
        logger "ugreen-leds: Error: Disk /dev/$dev went offline at $(date +%Y-%m-%d' '%H:%M:%S)"
        continue
      fi
    done
    sleep ${CHECK_DISK_ONLINE_INTERVAL}s
  done
) &
disk_online_check_pid=$!

# monitor disk activities
declare -A diskio_data_rw
while true; do
  for led in "${!devices[@]}"; do
    # if $dev does not exist, diskio_new_rw="", which will be safe
    diskio_new_rw="$(cat /sys/block/${devices[$led]}/stat 2>/dev/null)"
    if [ "${diskio_data_rw[$led]}" != "${diskio_new_rw}" ]; then
      echo 1 > /sys/class/leds/$led/shot
    fi
    diskio_data_rw[$led]=$diskio_new_rw
  done
  sleep ${LED_REFRESH_INTERVAL}s
done