#!/bin/bash

# ZFS pool health check - managed by Ansible (infra/nodito/32_zfs_pool_setup_playbook.yml)
#
# The exit code is the answer and systemd keeps it:
#     systemctl is-failed {{ zfs_systemd_health_service_name }}.service
# Reporting anywhere else is optional and generic.

LOG_FILE="{{ zfs_log_file }}"
POOL_NAME="{{ zfs_pool_name }}"
PUSH_URL="${HEALTHCHECK_PUSH_URL:-}"
HOSTNAME=$(hostname)

# Function to log messages
log_message() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') - $1" >> "$LOG_FILE"
}

# Function to check pool health using JSON output
check_pool_health() {
    local pool="$1"
    local issues_found=0
    
    # Get pool status as JSON
    local pool_json
    pool_json=$(zpool status -j "$pool" 2>&1)
    
    if [ $? -ne 0 ]; then
        log_message "ERROR: Failed to get pool status for $pool"
        log_message "  -> $pool_json"
        return 1
    fi
    
    # Check 1: Pool state must be ONLINE
    local pool_state
    pool_state=$(echo "$pool_json" | jq -r --arg pool "$pool" '.pools[$pool].state')
    
    if [ "$pool_state" != "ONLINE" ]; then
        log_message "ISSUE: Pool state is $pool_state (expected ONLINE)"
        issues_found=1
    else
        log_message "OK: Pool state is ONLINE"
    fi
    
    # Check 2: Check all vdevs and devices for non-ONLINE states
    local bad_states
    bad_states=$(echo "$pool_json" | jq -r --arg pool "$pool" '
        .pools[$pool].vdevs[] | 
        .. | objects | 
        select(.state? and .state != "ONLINE") | 
        "\(.name // "unknown"): \(.state)"
    ' 2>/dev/null)
    
    if [ -n "$bad_states" ]; then
        log_message "ISSUE: Found devices not in ONLINE state:"
        echo "$bad_states" | while read -r line; do
            log_message "  -> $line"
        done
        issues_found=1
    else
        log_message "OK: All devices are ONLINE"
    fi
    
    # Check 3: Check for resilvering in progress
    local scan_function scan_state
    scan_function=$(echo "$pool_json" | jq -r --arg pool "$pool" '.pools[$pool].scan_stats.function // "NONE"')
    scan_state=$(echo "$pool_json" | jq -r --arg pool "$pool" '.pools[$pool].scan_stats.state // "NONE"')
    
    if [ "$scan_function" = "RESILVER" ] && [ "$scan_state" = "SCANNING" ]; then
        local resilver_progress
        resilver_progress=$(echo "$pool_json" | jq -r --arg pool "$pool" '.pools[$pool].scan_stats.issued // "unknown"')
        log_message "ISSUE: Pool is currently resilvering (disk reconstruction in progress) - ${resilver_progress} processed"
        issues_found=1
    fi
    
    # Check 4: Check for read/write/checksum errors on all devices
    # Note: ZFS JSON output has error counts as strings, so convert to numbers for comparison
    local devices_with_errors
    devices_with_errors=$(echo "$pool_json" | jq -r --arg pool "$pool" '
        .pools[$pool].vdevs[] | 
        .. | objects | 
        select(.name? and ((.read_errors // "0" | tonumber) > 0 or (.write_errors // "0" | tonumber) > 0 or (.checksum_errors // "0" | tonumber) > 0)) | 
        "\(.name): read=\(.read_errors // 0) write=\(.write_errors // 0) cksum=\(.checksum_errors // 0)"
    ' 2>/dev/null)
    
    if [ -n "$devices_with_errors" ]; then
        log_message "ISSUE: Found devices with I/O errors:"
        echo "$devices_with_errors" | while read -r line; do
            log_message "  -> $line"
        done
        issues_found=1
    else
        log_message "OK: No read/write/checksum errors detected"
    fi
    
    # Check 5: Check for scan errors (from last scrub/resilver)
    local scan_errors
    scan_errors=$(echo "$pool_json" | jq -r --arg pool "$pool" '.pools[$pool].scan_stats.errors // "0"')
    
    if [ "$scan_errors" != "0" ] && [ "$scan_errors" != "null" ] && [ -n "$scan_errors" ]; then
        log_message "ISSUE: Last scan reported $scan_errors errors"
        issues_found=1
    else
        log_message "OK: No scan errors"
    fi
    
    return $issues_found
}

# Function to get last scrub info for status message
get_scrub_info() {
    local pool="$1"
    local pool_json
    pool_json=$(zpool status -j "$pool" 2>/dev/null)
    
    local scan_func scan_state scan_start
    scan_func=$(echo "$pool_json" | jq -r --arg pool "$pool" '.pools[$pool].scan_stats.function // "NONE"')
    scan_state=$(echo "$pool_json" | jq -r --arg pool "$pool" '.pools[$pool].scan_stats.state // "NONE"')
    scan_start=$(echo "$pool_json" | jq -r --arg pool "$pool" '.pools[$pool].scan_stats.start_time // ""')
    
    if [ "$scan_func" = "SCRUB" ] && [ "$scan_state" = "SCANNING" ]; then
        echo "scrub in progress (started $scan_start)"
    elif [ "$scan_func" = "SCRUB" ] && [ -n "$scan_start" ]; then
        echo "last scrub: $scan_start"
    else
        echo "no scrub history"
    fi
}

# Optional reporting to whatever is watching. No push URL is normal, not an
# error: the script's exit code is still a complete answer for anything reading
# unit state.
report() {
    local status="$1"
    local message="$2"
    
    [ -n "$PUSH_URL" ] || return 0
    
    log_message "Reporting ${status}: $message"
    
    # URL encode the message
    local encoded_message
    encoded_message=$(printf '%s\n' "$message" | sed 's/%/%25/g; s/ /%20/g; s/(/%28/g; s/)/%29/g; s/:/%3A/g; s/\//%2F/g')
    
    local response http_code
    response=$(curl -s --max-time 10 --retry 2 -w "\n%{http_code}" "${PUSH_URL}?status=${status}&msg=${encoded_message}&ping=" 2>&1)
    http_code=$(echo "$response" | tail -n1)
    
    if [ "$http_code" = "200" ] || [ "$http_code" = "201" ]; then
        log_message "Report sent successfully (HTTP $http_code)"
        return 0
    else
        log_message "ERROR: Failed to report (HTTP $http_code)"
        return 1
    fi
}

# Main health check logic
main() {
    log_message "=========================================="
    log_message "Starting ZFS health check for pool: $POOL_NAME on $HOSTNAME"
    
    # Run all health checks
    if check_pool_health "$POOL_NAME"; then
        local scrub_info
        scrub_info=$(get_scrub_info "$POOL_NAME")
        
        local message="Pool $POOL_NAME healthy ($scrub_info)"
        report "up" "$message"
        
        log_message "Health check completed: ALL OK"
        exit 0
    else
        log_message "Health check completed: ISSUES DETECTED"
        report "down" "Pool $POOL_NAME unhealthy - see $LOG_FILE"
        exit 1
    fi
}

# Run main function
main