#!/bin/bash
# disk-health-check.sh — argento health monitoring
# Checks SMART health (SATA and NVMe), ZFS pool state and error counters,
# ZFS dataset quotas, and non-ZFS disk space.
#
# Output is plain text to stdout (ANSI-colored when stdout is a TTY),
# or an HTML email to $MAILTO with --email.
#
# Usage:
#   disk-health-check.sh                        # full report to stdout
#   disk-health-check.sh --error-only           # problems only; silent if healthy
#   disk-health-check.sh --email                # full report by email
#   disk-health-check.sh --error-only --email   # email only if problems found
#
# Cron on the server runs --error-only --email daily and --email weekly.
#
# Dependencies: smartmontools, zfs, msmtp (as sendmail)

set -uo pipefail

MAILTO="timothykim@fastmail.fm"
HOSTNAME=$(hostname)
SPACE_THRESHOLD=85  # percent used — alert above this
ZFS_CAP_THRESHOLD=80  # ZFS write performance degrades on nearly-full pools
QUOTA_THRESHOLD=85  # percent of dataset quota — alert above this

ERROR_ONLY=false
USE_EMAIL=false
for arg in "$@"; do
    case "$arg" in
        --error-only) ERROR_ONLY=true ;;
        --email) USE_EMAIL=true ;;
        *) echo "Unknown option: $arg" >&2; exit 1 ;;
    esac
done

PROBLEMS=""
STATUS=""

# HTML table rows for the email body, built alongside the plain-text STATUS.
# Styles are inlined on every element because most mail clients strip
# <style> blocks. nowrap keeps device paths and sizes from wrapping.
TH='style="padding:2px 10px;text-align:left;border-bottom:1px solid #ccc;white-space:nowrap"'
TD='style="padding:2px 10px;text-align:left;white-space:nowrap"'
TDR='style="padding:2px 10px;text-align:right;white-space:nowrap"'
TABLE='style="border-collapse:collapse;margin:4px 0 16px"'
HTML_SATA=""
HTML_NVME=""
HTML_ZFS=""
HTML_SPACE=""

# Map a usage percentage to a hex color: green at 0%, amber at 50%,
# red at 100%, linearly interpolated per RGB channel. Endpoints are the
# same dark shades used by state_html so they stay readable on white.
# Done in bash because email clients can't compute colors in CSS.
usage_color() {
    local p=$1 r g b
    [ "$p" -ge 0 ] 2>/dev/null || { printf '#222222'; return; }
    [ "$p" -gt 100 ] && p=100
    if [ "$p" -le 50 ]; then
        # 1a7f37 (green) -> 9a6700 (amber)
        r=$(( 26 + (154 - 26) * p / 50 ))
        g=$(( 127 + (103 - 127) * p / 50 ))
        b=$(( 55 + (0 - 55) * p / 50 ))
    else
        # 9a6700 (amber) -> cf222e (red)
        r=$(( 154 + (207 - 154) * (p - 50) / 50 ))
        g=$(( 103 + (34 - 103) * (p - 50) / 50 ))
        b=$(( 0 + (46 - 0) * (p - 50) / 50 ))
    fi
    printf '#%02x%02x%02x' "$r" "$g" "$b"
}

# Green for healthy, red for anything else. UNKNOWN is amber because it
# means smartctl couldn't read the device, not that the device failed.
state_html() {
    case "$1" in
        PASSED|ONLINE) printf '<span style="color:#1a7f37;font-weight:bold">%s</span>' "$1" ;;
        UNKNOWN)       printf '<span style="color:#9a6700;font-weight:bold">%s</span>' "$1" ;;
        *)             printf '<span style="color:#cf222e;font-weight:bold">%s</span>' "$1" ;;
    esac
}

# Terminal colors (same palette as the HTML), only when stdout is a TTY
# so piped/redirected output stays clean.
COLOR=false
C_BOLD=""
C_RED=""
C_RESET=""
C_DIM=""
if [ -t 1 ]; then
    COLOR=true
    C_BOLD=$'\e[1m'
    C_DIM=$'\e[2m'
    C_RED=$'\e[1;38;2;207;34;46m'
    C_RESET=$'\e[0m'
fi

state_txt() {
    if [ "$COLOR" = true ]; then
        local c
        case "$1" in
            PASSED|ONLINE) c=$'\e[1;38;2;26;127;55m' ;;
            UNKNOWN)       c=$'\e[1;38;2;154;103;0m' ;;
            *)             c=$'\e[1;38;2;207;34;46m' ;;
        esac
        printf '%s%s%s' "$c" "$1" "$C_RESET"
    else
        printf '%s' "$1"
    fi
}

# $1 = percentage (drives the gradient), $2 = text to print. Callers pad
# $2 to its column width BEFORE calling — ANSI codes are invisible but
# still count as characters, so padding after coloring breaks alignment.
usage_txt() {
    if [ "$COLOR" = true ]; then
        local hex
        hex=$(usage_color "$1")
        printf '\e[38;2;%d;%d;%dm%s%s' \
            "$((16#${hex:1:2}))" "$((16#${hex:3:2}))" "$((16#${hex:5:2}))" \
            "$2" "$C_RESET"
    else
        printf '%s' "$2"
    fi
}

# ─── SMART Disk Health (SATA/SAS) ───────────────────────────────
STATUS+="${C_BOLD}SMART (SATA):${C_RESET}\n"
STATUS+="${C_DIM}$(printf '  %-13s %-24s %5s %7s  %s' Disk Model Temp Hours Health)${C_RESET}\n"
for disk in /dev/sd?; do
    [ -b "$disk" ] || continue

    # smartctl returns non-zero exit codes as a bitmask for disk
    # conditions (standby, threshold exceeded, etc.), so don't
    # treat non-zero as a script error — capture output and check it.
    health=$(smartctl -H "$disk" 2>/dev/null | grep "test result" | awk -F': ' '{print $2}') || true
    health="${health:-UNKNOWN}"
    temp=$(smartctl -A "$disk" 2>/dev/null | awk '$1 == 194 || $1 == 190 {print $10; exit}') || true
    temp="${temp:-?}"
    hours=$(smartctl -A "$disk" 2>/dev/null | awk '$1 == 9 {print $10; exit}') || true
    # Some Seagates report attribute 9 raw as "11170h+42m+51.807s" —
    # keep only the whole hours.
    hours=${hours%%[^0-9]*}
    hours="${hours:-?}"
    model=$(smartctl -i "$disk" 2>/dev/null | awk -F': +' '/Device Model|Model Number/{print $2; exit}') || true
    model="${model:-unknown}"

    STATUS+=$(printf '  %-13s %-24s %5s %7s  ' "$disk" "$model" "${temp}C" "${hours}h")$(state_txt "$health")"\n"
    HTML_SATA+="<tr><td $TD>$disk</td><td $TD>$model</td><td $TD>$(state_html "$health")</td><td $TDR>${temp}C</td><td $TDR>${hours}h</td></tr>"$'\n'

    if [ "$health" != "PASSED" ] && [ "$health" != "UNKNOWN" ]; then
        PROBLEMS+="[SMART] $disk: overall health check FAILED ($health)\n"
    fi

    # Check critical SMART attributes
    # 5   = Reallocated_Sector_Ct
    # 187 = Reported_Uncorrect
    # 197 = Current_Pending_Sector
    # 198 = Offline_Uncorrectable
    while IFS= read -r line; do
        id=$(echo "$line" | awk '{print $1}')
        name=$(echo "$line" | awk '{print $2}')
        raw=$(echo "$line" | awk '{print $10}')
        case "$id" in
            5|187|197|198)
                if [ "$raw" -gt 0 ] 2>/dev/null; then
                    PROBLEMS+="[SMART] $disk: $name = $raw (should be 0)\n"
                fi
                ;;
        esac
    done < <(smartctl -A "$disk" 2>/dev/null | awk 'NR>7 && /^[[:space:]]*[0-9]/' || true)
done

# ─── SMART Disk Health (NVMe) ───────────────────────────────────
STATUS+="\n${C_BOLD}SMART (NVMe):${C_RESET}\n"
STATUS+="${C_DIM}$(printf '  %-13s %-24s %5s %7s  %s' Disk Model Temp Hours Health)${C_RESET}\n"
for disk in /dev/nvme[0-9]*; do
    # Match only controller devices (nvme0, nvme1), not namespaces (nvme0n1)
    [[ "$disk" =~ ^/dev/nvme[0-9]+$ ]] || continue
    [ -c "$disk" ] || continue

    ns="${disk}n1"  # namespace device for smartctl
    [ -b "$ns" ] || continue

    health=$(smartctl -H "$ns" 2>/dev/null | grep "SMART overall-health" | awk -F': ' '{print $2}') || true
    health="${health:-UNKNOWN}"
    model=$(smartctl -i "$ns" 2>/dev/null | awk -F': +' '/Model Number/{print $2; exit}') || true
    model="${model:-unknown}"
    # smartctl prints "38 Celsius" — keep just the number to match the SATA column
    temp=$(smartctl -A "$ns" 2>/dev/null | awk -F': +' '/Temperature:/{print $2; exit}') || true
    temp=${temp%% *}
    temp="${temp:-?}"
    # strip thousands separator ("9,852")
    hours=$(smartctl -A "$ns" 2>/dev/null | awk -F': +' '/Power On Hours/{print $2; exit}' | tr -d ',') || true
    hours="${hours:-?}"
    # NVMe percentage used — 100% means full rated write endurance consumed
    pct_used=$(smartctl -A "$ns" 2>/dev/null | awk -F': +' '/Percentage Used/{print $2; exit}' | tr -d '%') || true

    STATUS+=$(printf '  %-13s %-24s %5s %7s  ' "$ns" "$model" "${temp}C" "${hours}h")$(state_txt "$health")"\n"
    HTML_NVME+="<tr><td $TD>$ns</td><td $TD>$model</td><td $TD>$(state_html "$health")</td><td $TDR>${temp}C</td><td $TDR>${hours}h</td></tr>"$'\n'

    if [ "$health" != "PASSED" ] && [ "$health" != "UNKNOWN" ]; then
        PROBLEMS+="[SMART] $ns: overall health check FAILED ($health)\n"
    fi

    if [ -n "$pct_used" ] && [ "$pct_used" -gt 90 ] 2>/dev/null; then
        PROBLEMS+="[SMART] $ns: NVMe endurance ${pct_used}% used (>90%)\n"
    fi
done

# ─── ZFS Pool Health ─────────────────────────────────────────────
STATUS+="\n${C_BOLD}ZFS Pools:${C_RESET}\n"
STATUS+="${C_DIM}$(printf '  %-12s %7s %7s %7s %5s  %s' Pool Size Alloc Free Cap State)${C_RESET}\n"
while IFS= read -r pool; do
    [ -n "$pool" ] || continue
    state=$(zpool list -H -o health "$pool" 2>/dev/null) || true
    state="${state:-UNKNOWN}"
    pool_info=$(zpool list -H -o name,size,alloc,free,cap "$pool" 2>/dev/null) || true

    # zpool -H output is tab-separated
    IFS=$'\t' read -r p_name p_size p_alloc p_free p_cap <<< "$pool_info" || true
    STATUS+=$(printf '  %-12s %7s %7s %7s ' "${p_name:-$pool}" "$p_size" "$p_alloc" "$p_free")$(usage_txt "${p_cap%\%}" "$(printf '%5s' "$p_cap")")"  "$(state_txt "$state")"\n"
    cap_cell="<span style=\"color:$(usage_color "${p_cap%\%}")\">$p_cap</span>"
    HTML_ZFS+="<tr><td $TD>${p_name:-$pool}</td><td $TDR>$p_size</td><td $TDR>$p_alloc</td><td $TDR>$p_free</td><td $TDR>$cap_cell</td><td $TD>$(state_html "$state")</td></tr>"$'\n'

    if [ "$state" != "ONLINE" ]; then
        PROBLEMS+="[ZFS] Pool '$pool' state: $state\n"
    fi

    # Datasets are excluded from the df check below, so pool cap is the
    # only capacity alert for ZFS data.
    if [ "${p_cap%\%}" -gt "$ZFS_CAP_THRESHOLD" ] 2>/dev/null; then
        PROBLEMS+="[ZFS] Pool '$pool' is ${p_cap} full\n"
    fi

    # Check for errors via zpool status -p (parseable).
    # Parse the per-vdev lines: columns are NAME STATE READ WRITE CKSUM.
    # Flag any vdev that is not ONLINE, or has non-zero error counters.
    while IFS= read -r vdev_line; do
        vdev_state=$(echo "$vdev_line" | awk '{print $2}')
        vdev_read=$(echo "$vdev_line" | awk '{print $3}')
        vdev_write=$(echo "$vdev_line" | awk '{print $4}')
        vdev_cksum=$(echo "$vdev_line" | awk '{print $5}')
        vdev_name=$(echo "$vdev_line" | awk '{print $1}')

        if [ "$vdev_state" != "ONLINE" ] && [ "$vdev_state" != "state" ]; then
            PROBLEMS+="[ZFS] Pool '$pool' vdev '$vdev_name' state: $vdev_state\n"
        fi
        for val in "$vdev_read" "$vdev_write" "$vdev_cksum"; do
            if [ "$val" -gt 0 ] 2>/dev/null; then
                PROBLEMS+="[ZFS] Pool '$pool' vdev '$vdev_name' has errors (R:${vdev_read} W:${vdev_write} C:${vdev_cksum})\n"
                break
            fi
        done
    done < <(zpool status -p "$pool" 2>/dev/null | awk '/NAME.*STATE.*READ/{found=1; next} found && /^[[:space:]]+[^ ]/{print} /^$/{found=0}' || true)
done < <(zpool list -H -o name 2>/dev/null || true)

# ─── ZFS Dataset Quotas ──────────────────────────────────────────
# Quotas are hard limits — writes fail at 100% — so warn before the kids
# hit the wall. Lists every dataset with a quota set (with -p, an unset
# quota reads as 0).
STATUS+="\n${C_BOLD}Quotas:${C_RESET}\n"
STATUS+="${C_DIM}$(printf '  %-18s %4s %7s %7s' Dataset 'Use%' Used Quota)${C_RESET}\n"
HTML_QUOTA=""
while IFS=$'\t' read -r ds used quota; do
    [ -n "$ds" ] || continue
    [ "$quota" -gt 0 ] 2>/dev/null || continue
    pct=$(( used * 100 / quota ))
    used_h=$(numfmt --to=iec "$used" 2>/dev/null || echo "$used")
    quota_h=$(numfmt --to=iec "$quota" 2>/dev/null || echo "$quota")
    STATUS+=$(printf '  %-18s ' "$ds")$(usage_txt "$pct" "$(printf '%4s' "${pct}%")")$(printf ' %7s %7s' "$used_h" "$quota_h")"\n"
    weight=""
    if [ "$pct" -gt "$QUOTA_THRESHOLD" ]; then
        PROBLEMS+="[QUOTA] $ds is at ${pct}% of its ${quota_h} quota\n"
        weight=";font-weight:bold"
    fi
    HTML_QUOTA+="<tr><td $TD>$ds</td><td $TDR>$used_h</td><td $TDR>$quota_h</td><td $TDR><span style=\"color:$(usage_color "$pct")$weight\">${pct}%</span></td></tr>"$'\n'
done < <(zfs list -Hp -o name,used,quota -t filesystem 2>/dev/null || true)

# ─── Disk Space ──────────────────────────────────────────────────
STATUS+="\n${C_BOLD}Disk Space:${C_RESET}\n"
STATUS+="${C_DIM}$(printf '  %-12s %4s %7s %7s %7s' Mount 'Use%' Used Size Avail)${C_RESET}\n"
while IFS= read -r line; do
    [ -n "$line" ] || continue
    usage=$(echo "$line" | awk '{print $5}' | tr -d '%')
    mount=$(echo "$line" | awk '{print $6}')
    size=$(echo "$line" | awk '{print $2}')
    used=$(echo "$line" | awk '{print $3}')
    avail=$(echo "$line" | awk '{print $4}')
    STATUS+=$(printf '  %-12s ' "$mount")$(usage_txt "$usage" "$(printf '%4s' "${usage}%")")$(printf ' %7s %7s %7s' "$used" "$size" "$avail")"\n"
    weight=""
    if [ "$usage" -gt "$SPACE_THRESHOLD" ] 2>/dev/null; then
        PROBLEMS+="[SPACE] $mount is ${usage}% full\n"
        weight=";font-weight:bold"
    fi
    usage_cell="<span style=\"color:$(usage_color "$usage")$weight\">${usage}%</span>"
    HTML_SPACE+="<tr><td $TD>$mount</td><td $TDR>$usage_cell</td><td $TDR>$used</td><td $TDR>$size</td><td $TDR>$avail</td></tr>"$'\n'
# ZFS datasets are excluded: their df percentages are computed against
# shared pool free space, so they're misleading — pool cap above covers them.
done < <(df -h --output=source,size,used,avail,pcent,target -x tmpfs -x devtmpfs -x overlay -x efivarfs -x zfs 2>/dev/null | tail -n +2 || true)

# ─── Output ──────────────────────────────────────────────────────
if [ "$ERROR_ONLY" = true ]; then
    if [ -z "$PROBLEMS" ]; then
        exit 0
    fi
    SUBJECT="[ALERT] $HOSTNAME health check failed"
    BODY="Health check found issues on $HOSTNAME at $(date):\n\n${C_RED}$PROBLEMS${C_RESET}\n---\n"
else
    if [ -n "$PROBLEMS" ]; then
        SUBJECT="[REPORT] $HOSTNAME health — ISSUES FOUND"
    else
        SUBJECT="[REPORT] $HOSTNAME health — all clear"
    fi
    BODY="Health report for $HOSTNAME at $(date):\n\n"
    if [ -n "$PROBLEMS" ]; then
        BODY+="${C_RED}*** ISSUES ***\n\n$PROBLEMS${C_RESET}\n"
    fi
    BODY+="--- Status ---\n\n$STATUS\n---\n"
fi

if [ "$USE_EMAIL" = true ]; then
    HTML_PROBLEMS=""
    while IFS= read -r p; do
        [ -n "$p" ] || continue
        HTML_PROBLEMS+="<li style=\"color:#cf222e;font-weight:bold;margin:2px 0\">$p</li>"$'\n'
    done < <(printf '%b' "$PROBLEMS")

    # Built with a newline after each row/section — SMTP caps lines at
    # ~1000 chars and one giant line can get mangled in transit.
    HTML_BODY="<div style=\"font-family:Arial,Helvetica,sans-serif;font-size:14px;color:#222\">"$'\n'
    if [ "$ERROR_ONLY" = true ]; then
        HTML_BODY+="<p>Health check found issues on <b>$HOSTNAME</b> at $(date):</p>"$'\n'
        HTML_BODY+="<ul style=\"padding-left:20px\">"$'\n'"$HTML_PROBLEMS</ul>"$'\n'
    else
        HTML_BODY+="<p>Health report for <b>$HOSTNAME</b> at $(date):</p>"$'\n'
        if [ -n "$HTML_PROBLEMS" ]; then
            HTML_BODY+="<h3 style=\"color:#cf222e\">Issues</h3>"$'\n'
            HTML_BODY+="<ul style=\"padding-left:20px\">"$'\n'"$HTML_PROBLEMS</ul>"$'\n'
        fi
        HTML_BODY+="<h3>SMART (SATA)</h3><table $TABLE>"$'\n'
        HTML_BODY+="<tr><th $TH>Disk</th><th $TH>Model</th><th $TH>Health</th><th $TH>Temp</th><th $TH>Hours</th></tr>"$'\n'
        HTML_BODY+="$HTML_SATA</table>"$'\n'
        HTML_BODY+="<h3>SMART (NVMe)</h3><table $TABLE>"$'\n'
        HTML_BODY+="<tr><th $TH>Disk</th><th $TH>Model</th><th $TH>Health</th><th $TH>Temp</th><th $TH>Hours</th></tr>"$'\n'
        HTML_BODY+="$HTML_NVME</table>"$'\n'
        HTML_BODY+="<h3>ZFS Pools</h3><table $TABLE>"$'\n'
        HTML_BODY+="<tr><th $TH>Pool</th><th $TH>Size</th><th $TH>Alloc</th><th $TH>Free</th><th $TH>Cap</th><th $TH>State</th></tr>"$'\n'
        HTML_BODY+="$HTML_ZFS</table>"$'\n'
        HTML_BODY+="<h3>Quotas</h3><table $TABLE>"$'\n'
        HTML_BODY+="<tr><th $TH>Dataset</th><th $TH>Used</th><th $TH>Quota</th><th $TH>Used %</th></tr>"$'\n'
        HTML_BODY+="$HTML_QUOTA</table>"$'\n'
        HTML_BODY+="<h3>Disk Space</h3><table $TABLE>"$'\n'
        HTML_BODY+="<tr><th $TH>Mount</th><th $TH>Used %</th><th $TH>Used</th><th $TH>Size</th><th $TH>Free</th></tr>"$'\n'
        HTML_BODY+="$HTML_SPACE</table>"$'\n'
    fi
    HTML_BODY+="</div>"

    {
        echo "Subject: $SUBJECT"
        echo "From: $HOSTNAME <argento@fastmail.com>"
        echo "To: $MAILTO"
        echo "MIME-Version: 1.0"
        echo "Content-Type: text/html; charset=utf-8"
        echo ""
        printf '%s\n' "$HTML_BODY"
    } | sendmail "$MAILTO"
else
    printf '%b' "$BODY"
fi