#!/bin/bash
# -----------------------------------------------------------------------------------------------
# --------------------------------- SMART Health Monitor ---------------------------------------
# -----------------------------------------------------------------------------------------------
# Checks SMART health attributes for all drives on the system.
# Reads data live from each drive via smartctl — no persistent writes.
#
# Monitored attributes:
#   Reallocated_Sector_Ct   — bad sectors remapped — any > 0 is concerning
#   Current_Pending_Sector  — sectors waiting for reallocation — any > 0 is concerning
#   Offline_Uncorrectable   — sectors that could not be corrected — any > 0 is critical
#   Temperature_Celsius     — drive temperature vs SMART_TEMP_WARN / SMART_TEMP_CRIT
#   Power_On_Hours          — informational — drive age estimation
#   SMART overall status    — pass/fail per drive
#
# Discovers drives automatically — no configuration needed for drive list.
# SMART_IGNORE_DRIVES allows skipping specific drives (e.g. USB flash drives).
#
# All configuration in Master.conf under SMART Health section.
# Supports --dry-run to show which drives would be checked without running smartctl.
# -----------------------------------------------------------------------------------------------

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

source "$SCRIPT_DIR/../Master.conf"
source "$SCRIPT_DIR/../common.sh"

parse_args "$@"

# -----------------------------------------------------------------------------------------------
# ━━━ $ICON_GEAR Setup ━━━
# -----------------------------------------------------------------------------------------------
echo ""
echo "━━━ $ICON_GEAR Setup ━━━"

if [[ "$EUID" -ne 0 ]]; then
    error "Must be run as root"
    exit 1
fi

success "Running as root"

if ! command -v smartctl >/dev/null 2>&1; then
    error "smartctl not found — install smartmontools"
    notify "SMART health check failed on $(hostname) — smartmontools not installed" "SMART Health" "warning"
    exit 1
fi

success "smartctl available"

# -----------------------------------------------------------------------------------------------
# ━━━ $ICON_SUMMARY Status ━━━
# -----------------------------------------------------------------------------------------------
if [[ "$SHOW_STATUS" == true ]]; then
    echo ""
    echo "━━━━━ $ICON_SUMMARY STATUS ━━━━━"
    echo "$ICON_SMART Temp warn:     ${SMART_TEMP_WARN}°C"
    echo "$ICON_SMART Temp crit:     ${SMART_TEMP_CRIT}°C"
    echo "$ICON_SMART Ignore drives: ${SMART_IGNORE_DRIVES[*]:-none}"
    echo "$ICON_GEAR  Dry Run:       $DRY_RUN"
    echo ""

    # Show drives that would be checked
    echo "━━━ Discovered Drives ━━━"
    for drive in /dev/sd? /dev/nvme?; do
        [[ ! -e "$drive" ]] && continue
        drive_name=$(basename "$drive")
        ignored=false
        for ignore in "${SMART_IGNORE_DRIVES[@]}"; do
            [[ "$drive_name" == "$ignore" ]] && ignored=true && break
        done
        if [[ "$ignored" == true ]]; then
            echo "  $ICON_WARN $drive — ignored"
        else
            echo "  $ICON_SMART $drive — would check"
        fi
    done
    echo "━━━━━━━━━━━━━━━━━━━━━━━"
    exit 0
fi

[[ "$DRY_RUN" == true ]] && warn "DRY RUN — showing drive list only, no SMART data read"

# -----------------------------------------------------------------------------------------------
# HELPER — extract SMART attribute value
# Usage: get_smart_attr "/dev/sda" "Reallocated_Sector_Ct"
# -----------------------------------------------------------------------------------------------
get_smart_attr() {
    local drive="$1" attr="$2"
    smartctl -A "$drive" 2>/dev/null | \
        awk -v attr="$attr" '$2 == attr {print $10}'
}

# -----------------------------------------------------------------------------------------------
# ━━━ $ICON_SMART SMART Health Check ━━━
# -----------------------------------------------------------------------------------------------
echo ""
echo "━━━ $ICON_SMART SMART Health Check — $(date '+%Y-%m-%d %H:%M:%S') ━━━"
echo ""

START=$(date +%s)
DRIVES_OK=()
DRIVES_WARN=()
DRIVES_CRIT=()
DRIVES_SKIP=()

for drive in /dev/sd? /dev/nvme?; do
    [[ ! -e "$drive" ]] && continue
    drive_name=$(basename "$drive")

    # Check ignore list
    ignored=false
    for ignore in "${SMART_IGNORE_DRIVES[@]}"; do
        [[ "$drive_name" == "$ignore" ]] && ignored=true && break
    done

    if [[ "$ignored" == true ]]; then
        info "$drive_name — ignored (in SMART_IGNORE_DRIVES)"
        DRIVES_SKIP+=("$drive_name")
        continue
    fi

    echo "━━━ $ICON_SMART $drive_name ━━━"

    if [[ "$DRY_RUN" == true ]]; then
        warn "DRY RUN — would check $drive_name"
        echo ""
        continue
    fi

    # Check if drive supports SMART
    if ! smartctl -i "$drive" 2>/dev/null | grep -q "SMART support is: Enabled"; then
        warn "$drive_name — SMART not enabled or not supported"
        DRIVES_SKIP+=("$drive_name")
        echo ""
        continue
    fi

    # Overall SMART status
    SMART_STATUS=$(smartctl -H "$drive" 2>/dev/null | grep "overall-health" | awk '{print $NF}')
    if [[ "$SMART_STATUS" == "PASSED" ]]; then
        success "Overall status: PASSED"
    else
        error "Overall status: $SMART_STATUS"
    fi

    # Key attributes
    DRIVE_WARN=false
    DRIVE_CRIT=false

    # Reallocated sectors
    REALLOC=$(get_smart_attr "$drive" "Reallocated_Sector_Ct")
    if [[ -n "$REALLOC" ]]; then
        if [[ "$REALLOC" -gt 0 ]]; then
            warn "$ICON_SMART Reallocated sectors: $REALLOC — drive showing wear"
            DRIVE_WARN=true
        else
            success "$ICON_SMART Reallocated sectors: $REALLOC"
        fi
    fi

    # Pending sectors
    PENDING=$(get_smart_attr "$drive" "Current_Pending_Sector")
    if [[ -n "$PENDING" ]]; then
        if [[ "$PENDING" -gt 0 ]]; then
            warn "$ICON_SMART Pending sectors: $PENDING — sectors awaiting reallocation"
            DRIVE_WARN=true
        else
            success "$ICON_SMART Pending sectors: $PENDING"
        fi
    fi

    # Uncorrectable sectors
    UNCORR=$(get_smart_attr "$drive" "Offline_Uncorrectable")
    if [[ -n "$UNCORR" ]]; then
        if [[ "$UNCORR" -gt 0 ]]; then
            error "$ICON_SMART Uncorrectable sectors: $UNCORR — CRITICAL"
            DRIVE_CRIT=true
        else
            success "$ICON_SMART Uncorrectable sectors: $UNCORR"
        fi
    fi

    # Temperature
    TEMP=$(get_smart_attr "$drive" "Temperature_Celsius")
    # NVMe uses different attribute name
    [[ -z "$TEMP" ]] && TEMP=$(smartctl -A "$drive" 2>/dev/null | \
        awk '/Temperature/{print $2}' | head -1)

    if [[ -n "$TEMP" ]]; then
        if [[ "$TEMP" -ge "$SMART_TEMP_CRIT" ]]; then
            error "$ICON_SMART Temperature: ${TEMP}°C — CRITICAL (threshold: ${SMART_TEMP_CRIT}°C)"
            DRIVE_CRIT=true
        elif [[ "$TEMP" -ge "$SMART_TEMP_WARN" ]]; then
            warn "$ICON_SMART Temperature: ${TEMP}°C — warning (threshold: ${SMART_TEMP_WARN}°C)"
            DRIVE_WARN=true
        else
            success "$ICON_SMART Temperature: ${TEMP}°C"
        fi
    fi

    # Power on hours — informational
    POH=$(get_smart_attr "$drive" "Power_On_Hours")
    if [[ -n "$POH" ]]; then
        POH_DAYS=$(( POH / 24 ))
        info "$ICON_SMART Power on hours: $POH (${POH_DAYS} days)"
    fi

    # Classify drive
    if [[ "$DRIVE_CRIT" == true ]]; then
        DRIVES_CRIT+=("$drive_name")
    elif [[ "$DRIVE_WARN" == true ]]; then
        DRIVES_WARN+=("$drive_name")
    else
        DRIVES_OK+=("$drive_name")
    fi

    echo ""
done

END=$(date +%s)

# -----------------------------------------------------------------------------------------------
# ━━━ $ICON_SUMMARY Summary ━━━
# -----------------------------------------------------------------------------------------------
echo "━━━━━ $ICON_SUMMARY SMART HEALTH SUMMARY ━━━━━"
echo "$ICON_TIME  Duration:  $(format_duration $((END - START)))"
echo ""
echo "  $ICON_SUCCESS Healthy:  ${#DRIVES_OK[@]}  $ICON_WARN Warning: ${#DRIVES_WARN[@]}  $ICON_ERROR Critical: ${#DRIVES_CRIT[@]}  skipped: ${#DRIVES_SKIP[@]}"
echo ""

[[ ${#DRIVES_OK[@]} -gt 0 ]]   && echo "  $ICON_SUCCESS ${DRIVES_OK[*]}"
[[ ${#DRIVES_WARN[@]} -gt 0 ]] && echo "  $ICON_WARN   ${DRIVES_WARN[*]}"
[[ ${#DRIVES_CRIT[@]} -gt 0 ]] && echo "  $ICON_ERROR  ${DRIVES_CRIT[*]}"

echo ""
if [[ "$DRY_RUN" == true ]]; then
    echo "$ICON_WARN Status: DRY RUN"
elif [[ ${#DRIVES_CRIT[@]} -gt 0 ]]; then
    echo "$ICON_ERROR Status: CRITICAL — ${DRIVES_CRIT[*]}"
    notify "SMART CRITICAL on $(hostname) — drives need immediate attention: ${DRIVES_CRIT[*]}" "SMART Health" "warning"
elif [[ ${#DRIVES_WARN[@]} -gt 0 ]]; then
    echo "$ICON_WARN Status: WARNING — ${DRIVES_WARN[*]}"
    notify "SMART WARNING on $(hostname) — drives showing wear: ${DRIVES_WARN[*]}" "SMART Health" "warning"
else
    echo "$ICON_DONE Status: $ICON_SUCCESS ALL DRIVES HEALTHY"
    notify "SMART health check passed on $(hostname) — ${#DRIVES_OK[@]} drives healthy" "SMART Health" "normal"
fi
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"