#!/bin/bash
# ==============================================================================================
# ================================= ZFS Pool Scrub ============================================
# ==============================================================================================
#
# PURPOSE
# ─────────────────────────────────────────────────────────────────────────────
# Triggers a ZFS scrub on all pools (or a specific pool), waits for completion,
# and sends a notification with any errors found. Reads every block on every pool
# and verifies checksums — catches silent corruption that would otherwise only
# surface when the corrupted data is read (possibly after redundancy can no
# longer help). Monthly recommended for all pools; quarterly minimum for large pools.
#
# ==============================================================================================
# OPERATIONAL MODEL
# ==============================================================================================
#
# Starts a scrub on each pool, then polls every 60 seconds until all complete.
# Progress is shown every poll — safe to leave running or interrupt. ZFS scrub
# continues in the kernel even if the script is stopped — it does not depend on
# this script remaining alive.
#
# Scrub is safe to run while the pool is in use. It does consume I/O bandwidth —
# schedule during off-peak hours or maintenance windows.
#
# Pools in HOST*_ZFS_REPORT_IGNORE_POOLS are skipped automatically (single-disk
# VM pools, temp pools, etc.). Specifying a pool by name bypasses the ignore list.
#
# ==============================================================================================
# OPERATIONAL SAFEGUARDS
# ==============================================================================================
#
# Single Instance Lock
#   acquire_lock prevents concurrent scrub starts on the same server.
#
# Scrub-in-Progress Check
#   Detects pools already scrubbing and skips them rather than erroring — safe
#   to run when a scrub may have been started by another path.
#
# SIGTERM Trap
#   The poll loop exits cleanly on signal. The ZFS scrub continues regardless.
#
# Tool Validation
#   validate_unraid_cmd confirms zpool and the notify script are present before use.
#
# Silent When Clean
#   Only errors produce visible output and a notification.
#
# ==============================================================================================
# CONFIGURATION
# ==============================================================================================
#
# host*.conf
#
#   HOST*_ZFS_REPORT_IGNORE_POOLS
#       Pools to exclude from automatic scrub. Typically single-disk VM pools
#       or temporary pools that do not need integrity checking.
#       Aliased by detect_hosts() → ZFS_REPORT_IGNORE_POOLS.
#
# ==============================================================================================
# RUNTIME MODES
# ==============================================================================================
#
# zfs_pool_scrub.sh
#     Scrub all pools not in ZFS_REPORT_IGNORE_POOLS. Wait for completion.
#
# zfs_pool_scrub.sh poolname
#     Scrub a specific pool by name. Bypasses the ignore list.
#
# zfs_pool_scrub.sh --status
#     Show current scrub status for all pools and exit.
#
# zfs_pool_scrub.sh --dry-run
#     Show which pools would be scrubbed. No scrub started.
#
# zfs_pool_scrub.sh --log
#     Verbose progress output every poll cycle.
#
# ==============================================================================================

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

source "$SCRIPT_DIR/../load_config.sh"

parse_args "$@"

TARGET_POOL="${PARSED_ARGS[0]:-}"
SCRUB_RUNNING=true

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

validate_unraid_cmd \
    "$(command -v zpool 2>/dev/null || echo /sbin/zpool)" \
    "--version" "" \
    "zpool" || {
        error "ZFS not available on this system — zpool not found"
        exit 1
    }

validate_unraid_cmd \
    "/usr/local/emhttp/plugins/dynamix/scripts/notify" \
    "" "" \
    "unRAID notify script" || warn "unRAID notify script not found — native notifications disabled"

acquire_lock

# detect_hosts() sets MY_ID and aliases HOST*_ZFS_REPORT_IGNORE_POOLS
detect_hosts

# Build ignore pool map
declare -A IGNORE_MAP
for pool in "${ZFS_REPORT_IGNORE_POOLS[@]:-}"; do
    [[ -n "$pool" ]] && IGNORE_MAP["$pool"]=1
done

log "Identity:     $MY_ID ($LOCAL_SERVER_NAME)"
log "Ignore pools: ${ZFS_REPORT_IGNORE_POOLS[*]:-none}"

[[ "$DRY_RUN" == true ]] && warn "DRY RUN — no scrubs will be started"

# SIGTERM trap — exit poll loop cleanly
trap 'warn "ZFS scrub script interrupted — scrub continues in background"; SCRUB_RUNNING=false; exit 0' \
    SIGTERM SIGINT

# ==============================================================================================
# ━━━ Status ━━━
# ==============================================================================================
if [[ "$SHOW_STATUS" == true ]]; then
    echo ""
    echo "━━━━━ $ICON_SUMMARY ZFS SCRUB STATUS ━━━━━"
    echo "$ICON_HOST Identity: $MY_ID ($LOCAL_SERVER_NAME)"
    echo ""
    while IFS= read -r pool; do
        [[ -z "$pool" ]] && continue
        SCAN=$(zpool status "$pool" 2>/dev/null | grep "scan:")
        IGNORED=""
        [[ -n "${IGNORE_MAP[$pool]:-}" ]] && IGNORED=" (ignored)"
        echo "  $ICON_ZFS $pool${IGNORED} — ${SCAN:-no scan data}"
    done < <(zpool list -H -o name 2>/dev/null)
    echo ""
    echo "  Ignored pools: ${ZFS_REPORT_IGNORE_POOLS[*]:-none}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━"
    exit 0
fi

# ==============================================================================================
# ── Build pool list ────────────────────────────────────────────────────────────────────────────
# ==============================================================================================
POOLS_TO_SCRUB=()

if [[ -n "$TARGET_POOL" ]]; then
    # Specific pool — bypass ignore list, validate exists
    if ! zpool list "$TARGET_POOL" >/dev/null 2>&1; then
        error "Pool not found: $TARGET_POOL"
        exit 1
    fi
    POOLS_TO_SCRUB=("$TARGET_POOL")
else
    # All pools — skip ignored ones
    while IFS= read -r pool; do
        [[ -z "$pool" ]] && continue
        if [[ -n "${IGNORE_MAP[$pool]:-}" ]]; then
            log "Skipping $pool (in ZFS_REPORT_IGNORE_POOLS)"
            continue
        fi
        POOLS_TO_SCRUB+=("$pool")
    done < <(zpool list -H -o name 2>/dev/null)
fi

if [[ ${#POOLS_TO_SCRUB[@]} -eq 0 ]]; then
    warn "No pools to scrub — all pools may be on the ignore list"
    warn "Ignored: ${ZFS_REPORT_IGNORE_POOLS[*]:-none}"
    exit 0
fi

log "Pools to scrub: ${POOLS_TO_SCRUB[*]}"

# ==============================================================================================
# ━━━ Start Scrubs ━━━
# ==============================================================================================
echo ""
echo "━━━ $ICON_ZFS Starting ZFS Scrubs — $MY_ID ━━━"
START=$(date +%s)

STARTED=()
SKIPPED_POOLS=()

for pool in "${POOLS_TO_SCRUB[@]}"; do

    # Check if scrub already in progress
    ALREADY=$(zpool status "$pool" 2>/dev/null | grep "scan:" | grep -c "in progress" || true)
    if [[ "$ALREADY" -gt 0 ]]; then
        warn "$pool — scrub already in progress — joining existing scrub"
        STARTED+=("$pool")
        continue
    fi

    log "Starting scrub on $pool..."

    if [[ "$DRY_RUN" == true ]]; then
        warn "DRY RUN — would scrub: $pool"
        STARTED+=("$pool")
    elif zpool scrub "$pool" 2>/dev/null; then
        log "$pool scrub started ✅"
        STARTED+=("$pool")
    else
        error "Failed to start scrub on $pool"
        SKIPPED_POOLS+=("$pool")
    fi
done

if [[ "$DRY_RUN" == true ]]; then
    echo ""
    echo "━━━━━ $ICON_SUMMARY ZFS SCRUB SUMMARY ━━━━━"
    echo "$ICON_HOST  Identity: $MY_ID ($LOCAL_SERVER_NAME)"
    warn "DRY RUN — no scrubs started"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    exit 0
fi

if [[ ${#STARTED[@]} -eq 0 ]]; then
    error "No scrubs were started — check pool status"
    exit 1
fi

# ==============================================================================================
# ━━━ Poll Until Complete ━━━
# ==============================================================================================
echo ""
echo "━━━ $ICON_TIME Waiting for Scrubs to Complete ━━━"
log "Polling every 60 seconds — scrubs may take hours on large pools"
log "Safe to interrupt — scrubs continue in background if script is stopped"

while [[ "$SCRUB_RUNNING" == true ]]; do
    sleep 60

    STILL_RUNNING=false
    for pool in "${STARTED[@]}"; do
        IN_PROGRESS=$(zpool status "$pool" 2>/dev/null | \
            grep "scan:" | grep -c "in progress" || true)
        if [[ "$IN_PROGRESS" -gt 0 ]]; then
            STILL_RUNNING=true
            # Show progress — always visible so user knows it's running
            PROGRESS=$(zpool status "$pool" 2>/dev/null | \
                grep "scan:" | grep -oE "[0-9]+\.[0-9]+% done")
            REPAIRED=$(zpool status "$pool" 2>/dev/null | \
                grep "scan:" | grep -oE "[0-9]+ repaired")
            warn "$pool — scrub in progress ${PROGRESS:+$PROGRESS}${REPAIRED:+ ($REPAIRED)}"
        fi
    done

    [[ "$STILL_RUNNING" == false ]] && SCRUB_RUNNING=false
done

END=$(date +%s)
warn "All scrubs complete — $(format_duration $(( END - START )))"

# ==============================================================================================
# ━━━ Results ━━━
# ==============================================================================================
echo ""
echo "━━━ $ICON_ZFS Scrub Results ━━━"

POOLS_OK=()
POOLS_ERRORS=()

for pool in "${STARTED[@]}"; do
    SCAN_LINE=$(zpool status "$pool" 2>/dev/null | grep "scan:")
    ERRORS=$(zpool status "$pool" 2>/dev/null | \
        grep "errors:" | grep -v "No known data errors")

    if [[ -n "$ERRORS" ]]; then
        error "$pool — ERRORS FOUND"
        error "  $SCAN_LINE"
        error "  $ERRORS"
        POOLS_ERRORS+=("$pool")
    else
        echo "$pool — $SCAN_LINE"
        POOLS_OK+=("$pool")
    fi
done

# ==============================================================================================
# ━━━ Summary ━━━
# ==============================================================================================
echo ""
echo "━━━━━ $ICON_SUMMARY ZFS SCRUB SUMMARY ━━━━━"
echo "$ICON_HOST   Identity:     $MY_ID ($LOCAL_SERVER_NAME)"
echo "$ICON_ZFS    Pools:        ${#POOLS_TO_SCRUB[@]} to scrub"
echo "$ICON_SUCCESS Clean:       ${#POOLS_OK[@]}"
[[ ${#POOLS_ERRORS[@]} -gt 0 ]] && echo "$ICON_ERROR  Errors:      ${#POOLS_ERRORS[@]}"
[[ ${#SKIPPED_POOLS[@]} -gt 0 ]] && warn "Failed start: ${SKIPPED_POOLS[*]}"
echo "$ICON_TIME   Duration:     $(format_duration $(( END - START )))"
echo ""

if [[ ${#POOLS_ERRORS[@]} -gt 0 ]]; then
    echo "$ICON_ERROR Status: ERRORS FOUND — ${POOLS_ERRORS[*]}"
    notify "ZFS scrub errors on $(hostname) ($MY_ID) — pools with errors: ${POOLS_ERRORS[*]}" \
        "ZFS Scrub" "warning"
elif [[ ${#POOLS_OK[@]} -gt 0 ]]; then
    echo "$ICON_DONE Status: all ${#POOLS_OK[@]} pools clean ✅"
fi
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

[[ ${#POOLS_ERRORS[@]} -gt 0 ]] && exit 1
exit 0