#!/usr/bin/env bash
#
# repair-hdd.sh — overwrite a disk to trigger remapping of pending/bad sectors,
# then run a SMART extended self-test and print a final report.
# Full documentation: see repair-hdd.md in this same folder.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=lib/common.sh
source "${SCRIPT_DIR}/../lib/common.sh"

SCRIPT_HELP=$(cat <<'HELP'
Usage: sudo ./HDD/repair-hdd.sh DEVICE [zero|null|random]

Overwrites the selected disk, then runs a SMART extended self-test. This
destroys all data on DEVICE. Use only with a directly connected SATA disk;
USB adapters may require adding -d sat to the smartctl calls.
At startup, asks you to confirm that you can keep the computer and drive
powered and connected for the hours this operation may take.

Arguments:
    DEVICE             Whole-disk block device, for example /dev/sdX
    zero               Write zeros (default; also accepted as "null")
    random             Write pseudorandom data

Options:
    -h, --help         Show this help message and exit

Examples:
    sudo ./HDD/repair-hdd.sh /dev/sdX
    sudo ./HDD/repair-hdd.sh /dev/sdX random
HELP
)

common::show_help_if_requested "$@"

[[ $# -le 2 ]] || common::die "too many arguments; run with --help for usage"
LOG_DIR="/var/log/hdd-repair"
DEVICE="${1:-}"
FILL_MODE="${2:-zero}"

# --- Валидация на входа ---------------------------------------------------

if [[ -z "$DEVICE" ]]; then
    common::die "a whole-disk block device is required; run with --help for usage"
fi

case "$FILL_MODE" in
    zero|null)
        FILL_MODE="zero"
        ;;
    random)
        ;;
    *)
        common::die "fill mode must be 'zero', 'null', or 'random'"
        ;;
esac

common::log "WARNING: This operation may take hours. Do not power off, reboot, or suspend the computer, or disconnect power or the drive."
common::confirm "Can the computer and drive remain powered and connected until this operation finishes?" || exit 0

for required_command in findmnt lsblk readlink swapon mkdir date tee smartctl grep umount dd sync sleep head base64; do
    common::require_cmd "$required_command"
done

if [[ "$EUID" -ne 0 ]]; then
    common::die "this script must be run as root (for example, with sudo)"
fi

if [[ ! -b "$DEVICE" ]]; then
    common::die "$DEVICE does not exist or is not a block device"
fi

# Refuse the root filesystem device and any block device it depends on.
TARGET_DEVICE="$(readlink -f -- "$DEVICE")" || common::die "cannot resolve device path: $DEVICE"
ROOT_SOURCE="$(findmnt -n -o SOURCE /)" || common::die "cannot determine the root filesystem device"
ROOT_SOURCE="${ROOT_SOURCE%%\[*}"
ROOT_DEVICE="$(readlink -f -- "$ROOT_SOURCE")" || common::die "cannot resolve root device path: $ROOT_SOURCE"
ROOT_DEVICE_TREE="$(lsblk --inverse --list --noheadings --paths --output PATH "$ROOT_DEVICE")" \
    || common::die "cannot determine the devices backing the root filesystem"

while IFS= read -r root_device; do
    [[ -n "$root_device" ]] || continue
    if [[ "$(readlink -f -- "$root_device")" == "$TARGET_DEVICE" ]]; then
        common::die "$DEVICE backs the root filesystem; refusing to overwrite it"
    fi
done <<< "$ROOT_DEVICE_TREE"

DEVICE_TREE="$(lsblk --list --noheadings --paths --output PATH "$TARGET_DEVICE")" \
    || common::die "cannot inspect block devices below $DEVICE"

while IFS= read -r child_device; do
    [[ -n "$child_device" ]] || continue
    while IFS= read -r swap_device; do
        [[ -n "$swap_device" ]] || continue
        if [[ "$(readlink -f -- "$swap_device")" == "$(readlink -f -- "$child_device")" ]]; then
            common::die "$child_device is active swap; disable it before running this script"
        fi
    done < <(swapon --show --noheadings --raw --output NAME)
done <<< "$DEVICE_TREE"

if [[ "$FILL_MODE" == "random" ]] && ! command -v openssl >/dev/null 2>&1; then
    common::log "OpenSSL is unavailable; random mode will use /dev/urandom."
fi

mkdir -p "$LOG_DIR"
TIMESTAMP="$(date +%Y%m%d-%H%M%S)"
MAIN_LOG="$LOG_DIR/hdd-repair-${DEVICE##*/}-${TIMESTAMP}.log"

# Send all output to both the console and the log file.
exec > >(tee -a "$MAIN_LOG") 2>&1

common::log "=============================================================="
common::log "HDD Repair — $DEVICE"
common::log "Write mode: $FILL_MODE"
common::log "Started: $(date '+%Y-%m-%d %H:%M:%S')"
common::log "Log: $MAIN_LOG"
common::log "=============================================================="

# --- Disk information before the operation --------------------------------

common::log "--- Device identification ---"
smartctl -i "$DEVICE" || true

common::log "--- SMART values before overwrite ---"
smartctl -A "$DEVICE" | grep -E \
    "Reallocated_Sector_Ct|Current_Pending_Sector|Offline_Uncorrectable|UDMA_CRC_Error_Count|Power_On_Hours" \
    || common::log "SMART read failed; check the connection."

common::log "WARNING: ALL DATA ON $DEVICE WILL BE DESTROYED."
read -r -p "Type uppercase YES to confirm: " CONFIRM
if [[ "$CONFIRM" != "YES" ]]; then
    common::die "operation cancelled; no changes were made"
fi

while IFS= read -r child_device; do
    [[ -n "$child_device" ]] || continue
    if findmnt -n -S "$child_device" >/dev/null 2>&1; then
        common::log "Unmounting $child_device..."
        umount "$child_device" || common::die "could not unmount $child_device; refusing to overwrite $DEVICE"
    fi
done <<< "$DEVICE_TREE"

# --- Phase 1: overwrite the disk to trigger sector remapping ---------------

common::log "--- Phase 1: overwriting $DEVICE ---"
common::log "Started: $(date '+%Y-%m-%d %H:%M:%S')"
common::log "This may take hours, depending on disk size and speed."

DD_START=$(date +%s)
DD_RESULT=0

if [[ "$FILL_MODE" == "zero" ]]; then
    RANDOM_SOURCE_DESC="/dev/zero"
    dd if=/dev/zero of="$DEVICE" bs=4M conv=fsync,noerror status=progress || DD_RESULT=$?
else
    # Use OpenSSL for faster pseudorandom generation when available.
    if command -v openssl >/dev/null 2>&1; then
        RANDOM_SOURCE_DESC="openssl enc (AES-256-CTR PRNG)"
        openssl enc -aes-256-ctr -pass pass:"$(head -c 32 /dev/urandom | base64)" \
                -nosalt < /dev/zero 2>/dev/null | \
            dd of="$DEVICE" bs=4M conv=fsync,noerror status=progress || DD_RESULT=$?
    else
        RANDOM_SOURCE_DESC="/dev/urandom (по-бавно, openssl не е намерен)"
        dd if=/dev/urandom of="$DEVICE" bs=4M conv=fsync,noerror status=progress || DD_RESULT=$?
    fi
fi

common::log "Data source: $RANDOM_SOURCE_DESC"
if [[ "$DD_RESULT" -eq 0 ]]; then
    common::log "dd completed successfully (no write I/O errors)."
else
    common::log "dd exited with status $DD_RESULT; the disk may have reached its end."
fi
DD_END=$(date +%s)
common::log "Overwrite finished: $(date '+%Y-%m-%d %H:%M:%S')"
common::log "Duration: $(( (DD_END - DD_START) / 3600 ))h $(( ((DD_END - DD_START) % 3600) / 60 ))m"

sync
common::log "Waiting 10 seconds before SMART verification."
sleep 10

common::log "--- SMART values after overwrite ---"
smartctl -A "$DEVICE" | grep -E \
    "Reallocated_Sector_Ct|Current_Pending_Sector|Offline_Uncorrectable|UDMA_CRC_Error_Count" \
    || common::log "SMART read failed."

# --- Phase 2: SMART extended self-test -------------------------------------

common::log "--- Phase 2: starting SMART extended self-test ---"
smartctl -t long "$DEVICE"

# Read the estimated duration from smartctl, falling back to 360 minutes.
EST_MIN="$(smartctl -c "$DEVICE" 2>/dev/null | grep -oP 'Extended self-test routine\s*\K[0-9]+' || true)"
EST_MIN="${EST_MIN:-360}"
WAIT_SECONDS=$(( (EST_MIN + 10) * 60 ))

common::log "Expected self-test duration: about ${EST_MIN} minutes."
common::log "Waiting $(( WAIT_SECONDS / 60 )) minutes before reading the result."
common::log "The drive firmware runs the self-test while this script waits."

sleep "$WAIT_SECONDS"

common::log "--- Self-test result ---"
smartctl -l selftest "$DEVICE"

# --- Final report ----------------------------------------------------------

common::log "=============================================================="
common::log "Final SMART report — $DEVICE"
common::log "Finished: $(date '+%Y-%m-%d %H:%M:%S')"
common::log "=============================================================="
smartctl -A "$DEVICE" | grep -E \
    "Reallocated_Sector_Ct|Current_Pending_Sector|Offline_Uncorrectable|UDMA_CRC_Error_Count|Reallocated_Event_Count"
smartctl -H "$DEVICE"
common::log "Full log: $MAIN_LOG"
common::log "Compare these values with the initial values in the log:"
common::log "Current_Pending_Sector should decrease (ideally to 0)."
common::log "Reallocated_Sector_Ct may increase as sectors are remapped."
common::log "If pending sectors remain, the drive may be unable to write them."
common::log "=============================================================="