#!/usr/bin/env bash

MQTT_HOST='192.168.1.1'
MQTT_PORT=1883
MQTT_USER='mqttuser'
MQTT_PASSWORD='password'

if [ -t 1 ]; then
    ANSI_RESET="$(tput sgr0)"
    ANSI_RED="`[ $(tput colors) -ge 16 ] && tput setaf 9 || tput setaf 1 bold`"
    ANSI_GREEN="`[ $(tput colors) -ge 16 ] && tput setaf 10 || tput setaf 2 bold`"
fi
trap "echo -n \"$ANSI_RESET\"" EXIT
trap "exit 113" HUP INT QUIT PIPE TERM

DISKS=$( lsblk -dno NAME | grep -v '^zd' )
for DISK in $DISKS; do
    MODEL=$( lsblk /dev/$DISK -dno MODEL )
    if [[ "$MODEL" == "" ]]; then
        echo "${ANSI_RED}/dev/$DISK: Cannot determine disk model${ANSI_RESET}" >/dev/stderr
        continue
    fi

    SERIAL=$( lsblk /dev/$DISK -dno SERIAL )
    if [[ "$SERIAL" == "" ]]; then
        echo "${ANSI_RED}/dev/$DISK ($MODEL): Cannot determine disk serial${ANSI_RESET}"
        continue
    fi

    TEMPERATURE=$( smartctl -A /dev/$DISK | grep -E "^194 Temperature_Celsius" | awk '{print $10}' | cut -dh -f1 )
    if [[ "$TEMPERATURE" == "" ]]; then
        TEMPERATURE=$( smartctl -A /dev/$DISK | grep -E '^Temperature Sensor' | cut -d: -f2 | awk '{print $1}' )
    fi
    if [[ "$TEMPERATURE" == "" ]]; then
        echo "${ANSI_RED}/dev/$DISK ($MODEL $SERIAL): Cannot determine disk temperature${ANSI_RESET}"
        continue
    fi

    DEVICE_ID="virtual_disk_${MODEL// /}_${SERIAL// /}"
    DEVICE_MANUFACTURER="Virtual ${HOSTNAME^}"

    TEMPERATURE_ID="virtual_disk_${MODEL// /}_${SERIAL// /}_temperature"
    TEMPERATURE_TOPIC="homeassistant/sensor/$DEVICE_ID/temperature/config"
    TEMPERATURE_STATE_TOPIC="homeassistant/sensor/$DEVICE_ID/temperature/state"

    mosquitto_pub -h $MQTT_HOST -p $MQTT_PORT \
                  -u $MQTT_USER -P $MQTT_PASSWORD \
                  -t $TEMPERATURE_TOPIC -r -m \
        "{
           \"name\": \"Temperature\",
           \"unique_id\": \"$TEMPERATURE_ID\",
           \"state_topic\": \"$TEMPERATURE_STATE_TOPIC\",
           \"device_class\": \"temperature\",
           \"state_class\": \"measurement\",
           \"unit_of_measurement\": \"°C\",
           \"device\": {
             \"identifiers\": [\"$DEVICE_ID\"],
             \"name\": \"Disk $MODEL $SERIAL\",
             \"manufacturer\": \"$DEVICE_MANUFACTURER\",
             \"model\": \"$MODEL\"
           }
        }"
    if [[ $? -ne 0 ]]; then
        echo "${ANSI_RED}/dev/$DISK ($MODEL $SERIAL): Cannot publish temperature device${ANSI_RESET} ($TEMPERATURE_TOPIC)" >/dev/stderr
        continue
    fi

    mosquitto_pub -h $MQTT_HOST -p $MQTT_PORT \
                  -u $MQTT_USER -P $MQTT_PASSWORD \
                  -t "$TEMPERATURE_STATE_TOPIC" -r -m "$TEMPERATURE"
    if [[ $? -ne 0 ]]; then
        echo "${ANSI_RED}/dev/$DISK ($MODEL $SERIAL): Cannot publish temperature state${ANSI_RESET} ($TEMPERATURE_STATE_TOPIC)" >/dev/stderr
        continue
    fi

    HOSTNAME_ID="virtual_disk_${MODEL// /}_${SERIAL// /}_hostname"
    HOSTNAME_TOPIC="homeassistant/sensor/$DEVICE_ID/hostname/config"
    HOSTNAME_STATE_TOPIC="homeassistant/sensor/$DEVICE_ID/hostname/state"

    mosquitto_pub -h $MQTT_HOST -p $MQTT_PORT \
                  -u $MQTT_USER -P $MQTT_PASSWORD \
                  -t $HOSTNAME_TOPIC -r -m \
        "{
           \"name\": \"Host name\",
           \"unique_id\": \"$HOSTNAME_ID\",
           \"state_topic\": \"$HOSTNAME_STATE_TOPIC\",
           \"icon\": \"mdi:server-network-outline\",
           \"device\": {
             \"identifiers\": [\"$DEVICE_ID\"],
             \"name\": \"Disk $MODEL $SERIAL\",
             \"manufacturer\": \"$DEVICE_MANUFACTURER\",
             \"model\": \"$MODEL\"
           }
        }"
    if [[ $? -ne 0 ]]; then
        echo "${ANSI_RED}/dev/$DISK ($MODEL $SERIAL): Cannot publish hostname device${ANSI_RESET} ($HOSTNAME_TOPIC)" >/dev/stderr
        continue
    fi

    mosquitto_pub -h $MQTT_HOST -p $MQTT_PORT \
                  -u $MQTT_USER -P $MQTT_PASSWORD \
                  -t "$HOSTNAME_STATE_TOPIC" -r -m "$( hostname -s )"
    if [[ $? -ne 0 ]]; then
        echo "${ANSI_RED}/dev/$DISK ($MODEL $SERIAL): Cannot publish hostname state${ANSI_RESET} ($HOSTNAME_STATE_TOPIC)" >/dev/stderr
        continue
    fi

    PATH_ID="virtual_disk_${MODEL// /}_${SERIAL// /}_path"
    PATH_TOPIC="homeassistant/sensor/$DEVICE_ID/path/config"
    PATH_STATE_TOPIC="homeassistant/sensor/$DEVICE_ID/path/state"

    mosquitto_pub -h $MQTT_HOST -p $MQTT_PORT \
                  -u $MQTT_USER -P $MQTT_PASSWORD \
                  -t $PATH_TOPIC -r -m \
        "{
           \"name\": \"Path\",
           \"unique_id\": \"$PATH_ID\",
           \"state_topic\": \"$PATH_STATE_TOPIC\",
           \"icon\": \"mdi:harddisk\",
           \"device\": {
             \"identifiers\": [\"$DEVICE_ID\"],
             \"name\": \"Disk $MODEL $SERIAL\",
             \"manufacturer\": \"$DEVICE_MANUFACTURER\",
             \"model\": \"$MODEL\"
           }
        }"
    if [[ $? -ne 0 ]]; then
        echo "${ANSI_RED}/dev/$DISK ($MODEL $SERIAL): Cannot publish path device${ANSI_RESET} ($PATH_TOPIC)" >/dev/stderr
        continue
    fi

    mosquitto_pub -h $MQTT_HOST -p $MQTT_PORT \
                  -u $MQTT_USER -P $MQTT_PASSWORD \
                  -t "$PATH_STATE_TOPIC" -r -m "/dev/$DISK"
    if [[ $? -ne 0 ]]; then
        echo "${ANSI_RED}/dev/$DISK ($MODEL $SERIAL): Cannot publish path state${ANSI_RESET} ($PATH_STATE_TOPIC)" >/dev/stderr
        continue
    fi

    echo "${ANSI_GREEN}/dev/$DISK ($MODEL $SERIAL): $TEMPERATURE °C${ANSI_RESET}"
done