#!/bin/sh

# Read-only disk health diagnostics for ZimaOS.
# Exit codes: 0 = no issue found, 1 = warning/incomplete, 2 = unhealthy.

set -u

PROGRAM_NAME=${0##*/}
CHECK_ALL=0
READ_SCAN=0
OVERALL_STATUS=0
SYSTEM_SLOT=""
SYSTEM_DISK=""
DISKS=""
EMMC_CHECKED=0

usage() {
    cat <<EOF
Usage: $PROGRAM_NAME [--all] [--read-scan]

  --all        Check every physical disk, not only the ZimaOS system disk.
  --read-scan  Sequentially read the casaos-system0 slot (no data is written).
  --help       Show this help.

Exit codes: 0 = no issue found, 1 = warning/incomplete, 2 = unhealthy.
EOF
}

info() {
    printf '[INFO] %s\n' "$*"
}

pass() {
    printf '[PASS] %s\n' "$*"
}

warn() {
    printf '[WARN] %s\n' "$*"
    if [ "$OVERALL_STATUS" -lt 1 ]; then
        OVERALL_STATUS=1
    fi
}

fail() {
    printf '[FAIL] %s\n' "$*"
    OVERALL_STATUS=2
}

have() {
    command -v "$1" >/dev/null 2>&1
}

add_disk() {
    candidate=$1
    [ -n "$candidate" ] || return 0
    [ -b "$candidate" ] || return 0

    case "
$DISKS
" in
        *"
$candidate
"*) return 0 ;;
    esac

    if [ -n "$DISKS" ]; then
        DISKS="$DISKS
$candidate"
    else
        DISKS=$candidate
    fi
}

physical_disk_for() {
    device=$1
    resolved=$(readlink -f "$device" 2>/dev/null || true)
    [ -b "$resolved" ] || return 1

    current=$resolved
    while :; do
        parent=$(lsblk -npo PKNAME "$current" 2>/dev/null | sed -n '1p')
        [ -n "$parent" ] || break
        current=$parent
    done

    printf '%s\n' "$current"
}

find_system_disk() {
    if [ -e /dev/disk/by-partlabel/casaos-system0 ]; then
        SYSTEM_SLOT=$(readlink -f /dev/disk/by-partlabel/casaos-system0 2>/dev/null || true)
        SYSTEM_DISK=$(physical_disk_for "$SYSTEM_SLOT" 2>/dev/null || true)
    fi

    if [ -z "$SYSTEM_DISK" ] && have findmnt; then
        root_source=$(findmnt -rn -o SOURCE / 2>/dev/null | sed -n '1p' | sed 's/\[.*$//')
        SYSTEM_DISK=$(physical_disk_for "$root_source" 2>/dev/null || true)
    fi
}

number_after_colon() {
    label=$1
    file=$2
    sed -n "s/^[[:space:]]*${label}[[:space:]]*:[[:space:]]*//p" "$file" \
        | sed -n '1p' | tr -cd '0-9\n' | sed -n '1p'
}

check_smart() {
    disk=$1
    output=$(mktemp /tmp/zimaos-smart.XXXXXX) || {
        warn "Cannot create a temporary file for SMART output."
        return
    }

    smartctl -x "$disk" >"$output" 2>&1
    smart_rc=$?

    printf '%s\n' "  SMART/NVMe health:"
    grep -E 'SMART overall-health|SMART Health Status|SMART support is:|Critical Warning:|Available Spare:|Available Spare Threshold:|Percentage Used:|Media and Data Integrity Errors:|Current Drive Temperature:|Temperature:' "$output" \
        | sed 's/^[[:space:]]*/    /' || true

    if grep -Eiq 'SMART support is:[[:space:]]*(Disabled|Unavailable)' "$output"; then
        warn "SMART health reporting is disabled or unavailable on $disk."
    elif ! grep -Eiq 'SMART overall-health|SMART Health Status:|Critical Warning:' "$output"; then
        warn "smartctl returned no recognized overall-health field for $disk."
    fi

    if grep -Eiq 'SMART overall-health.*FAILED|SMART Health Status:[[:space:]]*(BAD|FAILED)' "$output"; then
        fail "$disk reports a failed SMART health status."
    elif grep -Eiq 'SMART overall-health.*PASSED|SMART Health Status:[[:space:]]*(OK|PASSED)' "$output"; then
        pass "$disk reports a passing SMART health status."
    fi

    critical_warning=$(sed -n 's/^[[:space:]]*Critical Warning:[[:space:]]*//p' "$output" | sed -n '1p' | awk '{print $1}')
    case "$critical_warning" in
        ""|0|0x0|0x00) ;;
        *) fail "$disk reports NVMe Critical Warning $critical_warning." ;;
    esac

    media_errors=$(number_after_colon 'Media and Data Integrity Errors' "$output")
    case "$media_errors" in
        ""|0) ;;
        *) fail "$disk reports $media_errors NVMe media/data-integrity errors." ;;
    esac

    percentage_used=$(number_after_colon 'Percentage Used' "$output")
    if [ -n "$percentage_used" ]; then
        if [ "$percentage_used" -ge 100 ] 2>/dev/null; then
            fail "$disk has reached its rated NVMe endurance ($percentage_used% used)."
        elif [ "$percentage_used" -ge 90 ] 2>/dev/null; then
            warn "$disk is close to its rated NVMe endurance ($percentage_used% used)."
        fi
    fi

    for attribute in Current_Pending_Sector Offline_Uncorrectable; do
        value=$(awk -v name="$attribute" '$2 == name {print $10; exit}' "$output")
        case "$value" in
            ""|0) ;;
            *[!0-9]*) warn "$disk reports an unreadable $attribute value: $value." ;;
            *) fail "$disk reports $attribute=$value." ;;
        esac
    done

    reallocated=$(awk '$2 == "Reallocated_Sector_Ct" {print $10; exit}' "$output")
    case "$reallocated" in
        ""|0) ;;
        *[!0-9]*) warn "$disk reports an unreadable Reallocated_Sector_Ct value: $reallocated." ;;
        *) warn "$disk has $reallocated reallocated sectors." ;;
    esac

    crc_errors=$(awk '$2 == "UDMA_CRC_Error_Count" {print $10; exit}' "$output")
    case "$crc_errors" in
        ""|0) ;;
        *[!0-9]*) warn "$disk reports an unreadable UDMA_CRC_Error_Count value: $crc_errors." ;;
        *) warn "$disk has recorded $crc_errors SATA link CRC errors; check its connector and power path." ;;
    esac

    # smartctl uses a bitmask. Bits 2 and 3 mean failed health or a failing attribute.
    if [ $((smart_rc & 12)) -ne 0 ]; then
        fail "$disk returned SMART failure status (smartctl exit code $smart_rc)."
    elif [ $((smart_rc & 3)) -ne 0 ]; then
        warn "SMART data for $disk could not be read completely (smartctl exit code $smart_rc)."
        sed -n '1,5p' "$output" | sed 's/^/    /'
    elif [ $((smart_rc & 240)) -ne 0 ]; then
        warn "$disk has entries in its SMART error or self-test history (smartctl exit code $smart_rc)."
    fi

    rm -f "$output"
}

check_nvme_cli() {
    disk=$1
    output=$(mktemp /tmp/zimaos-nvme.XXXXXX) || {
        warn "Cannot create a temporary file for NVMe output."
        return
    }

    if ! nvme smart-log "$disk" >"$output" 2>&1; then
        warn "nvme-cli could not read health data from $disk."
        sed -n '1,5p' "$output" | sed 's/^/    /'
        rm -f "$output"
        return
    fi

    printf '%s\n' "  NVMe health:"
    grep -Ei 'critical_warning|available_spare[[:space:]]|available_spare_threshold|percentage_used|media_errors|num_err_log_entries|temperature[[:space:]]' "$output" \
        | sed 's/^[[:space:]]*/    /' || true

    critical_warning=$(awk -F: '/^critical_warning/ {gsub(/[[:space:]]/, "", $2); print $2; exit}' "$output")
    case "$critical_warning" in
        ""|0|0x0|0x00) ;;
        *) fail "$disk reports NVMe Critical Warning $critical_warning." ;;
    esac

    media_errors=$(awk -F: '/^media_errors/ {gsub(/[,[:space:]]/, "", $2); print $2; exit}' "$output")
    case "$media_errors" in
        ""|0) ;;
        *[!0-9]*) warn "$disk reports an unreadable NVMe media_errors value: $media_errors." ;;
        *) fail "$disk reports $media_errors NVMe media errors." ;;
    esac

    percentage_used=$(awk -F: '/^percentage_used/ {gsub(/[%[:space:]]/, "", $2); print $2; exit}' "$output")
    if [ -n "$percentage_used" ]; then
        if [ "$percentage_used" -ge 100 ] 2>/dev/null; then
            fail "$disk has reached its rated NVMe endurance ($percentage_used% used)."
        elif [ "$percentage_used" -ge 90 ] 2>/dev/null; then
            warn "$disk is close to its rated NVMe endurance ($percentage_used% used)."
        fi
    fi

    rm -f "$output"
}

check_emmc() {
    disk=$1
    disk_name=${disk##*/}
    device_dir=/sys/block/$disk_name/device
    [ -d "$device_dir" ] || return

    pre_eol_file=$(find "$device_dir" -name pre_eol_info -type f 2>/dev/null | sed -n '1p')
    life_time_file=$(find "$device_dir" -name life_time -type f 2>/dev/null | sed -n '1p')
    [ -n "$pre_eol_file$life_time_file" ] || return
    EMMC_CHECKED=1

    printf '%s\n' "  eMMC health:"
    if [ -n "$pre_eol_file" ]; then
        pre_eol=$(tr -d '[:space:]' <"$pre_eol_file")
        printf '    Pre-EOL: %s\n' "$pre_eol"
        case "$pre_eol" in
            0x01|0x1|1) pass "$disk eMMC pre-EOL status is normal." ;;
            0x02|0x2|2) warn "$disk eMMC pre-EOL status indicates a warning." ;;
            0x03|0x3|3) fail "$disk eMMC pre-EOL status is urgent." ;;
            *) warn "$disk returned an unknown eMMC pre-EOL value: $pre_eol." ;;
        esac
    fi

    if [ -n "$life_time_file" ]; then
        life_time=$(tr -s '[:space:]' ' ' <"$life_time_file" | sed 's/^ //; s/ $//')
        printf '    Lifetime estimates (A/B): %s\n' "$life_time"
        case "$life_time" in
            *0x0[Bb]*|*0x[Bb]*) fail "$disk eMMC has exceeded its estimated lifetime." ;;
            *0x09*|*0x0[Aa]*|*0x9*|*0x[Aa]*) warn "$disk eMMC is close to its estimated lifetime." ;;
        esac
    fi
}

check_disk() {
    disk=$1
    disk_name=${disk##*/}

    printf '\n=== Disk %s ===\n' "$disk"
    lsblk -dn -o NAME,SIZE,RO,TRAN,MODEL "$disk" 2>/dev/null | sed 's/^/  /' || true

    read_only=$(lsblk -dn -o RO "$disk" 2>/dev/null | tr -d '[:space:]')
    if [ "$read_only" = 1 ]; then
        fail "$disk is read-only at the block-device layer."
    else
        pass "$disk is writable at the block-device layer."
    fi

    smart_checked=0
    EMMC_CHECKED=0
    if have smartctl; then
        check_smart "$disk"
        smart_checked=1
    elif [ "${disk_name#nvme}" != "$disk_name" ] && have nvme; then
        check_nvme_cli "$disk"
        smart_checked=1
    fi

    check_emmc "$disk"

    if [ "$smart_checked" -eq 0 ] && [ "$EMMC_CHECKED" -eq 0 ]; then
        warn "No SMART/NVMe health tool is available for $disk; hardware health is incomplete."
    fi
}

check_kernel_log() {
    output=$(mktemp /tmp/zimaos-kernel-io.XXXXXX) || {
        warn "Cannot create a temporary file for kernel log analysis."
        return
    }

    log_available=0
    if have journalctl && journalctl -k -b --no-pager >"$output" 2>/dev/null; then
        log_available=1
    elif have dmesg && dmesg >"$output" 2>/dev/null; then
        log_available=1
    fi

    if [ "$log_available" -eq 0 ]; then
        warn "Kernel logs could not be read with journalctl or dmesg; I/O errors were not checked."
        rm -f "$output"
        return
    fi

    errors=$(grep -Ei 'I/O error|Buffer I/O|blk_update_request|end_request.*I/O|medium error|uncorrectable|rejecting I/O|device offline|remounting filesystem read-only|nvme.*(reset|timeout|abort)|ata[0-9]+.*(reset|failed command|hard resetting link)' "$output" | tail -n 30 || true)
    if [ -n "$errors" ]; then
        fail "Kernel storage errors were found during the current boot:"
        printf '%s\n' "$errors" | sed 's/^/    /'
    else
        pass "No kernel storage I/O errors were found during the current boot."
    fi

    rm -f "$output"
}

run_read_scan() {
    if [ -z "$SYSTEM_SLOT" ] || [ ! -b "$SYSTEM_SLOT" ]; then
        warn "casaos-system0 was not found; the optional read scan was skipped."
        return
    fi

    printf '\n=== Read scan %s ===\n' "$SYSTEM_SLOT"
    info "Reading the complete system slot. This may take several minutes; no data will be written."
    output=$(mktemp /tmp/zimaos-read-scan.XXXXXX) || {
        warn "Cannot create a temporary file for read-scan output."
        return
    }

    direct_flag=""
    if dd if="$SYSTEM_SLOT" of=/dev/null bs=4096 count=0 iflag=direct 2>/dev/null; then
        direct_flag=iflag=direct
    fi

    # direct_flag is intentionally split into either zero or one dd argument.
    # shellcheck disable=SC2086
    if dd if="$SYSTEM_SLOT" of=/dev/null bs=16M $direct_flag 2>"$output"; then
        pass "The complete $SYSTEM_SLOT slot was read without an I/O error."
    else
        fail "The read scan of $SYSTEM_SLOT failed."
        tail -n 10 "$output" | sed 's/^/    /'
    fi
    rm -f "$output"
}

while [ "$#" -gt 0 ]; do
    case "$1" in
        --all) CHECK_ALL=1 ;;
        --read-scan) READ_SCAN=1 ;;
        --help|-h) usage; exit 0 ;;
        *) printf 'Unknown option: %s\n' "$1" >&2; usage >&2; exit 64 ;;
    esac
    shift
done

if [ "$(id -u)" -ne 0 ]; then
    warn "Run as root for complete SMART and kernel diagnostics."
fi

if ! have lsblk; then
    fail "lsblk is required but was not found."
    exit "$OVERALL_STATUS"
fi

printf 'ZimaOS disk health check\n'
printf 'Time: %s\n' "$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
printf 'Mode: read-only diagnostics\n'

find_system_disk
if [ -n "$SYSTEM_DISK" ]; then
    info "System disk: $SYSTEM_DISK"
    add_disk "$SYSTEM_DISK"
else
    warn "Could not identify the ZimaOS system disk."
fi

if [ "$CHECK_ALL" -eq 1 ]; then
    while IFS= read -r disk; do
        add_disk "$disk"
    done <<EOF
$(lsblk -dnpo NAME,TYPE 2>/dev/null | awk '$2 == "disk" {print $1}')
EOF
fi

if [ -z "$DISKS" ]; then
    fail "No physical disk was found to check."
else
    while IFS= read -r disk; do
        [ -n "$disk" ] && check_disk "$disk"
    done <<EOF
$DISKS
EOF
fi

printf '\n=== Kernel log ===\n'
check_kernel_log

if [ "$READ_SCAN" -eq 1 ]; then
    run_read_scan
fi

printf '\n=== Result ===\n'
case "$OVERALL_STATUS" in
    0) printf 'HEALTHY: no disk-health issue was detected.\n' ;;
    1) printf 'WARNING: the check is incomplete or found warning indicators.\n' ;;
    2) printf 'UNHEALTHY: disk or storage I/O errors were detected. Back up data before retrying an upgrade.\n' ;;
esac
printf 'Note: a passing result reduces risk but cannot guarantee that a disk will not fail.\n'

exit "$OVERALL_STATUS"