#!/usr/bin/env bash

set -euo pipefail

SMARTCTL="/usr/sbin/smartctl"
BADBLOCKS="/sbin/badblocks"

usage() {
    echo "Usage: $(basename "$0") <device>" >&2
}

check_dependencies() {
    for program in "$SMARTCTL" "$BADBLOCKS"; do
        if ! command -v "$program" &>/dev/null; then
            echo "Error: $program is not installed or not executable" >&2
            exit 1
        fi
    done
}

validate_device() {
    local device="$1"
    if [ ! -b "$device" ]; then
        echo "Error: $device is not a valid block device" >&2
        usage
        exit 1
    fi
}

enable_smart() {
    local device="$1"
    echo "Enabling SMART on $device..."
    $SMARTCTL -q silent -s on -S on "$device" || {
        echo "SMART enable failed. Retrying..." >&2
        $SMARTCTL -s on -S on "$device" || true
    }
}

check_smart_status() {
    local device="$1"
    echo "Checking SMART status..."
    $SMARTCTL -q silent -H "$device" || { echo "SMART Health Check Failed" >&2; exit 1; }
    $SMARTCTL -q silent -l error "$device" || { echo "SMART Error Log Check Failed" >&2; exit 1; }
    $SMARTCTL -q silent -l selftest "$device" || { echo "SMART Self-Test Log Check Failed" >&2; exit 1; }
}

run_self_tests() {
    local device="$1"
    local test_type="$2"
    local sleep_time="$3"

    echo "Running $test_type test (estimated time: $sleep_time minutes)..."
    $SMARTCTL -q silent -t "$test_type" "$device"
    sleep "${sleep_time}m"
    check_smart_status "$device"
}

perform_write_test() {
    local device="$1"
    echo "Performing non-destructive write test on $device..."
    $BADBLOCKS -s -n "$device" || { echo "Non-destructive write test failed" >&2; exit 1; }
    check_smart_status "$device"
}

parse_polling_times() {
    local device="$1"
    local tmp_file
    tmp_file=$(mktemp /tmp/drivetest.XXXXXXXXXX)

    $SMARTCTL -c "$device" | awk '
        /^(Short|Extended|Conveyance) self-test routine/ {
            test_type = toupper($1)
            getline
            time = $0 ~ /\(\s*[0-9]+\)/ ? $2 : 0
            printf "%s_SLEEP=%d\n", test_type, time
        }
    ' >"$tmp_file"
    source "$tmp_file"
    rm -f "$tmp_file"
}

main() {
    if [ $# -ne 1 ]; then
        usage
        exit 1
    fi

    local device="$1"
    validate_device "$device"
    check_dependencies

    if [[ "$(basename "$device" | sed 's/^\(..\).*$/\1/')" == "sd" ]]; then
        SMARTCTL="$SMARTCTL -d ata"
    fi

    enable_smart "$device"
    check_smart_status "$device"
    parse_polling_times "$device"

    if [ -n "${CONVEYANCE_SLEEP:-}" ]; then
        run_self_tests "$device" "conveyance" "$CONVEYANCE_SLEEP"
    elif [ -n "${SHORT_SLEEP:-}" ]; then
        run_self_tests "$device" "short" "$SHORT_SLEEP"
    else
        echo "Skipping conveyance/short tests..."
    fi

    perform_write_test "$device"

    if [ -n "${EXTENDED_SLEEP:-}" ]; then
        run_self_tests "$device" "long" "$EXTENDED_SLEEP"
    else
        echo "Skipping extended test..."
    fi

    echo "All tests completed successfully."
}

main "$@"