#!/bin/bash

# =============================================================================
# Drive Health Checker (v8)
#
# A "pretty" script to list drives, let the user select one,
# and run a SMART health check.
#
# v8 Update:
#   - Replaced string joining with printf for a perfectly
#     aligned, multi-column menu (ID, LABEL, DEVICE, SIZE, MODEL).
#   - Reworked list-building logic to use parallel arrays.
# =============================================================================

# --- Define Colors ---
if command -v tput &> /dev/null; then
    GREEN=$(tput setaf 2)
    RED=$(tput setaf 1)
    YELLOW=$(tput setaf 3)
    BLUE=$(tput setaf 4)
    BOLD=$(tput bold)
    NC=$(tput sgr0) # No Color
else
    GREEN='\033[0;32m'
    RED='\033[0;31m'
    YELLOW='\033[1;33m'
    BLUE='\033[0;34m'
    BOLD='\033[1m'
    NC='\033[0m'
fi

# --- Helper Function for Headers ---
print_header() {
    echo -e "\n${BLUE}${BOLD}--- $1 ---${NC}"
}

# --- 1. System Checks ---

# Check for root privileges
if [ "$EUID" -ne 0 ]; then
  echo -e "${RED}${BOLD}Error:${NC} This script must be run with sudo for drive access."
  echo -e "Please run as: ${GREEN}sudo $0${NC}"
  exit 1
fi

# Check for dependencies
if ! command -v smartctl &> /dev/null || ! command -v lsblk &> /dev/null; then
    echo -e "${YELLOW}Warning:${NC} 'smartctl' or 'lsblk' command not found."
    echo -e "This script requires ${BOLD}smartmontools${NC} and ${BOLD}util-linux${NC}."
    echo -e "As a Manjaro user, you can install them by running:"
    echo -e "${GREEN}sudo pacman -S smartmontools util-linux${NC}"
    exit 1
fi

# --- 2. Get Drive List (once) ---

# --- UPDATED SECTION ---
print_header "Building Drive List (this may take a second)..."

# Declare parallel arrays to hold drive components
declare -a drive_names=()
declare -a drive_labels=()
declare -a drive_sizes=()
declare -a drive_models=()

# Get a list of just the disk names (e.g., sda, sdb)
disks=$(lsblk -d -n -o NAME -e 7,11)

if [ -z "$disks" ]; then
    echo -e "${RED}No physical drives found.${NC}"
    exit 1
fi

# Loop through each disk to build the component lists
for disk in $disks; do
    # Add device name
    drive_names+=("$disk")
    
    # Get Size and Model
    # We use 'read' to handle models with spaces
    read -r size model < <(lsblk -d -n -o SIZE,MODEL "/dev/$disk")
    drive_sizes+=("$size")
    drive_models+=("$model")

    # Get the first available filesystem LABEL
    label=$(lsblk -n -o LABEL "/dev/$disk" | grep -v "^$" | head -n 1)
    
    # Format label for display (or use a placeholder)
    if [ -n "$label" ]; then
        drive_labels+=("($label)")
    else
        drive_labels+=("( )") # Placeholder for alignment
    fi
done
# --- END UPDATED SECTION ---


# --- 3. Main Loop ---

while true; do
    clear
    print_header "Available Physical Drives"
    
    # --- UPDATED MENU PRINTING ---
    # Print the header
    printf "  ${BOLD}%-4s %-16s %-8s %-7s %s${NC}\n" "ID)" "LABEL" "DEVICE" "SIZE" "MODEL"
    printf "  ${BOLD}%-4s %-16s %-8s %-7s %s${NC}\n" "---" "---------------" "------" "------" "-----------------"

    # Loop and print each drive using printf for alignment
    for ((i=0; i<${#drive_names[@]}; i++)); do
        id="$((i+1))"
        # Print the formatted line
        printf "  ${YELLOW}%-4s${NC} ${YELLOW}%-16s${NC} %-8s %-7s %s\n" \
            "$id)" \
            "${drive_labels[$i]}" \
            "${drive_names[$i]}" \
            "${drive_sizes[$i]}" \
            "${drive_models[$i]}"
    done
    # --- END UPDATED MENU ---

    echo -e "  ${YELLOW}q)${NC} Quit"
    echo "" # Add a newline

    # Prompt user
    read -p "${YELLOW}${BOLD}Select a drive (number) or 'q' to quit: ${NC}" selection

    # Process input
    case "$selection" in
        q|Q|quit|Quit)
            echo -e "${GREEN}Exiting.${NC}"
            break
            ;;
        *)
            # Check if it's a valid number
            if ! [[ "$selection" =~ ^[0-9]+$ ]] || [ "$selection" -lt 1 ] || [ "$selection" -gt ${#drive_names[@]} ]; then
                echo -e "\n${RED}Invalid selection.${NC} Press Enter to try again."
                read -p ""
                continue # Restart the loop
            fi

            # --- 4. It's a valid number, get the drive ---
            
            # --- UPDATED PARSING ---
            # Convert selection to zero-based index
            ((idx = selection - 1))
            
            # Get components from our parallel arrays
            device_name="${drive_names[$idx]}"
            model_name="${drive_models[$idx]}"
            label_name="${drive_labels[$idx]}"
            device_path="/dev/$device_name"

            # Build a cleaner header for the report
            print_header "Health Report for $device_path (${BOLD}${label_name} ${model_name}${NC})"
            # --- END UPDATED PARSING ---
            
            if ! smartctl -i "$device_path" | grep -q "SMART support is: Available"; then
                echo -e "${YELLOW}Warning:${NC} This drive (${BOLD}$device_path${NC}) does not report SMART support."
                echo "Cannot perform health check."
            else
                # --- 5. Run Health Check ---
                smartctl -s on "$device_path" &> /dev/null
                echo -e "Running SMART health check..."

                attributes=$(smartctl -A "$device_path")
                health_result=$(smartctl -H "$device_path" | grep 'overall-health' | awk '{print $6}')
                temp=$(echo "$attributes" | grep -i 'Temperature_Celsius' | awk '{print $10}')
                power_on_hours=$(echo "$attributes" | grep -i 'Power_On_Hours' | awk '{print $10}')
                serial_number=$(smartctl -i "$device_path" | grep 'Serial Number:' | awk '{print $3}')
                reallocated_sectors=$(echo "$attributes" | grep -i 'Reallocated_Sector_Ct' | awk '{print $10}')
                pending_sectors=$(echo "$attributes" | grep -i 'Current_Pending_Sector_Ct' | awk '{print $10}')
                uncorrectable_sectors=$(echo "$attributes" | grep -i 'Offline_Uncorrectable' | awk '{print $10}')
                
                [ -z "$reallocated_sectors" ] && reallocated_sectors=0
                [ -z "$pending_sectors" ] && pending_sectors=0
                [ -z "$uncorrectable_sectors" ] && uncorrectable_sectors=0
                
                # --- 6. Display Verdict ---
                echo -e "\n${BOLD}Quick Summary:${NC}"
                echo -e "  ${BLUE}Serial Number:${NC}      $serial_number"
                echo -e "  ${BLUE}Power On Hours:${NC}     $power_on_hours"
                echo -e "  ${BLUE}Temperature:${NC}        ${temp}°C"
                echo -e "\n  ${BOLD}Bad Sector Analysis:${NC}"
                echo -e "  ${BLUE}Reallocated Sectors:${NC}  ${reallocated_sectors}"
                echo -e "  ${BLUE}Pending Sectors:${NC}      ${pending_sectors}"
                echo -e "  ${BLUE}Uncorrectable Sectors:${NC}  ${uncorrectable_sectors}"

                echo -e "\n${BOLD}--- VERDICT ---${NC}"
                
                if [ "$health_result" = "PASSED" ] && [ "$pending_sectors" -eq 0 ] && [ "$uncorrectable_sectors" -eq 0 ]; then
                    echo -e "${GREEN}✅ PASSED${NC}"
                    echo -e "The drive's overall SMART health assessment is ${BOLD}PASSED${NC}."
                    if [ "$reallocated_sectors" -ne 0 ]; then
                        echo -e "${YELLOW}Note:${NC} The drive has reallocated ${BOLD}$reallocated_sectors${NC} sectors. This is not an"
                        echo -e "immediate failure, but you should monitor this value."
                    fi
                elif [ "$health_result" = "PASSED" ]; then
                    echo -e "${YELLOW}⚠️ WARNING${NC}"
                    echo -e "The drive's overall SMART health is ${BOLD}PASSED${NC}, but it reports critical issues!"
                    if [ "$pending_sectors" -ne 0 ]; then
                        echo -e "${RED}Critical:${NC} ${BOLD}$pending_sectors${NC} pending sectors. The drive is struggling to read data."
                    fi
                    if [ "$uncorrectable_sectors" -ne 0 ]; then
                        echo -e "${RED}Critical:${NC} ${BOLD}$uncorrectable_sectors${NC} uncorrectable sectors. ${BOLD}Data has been lost.${NC}"
                    fi
                    echo -e "${YELLOW}ACTION REQUIRED: Back up any critical data immediately.${NC}"
                else
                    echo -e "${RED}❌ FAILED${NC}"
                    echo -e "The drive's overall SMART health assessment is ${BOLD}${health_result}${NC}."
                    echo -e "${YELLOW}ACTION REQUIRED: Back up any critical data immediately"
                    echo -e "and consider replacing this drive.${NC}"
                fi

                # Ask to show full report
                echo "" # Newline
                read -p "${YELLOW}Show full SMART report for $device_path? (y/N): ${NC}" show_full
                if [[ "$show_full" =~ ^[Yy]$ ]]; then
                    print_header "Full SMART Report for $device_path"
                    echo "Look for 'Percentage_Used' or 'Media_Wearout_Indicator' on SSDs."
                    smartctl -a "$device_path" | less
                fi
            fi
            
            # --- 7. Pause before looping ---
            echo -e "\n${YELLOW}Press Enter to return to the drive list...${NC}"
            read -p ""
            ;;
    esac
done

exit 0