#!/bin/bash

# smartshot - SMART disk health checker by Evgeniy Shumilov <eashumilov@ya.ru>
# Version: 0.20 - FINAL WITH CORRECT HEALTH CALCULATION

set -euo pipefail

# Colors for output
if [ -t 1 ]; then
    RED='\033[0;31m'
    GREEN='\033[0;32m'
    YELLOW='\033[1;33m'
    BLUE='\033[0;34m'
    PURPLE='\033[0;35m'
    LIGHT_PURPLE='\033[1;95m'
    CYAN='\033[0;36m'
    LIGHT_CYAN='\033[1;96m'
    BOLD='\033[1m'
    NC='\033[0m'
else
    RED=''
    GREEN=''
    YELLOW=''
    BLUE=''
    PURPLE=''
    LIGHT_PURPLE=''
    CYAN=''
    LIGHT_CYAN=''
    BOLD=''
    NC=''
fi

VERSION="0.20"

check_smartctl() {
    if ! command -v smartctl &> /dev/null; then
        echo -e "${RED}Error: smartctl not found${NC}" >&2
        exit 1
    fi
}

# Get all disks via smartctl --scan
get_all_disks() {
    local disks=()
    while IFS= read -r line; do
        if [[ "$line" =~ ^(/dev/[^[:space:]]+) ]]; then
            disks+=("$(basename "${BASH_REMATCH[1]}")")
        fi
    done < <(smartctl --scan 2>/dev/null)
    printf '%s\n' "${disks[@]}" | sort -u
}

# Universal number cleaner - removes ALL non-digit characters
clean_number() {
    echo "$1" | sed -E 's/[^0-9]//g'
}

# Temperature color
get_temp_color() {
    local temp=$1
    if [ -z "$temp" ] || [ "$temp" -lt 0 ]; then
        echo "${NC}"
    elif [ "$temp" -ge 70 ]; then
        echo "${RED}"
    elif [ "$temp" -ge 60 ]; then
        echo "${YELLOW}"
    elif [ "$temp" -ge 50 ]; then
        echo "${YELLOW}"
    else
        echo "${GREEN}"
    fi
}

# Calculate health score based on drive type and metrics
calculate_health_score() {
    local drive_type=$1
    local reallocated=$2
    local vendor_health=$3
    local percentage_used=$4
    local media_errors=$5
    
    local health_score=100
    
    case "$drive_type" in
        "HDD")
            # HDD health based on reallocated sectors
            if [ -n "$reallocated" ]; then
                if [ "$reallocated" -eq 0 ]; then
                    health_score=100
                elif [ "$reallocated" -le 10 ]; then
                    health_score=70   # Fair
                elif [ "$reallocated" -le 100 ]; then
                    health_score=50   # Poor
                elif [ "$reallocated" -le 1000 ]; then
                    health_score=25   # Critical
                else
                    health_score=10   # Failing
                fi
            fi
            ;;
            
        "SATA SSD")
            # SATA SSD health from vendor attributes or wear leveling
            if [ -n "$vendor_health" ]; then
                health_score=$vendor_health
            else
                # Default to good if no vendor health
                health_score=100
            fi
            ;;
            
        "NVMe SSD")
            # NVMe health based on percentage used and media errors
            if [ -n "$percentage_used" ]; then
                health_score=$((100 - percentage_used))
                [ $health_score -lt 0 ] && health_score=0
            fi
            
            # Media errors reduce health score
            if [ -n "$media_errors" ] && [ "$media_errors" -gt 0 ]; then
                local media_health=100
                if [ "$media_errors" -le 10 ]; then
                    media_health=70
                elif [ "$media_errors" -le 100 ]; then
                    media_health=50
                elif [ "$media_errors" -le 1000 ]; then
                    media_health=25
                else
                    media_health=10
                fi
                health_score=$(( health_score < media_health ? health_score : media_health ))
            fi
            ;;
    esac
    
    echo "$health_score"
}

# NVMe function - UNIVERSAL
get_nvme_data() {
    local device="$1"
    local -n _data=$2
    
    _data[MODEL]=""
    _data[SERIAL]=""
    _data[FIRMWARE]=""
    _data[CAPACITY]=""
    _data[CAPACITY_RAW]=""
    _data[TYPE]="NVMe SSD"
    _data[SMART_STATUS]=""
    _data[TEMP]=""
    _data[PERCENTAGE_USED]=""
    _data[POWER_ON_HOURS]=""
    _data[POWER_CYCLES]=""
    _data[AVAIL_SPARE]=""
    _data[MEDIA_ERRORS]=""
    _data[HEALTH_SCORE]="100"
    
    local dev="/dev/$(basename "$device")"
    local smart_info=$(smartctl -i "$dev" 2>/dev/null)
    local smart_health=$(smartctl -H "$dev" 2>/dev/null)
    local smart_attr=$(smartctl -A "$dev" 2>/dev/null)
    
    _data[MODEL]=$(echo "$smart_info" | grep -i "Model Number:" | head -1 | cut -d: -f2- | sed -e 's/^[ \t]*//' -e 's/[ \t]*$//')
    _data[SERIAL]=$(echo "$smart_info" | grep -i "Serial Number:" | head -1 | cut -d: -f2- | sed -e 's/^[ \t]*//' -e 's/[ \t]*$//')
    _data[FIRMWARE]=$(echo "$smart_info" | grep -i "Firmware Version:" | head -1 | cut -d: -f2- | sed -e 's/^[ \t]*//' -e 's/[ \t]*$//')
    
    local cap_line=$(echo "$smart_info" | grep -i "Namespace 1 Size/Capacity:" | head -1)
    if [[ "$cap_line" =~ \[([0-9,.]+)[[:space:]]*([A-Z]+)\] ]]; then
        _data[CAPACITY]=$(echo "${BASH_REMATCH[1]}" | sed 's/[^0-9]//g')
        _data[CAPACITY_RAW]="$(echo "${BASH_REMATCH[1]}" | sed 's/[^0-9.]//g')${BASH_REMATCH[2]}"
    fi
    
    _data[SMART_STATUS]=$(echo "$smart_health" | grep -i "SMART overall-health" | cut -d: -f2- | sed -e 's/^[ \t]*//' -e 's/[ \t]*$//')
    
    local temp_line=$(echo "$smart_attr" | grep -i "^Temperature:" | head -1)
    [[ "$temp_line" =~ ([0-9]+) ]] && _data[TEMP]="${BASH_REMATCH[1]}"
    
    local used_line=$(echo "$smart_attr" | grep -i "Percentage Used:" | head -1)
    [[ "$used_line" =~ ([0-9]+)% ]] && _data[PERCENTAGE_USED]="${BASH_REMATCH[1]}"
    
    # UNIVERSAL hours extraction - take everything after colon and remove NON-digits
    local hours_line=$(echo "$smart_attr" | grep -i "Power On Hours:" | head -1)
    if [ -n "$hours_line" ]; then
        local hours_raw=$(echo "$hours_line" | cut -d: -f2- | sed -e 's/^[ \t]*//' -e 's/[ \t]*$//')
        _data[POWER_ON_HOURS]=$(clean_number "$hours_raw")
    fi
    
    # UNIVERSAL power cycles extraction
    local cycles_line=$(echo "$smart_attr" | grep -i "Power Cycles:" | head -1)
    if [ -n "$cycles_line" ]; then
        local cycles_raw=$(echo "$cycles_line" | cut -d: -f2- | sed -e 's/^[ \t]*//' -e 's/[ \t]*$//')
        _data[POWER_CYCLES]=$(clean_number "$cycles_raw")
    fi
    
    local spare_line=$(echo "$smart_attr" | grep -i "Available Spare:" | head -1)
    [[ "$spare_line" =~ ([0-9]+)% ]] && _data[AVAIL_SPARE]="${BASH_REMATCH[1]}"
    
    local media_line=$(echo "$smart_attr" | grep -i "Media and Data Integrity Errors:" | head -1)
    [[ "$media_line" =~ ([0-9]+) ]] && _data[MEDIA_ERRORS]="${BASH_REMATCH[1]}"
    
    # Calculate health score using universal function
    _data[HEALTH_SCORE]=$(calculate_health_score \
        "${_data[TYPE]}" \
        "" \
        "" \
        "${_data[PERCENTAGE_USED]}" \
        "${_data[MEDIA_ERRORS]}")
    
    [ -n "${_data[MODEL]}" ]
}

# ATA/SATA function - UNIVERSAL
get_ata_data() {
    local device="$1"
    local -n _data=$2
    
    _data[MODEL]=""
    _data[SERIAL]=""
    _data[FIRMWARE]=""
    _data[CAPACITY]=""
    _data[CAPACITY_RAW]=""
    _data[TYPE]=""
    _data[SMART_STATUS]=""
    _data[TEMP]=""
    _data[POWER_ON_HOURS]=""
    _data[POWER_CYCLES]=""
    _data[REALLOCATED_SECTORS]=""
    _data[PENDING_SECTORS]=""
    _data[CRC_ERRORS]=""
    _data[VENDOR_HEALTH]=""
    _data[HEALTH_SCORE]="100"
    
    local dev="/dev/$(basename "$device")"
    local smart_info=$(smartctl -i "$dev" 2>/dev/null)
    local smart_health=$(smartctl -H "$dev" 2>/dev/null)
    local smart_attr=$(smartctl -A "$dev" 2>/dev/null)
    
    _data[MODEL]=$(echo "$smart_info" | grep -i "Device Model:" | head -1 | cut -d: -f2- | sed -e 's/^[ \t]*//' -e 's/[ \t]*$//')
    _data[SERIAL]=$(echo "$smart_info" | grep -i "Serial Number:" | head -1 | cut -d: -f2- | sed -e 's/^[ \t]*//' -e 's/[ \t]*$//')
    _data[FIRMWARE]=$(echo "$smart_info" | grep -i "Firmware Version:" | head -1 | cut -d: -f2- | sed -e 's/^[ \t]*//' -e 's/[ \t]*$//')
    
    local cap_line=$(echo "$smart_info" | grep -i "User Capacity:" | head -1)
    if [[ "$cap_line" =~ \[([0-9,.]+)[[:space:]]*([A-Z]+)\] ]]; then
        local unit="${BASH_REMATCH[2]}"
        local value=$(echo "${BASH_REMATCH[1]}" | sed 's/[^0-9.]//g')
        if [ "$unit" = "TB" ]; then
            _data[CAPACITY_RAW]="${value}TB"
        elif [ "$unit" = "GB" ]; then
            _data[CAPACITY_RAW]="${value}GB"
        fi
    fi
    
    # FIXED: Proper drive type detection
    local rotation=$(echo "$smart_info" | grep -i "Rotation Rate:" | head -1 | cut -d: -f2- | sed -e 's/^[ \t]*//' -e 's/[ \t]*$//')
    
    if echo "$rotation" | grep -qi "solid state"; then
        _data[TYPE]="SATA SSD"
    elif echo "$rotation" | grep -qi "0 rpm"; then
        # Check for SSD-specific features
        if echo "$smart_info" | grep -qi "TRIM Command:"; then
            _data[TYPE]="SATA SSD"
        else
            _data[TYPE]="HDD"  # Enterprise HDDs often show 0 RPM
        fi
    else
        _data[TYPE]="HDD"
    fi
    
    _data[SMART_STATUS]=$(echo "$smart_health" | grep -i "SMART overall-health" | cut -d: -f2- | sed -e 's/^[ \t]*//' -e 's/[ \t]*$//')
    
    # Parse SMART attributes
    while IFS= read -r line; do
        if [[ "$line" =~ ^[[:space:]]*([0-9]+)[[:space:]]+([^[:space:]]+) ]]; then
            local id="${BASH_REMATCH[1]}"
            local values=($line)
            
            case "$id" in
                9)  _data[POWER_ON_HOURS]="${values[9]:-0}" ;;
                12) _data[POWER_CYCLES]="${values[9]:-0}" ;;
                194) _data[TEMP]="${values[9]:-0}" ;;
                5)  _data[REALLOCATED_SECTORS]="${values[9]:-0}" ;;
                197) _data[PENDING_SECTORS]="${values[9]:-0}" ;;
                199) _data[CRC_ERRORS]="${values[9]:-0}" ;;
                169) _data[VENDOR_HEALTH]="${values[4]:-100}" ;;  # WORST value often more reliable
                232) _data[VENDOR_HEALTH]="${values[4]:-100}" ;;  # Available Reserved Space
                177) _data[VENDOR_HEALTH]="${values[4]:-100}" ;;  # Wear Leveling Count
            esac
        fi
    done <<< "$smart_attr"
    
    # Calculate health score using universal function
    _data[HEALTH_SCORE]=$(calculate_health_score \
        "${_data[TYPE]}" \
        "${_data[REALLOCATED_SECTORS]}" \
        "${_data[VENDOR_HEALTH]}" \
        "" \
        "")
    
    [ -n "${_data[MODEL]}" ]
}

is_nvme() {
    [[ "$(basename "$1")" =~ ^nvme[0-9]+ ]]
}

get_health_rating() {
    local score=$1
    if [ "$score" -ge 90 ]; then echo "Excellent"
    elif [ "$score" -ge 70 ]; then echo "Good"
    elif [ "$score" -ge 50 ]; then echo "Fair"
    elif [ "$score" -ge 25 ]; then echo "Poor"
    else echo "Critical"
    fi
}

get_health_color() {
    local score=$1
    if [ "$score" -ge 90 ]; then echo "${GREEN}"
    elif [ "$score" -ge 70 ]; then echo "${GREEN}"
    elif [ "$score" -ge 50 ]; then echo "${YELLOW}"
    elif [ "$score" -ge 25 ]; then echo "${RED}"
    else echo "${RED}"
    fi
}

# Output functions
print_info() {
    local device=$1
    local -n _data=$2
    
    echo -e "${BOLD}Device:${NC} ${LIGHT_CYAN}$device${NC}"
    echo -e "${BOLD}Model:${NC} ${LIGHT_CYAN}${_data[MODEL]:-N/A}${NC}"
    echo -e "${BOLD}Serial:${NC} ${LIGHT_CYAN}${_data[SERIAL]:-N/A}${NC}"
    echo -e "${BOLD}Firmware:${NC} ${LIGHT_CYAN}${_data[FIRMWARE]:-N/A}${NC}"
    
    if [ -n "${_data[CAPACITY_RAW]}" ]; then
        echo -e "${BOLD}Capacity:${NC} ${LIGHT_CYAN}${_data[CAPACITY_RAW]}${NC}"
    else
        echo -e "${BOLD}Capacity:${NC} ${LIGHT_CYAN}${_data[CAPACITY]:-N/A} GB${NC}"
    fi
    
    echo -e "${BOLD}Type:${NC} ${LIGHT_CYAN}${_data[TYPE]:-N/A}${NC}"
    echo -e "${BOLD}SMART:${NC} ${_data[SMART_STATUS]:-N/A}"
    
    local temp_color=$(get_temp_color "${_data[TEMP]}")
    echo -e "${BOLD}Temperature:${NC} ${temp_color}${_data[TEMP]:-N/A}°C${NC}"
    echo -e "${BOLD}Power On Hours:${NC} ${_data[POWER_ON_HOURS]:-0} h"
    echo -e "${BOLD}Power Cycles:${NC} ${_data[POWER_CYCLES]:-0}"
    
    if [ "${_data[TYPE]}" = "NVMe SSD" ]; then
        echo -e "${BOLD}Percentage Used:${NC} ${_data[PERCENTAGE_USED]:-0}%"
        echo -e "${BOLD}Available Spare:${NC} ${_data[AVAIL_SPARE]:-100}%"
        echo -e "${BOLD}Media Errors:${NC} ${_data[MEDIA_ERRORS]:-0}"
    else
        local realloc_color="${GREEN}"
        if [ "${_data[REALLOCATED_SECTORS]}" -gt 0 ]; then
            if [ "${_data[REALLOCATED_SECTORS]}" -le 10 ]; then
                realloc_color="${YELLOW}"
            else
                realloc_color="${RED}"
            fi
        fi
        echo -e "${BOLD}Reallocated Sectors:${NC} ${realloc_color}${_data[REALLOCATED_SECTORS]:-0}${NC}"
        
        local pending_color="${GREEN}"
        [ "${_data[PENDING_SECTORS]}" -gt 0 ] && pending_color="${YELLOW}"
        [ "${_data[PENDING_SECTORS]}" -gt 10 ] && pending_color="${RED}"
        echo -e "${BOLD}Pending Sectors:${NC} ${pending_color}${_data[PENDING_SECTORS]:-0}${NC}"
        
        local crc_color="${GREEN}"
        [ "${_data[CRC_ERRORS]}" -gt 0 ] && crc_color="${YELLOW}"
        [ "${_data[CRC_ERRORS]}" -gt 100 ] && crc_color="${RED}"
        echo -e "${BOLD}CRC Errors:${NC} ${crc_color}${_data[CRC_ERRORS]:-0}${NC}"
        
        if [ -n "${_data[VENDOR_HEALTH]}" ]; then
            echo -e "${BOLD}Vendor Health:${NC} ${_data[VENDOR_HEALTH]}%"
        fi
    fi
    
    local health_score=${_data[HEALTH_SCORE]:-100}
    local health_color=$(get_health_color $health_score)
    local health_rating=$(get_health_rating $health_score)
    echo ""
    echo -e "${LIGHT_PURPLE}STATE:${NC} ${health_color}${health_rating}${NC}"
    echo -e "${LIGHT_PURPLE}HEALTH:${NC} ${health_color}${health_score}%${NC}"
}

print_state() {
    local -n _data=$1
    local score=${_data[HEALTH_SCORE]:-100}
    local color=$(get_health_color $score)
    echo -e "${color}$(get_health_rating $score)${NC}"
}

print_worst() {
    local device=$1
    local -n _data=$2
    
    local worst=""
    local value=""
    local comment=""
    local color="${YELLOW}"
    
    if [ "${_data[TYPE]}" = "NVMe SSD" ]; then
        if [ -n "${_data[PERCENTAGE_USED]}" ] && [ "${_data[PERCENTAGE_USED]}" -ge 70 ]; then
            worst="Percentage Used"
            value="${_data[PERCENTAGE_USED]}%"
            comment="Drive is aging"
            [ "${_data[PERCENTAGE_USED]}" -ge 90 ] && color="${RED}"
        fi
        if [ -n "${_data[MEDIA_ERRORS]}" ] && [ "${_data[MEDIA_ERRORS]}" -gt 0 ]; then
            worst="Media Errors"
            value="${_data[MEDIA_ERRORS]}"
            comment="Media errors detected"
            [ "${_data[MEDIA_ERRORS]}" -gt 100 ] && color="${RED}"
        fi
    else
        if [ -n "${_data[REALLOCATED_SECTORS]}" ] && [ "${_data[REALLOCATED_SECTORS]}" -gt 0 ]; then
            worst="Reallocated Sectors"
            value="${_data[REALLOCATED_SECTORS]}"
            if [ "${_data[REALLOCATED_SECTORS]}" -gt 100 ]; then
                comment="CRITICAL: Replace drive"
                color="${RED}"
            elif [ "${_data[REALLOCATED_SECTORS]}" -gt 10 ]; then
                comment="Drive degrading"
                color="${YELLOW}"
            else
                comment="Early signs of degradation"
                color="${YELLOW}"
            fi
        fi
    fi
    
    if [ -n "${_data[TEMP]}" ] && [ "${_data[TEMP]}" -ge 70 ]; then
        worst="Temperature"
        value="${_data[TEMP]}°C"
        comment="Critical temperature"
        color="${RED}"
    elif [ -n "${_data[TEMP]}" ] && [ "${_data[TEMP]}" -ge 60 ]; then
        worst="Temperature"
        value="${_data[TEMP]}°C"
        comment="High temperature"
        color="${YELLOW}"
    fi
    
    if [ -n "$worst" ]; then
        echo -e "${BOLD}${LIGHT_CYAN}${device}${NC}: $worst = ${value}"
        echo -e "  → ${color}${comment}${NC}"
    else
        echo -e "${GREEN}${LIGHT_CYAN}${device}${NC}: No critical issues detected${NC}"
    fi
}

draw_scale() {
    local value=$1
    local width=$2
    local filled=$((width * value / 100))
    local empty=$((width - filled))
    
    if [ "$value" -ge 90 ]; then echo -ne "${GREEN}"
    elif [ "$value" -ge 70 ]; then echo -ne "${GREEN}"
    elif [ "$value" -ge 50 ]; then echo -ne "${YELLOW}"
    else echo -ne "${RED}"
    fi
    
    printf '%*s' $filled | tr ' ' '█'
    echo -ne "${NC}"
    printf '%*s' $empty | tr ' ' '░'
}

print_metrics() {
    local -n _data=$1
    
    echo -e "${BOLD}Health:${NC}"
    echo -n "  "
    draw_scale ${_data[HEALTH_SCORE]:-100} 40
    echo -e " ${_data[HEALTH_SCORE]:-100}%"
    
    if [ "${_data[TYPE]}" = "NVMe SSD" ] && [ -n "${_data[PERCENTAGE_USED]}" ]; then
        echo -e "\n${BOLD}Usage:${NC}"
        echo -n "  "
        draw_scale $((100 - ${_data[PERCENTAGE_USED]})) 40
        echo -e " ${_data[PERCENTAGE_USED]}% used"
    fi
    
    if [ -n "${_data[TEMP]}" ]; then
        echo -e "\n${BOLD}Temperature:${NC}"
        local temp_score=100
        [ "${_data[TEMP]}" -gt 60 ] && temp_score=$((100 - (${_data[TEMP]} - 60) * 5))
        [ $temp_score -lt 0 ] && temp_score=0
        echo -n "  "
        draw_scale $temp_score 40
        echo -e " ${_data[TEMP]}°C"
    fi
}

print_table() {
    local -n _devices=$1
    local -n _data=$2
    
    # Top border
    echo -e "${BOLD}+-----------------------------------------------------------------------------------------------------------------------------------+${NC}"
    
    # Header
    echo -e "${BOLD}| DISK HEALTH SUMMARY                                                                                                               |${NC}"
    
    # Separator
    echo -e "${BOLD}|-----------------------------------------------------------------------------------------------------------------------------------|${NC}"
    
    # Column headers
    echo -e "${BOLD}| DEVICE     | MODEL                     | SERIAL               | CAPACITY | TYPE     | HEALTH | STATUS         | TEMP   | HOURS    |${NC}"
    
    # Separator
    echo -e "${BOLD}|-----------------------------------------------------------------------------------------------------------------------------------|${NC}"
    
    # Data
    for device in "${_devices[@]}"; do
        local dev_name=$(basename "$device")
        local health="${_data[${device}_HEALTH_SCORE]:-100}"
        local health_color=$(get_health_color $health)
        local rating=$(get_health_rating $health)
        local model="${_data[${device}_MODEL]:-N/A}"
        local serial="${_data[${device}_SERIAL]:-N/A}"
        local capacity="${_data[${device}_CAPACITY_RAW]:-N/A}"
        local type="${_data[${device}_TYPE]:-N/A}"
        local temp="${_data[${device}_TEMP]:-N/A}"
        local temp_color=$(get_temp_color "$temp")
        local hours="${_data[${device}_POWER_ON_HOURS]:-0}"
        
        # Trim model if too long
        if [ ${#model} -gt 25 ]; then
            model="${model:0:22}..."
        fi
        
        # Trim serial if too long
        if [ ${#serial} -gt 20 ]; then
            serial="${serial:0:17}..."
        fi
        
        # Output data row
        echo -ne "${BOLD}|${NC} "
        echo -ne "${LIGHT_CYAN}"
        printf "%-10s" "$dev_name"
        echo -ne "${NC} | "
        echo -ne "${LIGHT_CYAN}"
        printf "%-25s" "$model"
        echo -ne "${NC} | "
        echo -ne "${LIGHT_CYAN}"
        printf "%-20s" "$serial"
        echo -ne "${NC} | "
        printf "%-8s" "$capacity"
        echo -n " | "
        printf "%-8s" "$type"
        echo -n " | "
        echo -ne "${health_color}"
        printf "%-6s" "${health}%"
        echo -ne "${NC} | "
        echo -ne "${health_color}"
        printf "%-14s" "$rating"
        echo -ne "${NC} | "
        echo -ne "${temp_color}"
        printf "%-7s" "${temp}°C"
        echo -ne "${NC} | "
        printf "%-8s" "$hours"
        echo -e " ${BOLD}|${NC}"
    done
    
    # Bottom border
    echo -e "${BOLD}+-----------------------------------------------------------------------------------------------------------------------------------+${NC}"
}

show_help() {
    echo -e "${BOLD}smartshot v${VERSION} - SMART disk health checker${NC}"
    echo ""
    echo -e "${BOLD}Usage:${NC}"
    echo "  smartshot [options]"
    echo "  smartshot <device> [options]"
    echo "  smartshot <device1> <device2> ... [options]"
    echo ""
    echo -e "${BOLD}Options:${NC}"
    echo "  -i, --info      Show device information"
    echo "  -s, --state     Show health state only"
    echo "  -p, --percent   Show health percentage only"
    echo "  -w, --worst     Show the most critical metric"
    echo "  -t, --table     Show summary table for multiple devices"
    echo "  -h, --help      Show this help message"
    echo ""
    echo -e "${BOLD}Examples:${NC}"
    echo "  smartshot                # Show table for all drives"
    echo "  smartshot -i             # Show info for all drives"
    echo "  smartshot -s             # Show state for all drives"
    echo "  smartshot -p             # Show health percentage for all drives"
    echo "  smartshot -w             # Show worst metric for all drives"
    echo "  smartshot nvme0          # Basic info + state + percent for nvme0"
    echo "  smartshot sda            # Basic info for SATA drive"
    echo "  smartshot nvme0 -w       # Show worst metric for nvme0"
    echo "  smartshot sda sdb -t     # Table for specific drives"
    echo "  smartshot nvme0 -isp     # Combined flags (info + state + percent)"
}

# MAIN
check_smartctl

# Parse arguments
DEVICES=()
OPT_INFO=false
OPT_STATE=false
OPT_PERCENT=false
OPT_WORST=false
OPT_TABLE=false
OPT_HELP=false

# If no arguments - show table for all drives
if [ $# -eq 0 ]; then
    OPT_TABLE=true
fi

while [ $# -gt 0 ]; do
    case "$1" in
        -i|--info) OPT_INFO=true; shift ;;
        -s|--state) OPT_STATE=true; shift ;;
        -p|--percent) OPT_PERCENT=true; shift ;;
        -w|--worst) OPT_WORST=true; shift ;;
        -t|--table) OPT_TABLE=true; shift ;;
        -h|--help) OPT_HELP=true; shift ;;
        -*)
            if [[ "$1" =~ ^-[iswthp]+$ ]]; then
                flags="${1#-}"
                for ((i=0; i<${#flags}; i++)); do
                    case "${flags:$i:1}" in
                        i) OPT_INFO=true ;;
                        s) OPT_STATE=true ;;
                        p) OPT_PERCENT=true ;;
                        w) OPT_WORST=true ;;
                        t) OPT_TABLE=true ;;
                        h) OPT_HELP=true ;;
                    esac
                done
                shift
            else
                echo -e "${RED}Unknown option: $1${NC}" >&2
                exit 1
            fi
            ;;
        *)
            DEVICES+=("$1")
            shift
            ;;
    esac
done

[ "$OPT_HELP" = true ] && { show_help; exit 0; }

# If no options and not table - show basic information
if ! $OPT_INFO && ! $OPT_STATE && ! $OPT_PERCENT && ! $OPT_WORST && ! $OPT_TABLE; then
    OPT_INFO=true
    OPT_STATE=true
    OPT_PERCENT=true
fi

# If no devices specified - get all disks via smartctl --scan
if [ ${#DEVICES[@]} -eq 0 ]; then
    mapfile -t DEVICES < <(get_all_disks)
fi

declare -A DISK_DATA
VALID_DEVICES=()

# Quiet mode for -s and -p (only when specific devices are specified and no other options)
QUIET_MODE=false
if [ "$OPT_STATE" = true ] && [ "$OPT_INFO" = false ] && [ "$OPT_WORST" = false ] && [ "$OPT_TABLE" = false ] && [ ${#DEVICES[@]} -eq 1 ]; then
    QUIET_MODE=true
fi
if [ "$OPT_PERCENT" = true ] && [ "$OPT_INFO" = false ] && [ "$OPT_WORST" = false ] && [ "$OPT_TABLE" = false ] && [ ${#DEVICES[@]} -eq 1 ]; then
    QUIET_MODE=true
fi

for dev in "${DEVICES[@]}"; do
    dev_name=$(basename "$dev")
    dev_path="/dev/$dev_name"
    
    # Show "Reading..." only if not quiet mode and not table
    if [ "$QUIET_MODE" = false ] && [ "$OPT_TABLE" = false ]; then
        echo -e "${BLUE}Reading $dev_path...${NC}" >&2
    fi
    
    if [ ! -e "$dev_path" ]; then
        [ "$QUIET_MODE" = false ] && [ "$OPT_TABLE" = false ] && echo -e "${RED}Error: Device $dev_path does not exist${NC}" >&2
        continue
    fi
    
    declare -A device_data
    
    if is_nvme "$dev_path"; then
        if get_nvme_data "$dev_path" device_data; then
            VALID_DEVICES+=("$dev_path")
            for key in "${!device_data[@]}"; do
                DISK_DATA["${dev_path}_${key}"]="${device_data[$key]}"
            done
            [ "$QUIET_MODE" = false ] && [ "$OPT_TABLE" = false ] && echo -e "${GREEN}✓ ${device_data[TYPE]} detected${NC}" >&2
        else
            [ "$QUIET_MODE" = false ] && [ "$OPT_TABLE" = false ] && echo -e "${RED}✗ Failed to read NVMe SMART data${NC}" >&2
        fi
    else
        if get_ata_data "$dev_path" device_data; then
            VALID_DEVICES+=("$dev_path")
            for key in "${!device_data[@]}"; do
                DISK_DATA["${dev_path}_${key}"]="${device_data[$key]}"
            done
            [ "$QUIET_MODE" = false ] && [ "$OPT_TABLE" = false ] && echo -e "${GREEN}✓ ${device_data[TYPE]} detected${NC}" >&2
        else
            [ "$QUIET_MODE" = false ] && [ "$OPT_TABLE" = false ] && echo -e "${RED}✗ Failed to read ATA SMART data${NC}" >&2
        fi
    fi
    [ "$QUIET_MODE" = false ] && [ "$OPT_TABLE" = false ] && echo >&2
done

[ ${#VALID_DEVICES[@]} -eq 0 ] && { 
    if [ "$OPT_TABLE" = true ]; then
        echo -e "${RED}No valid disks found${NC}" >&2
    else
        echo -e "${RED}Error: No valid devices found${NC}" >&2
    fi
    exit 1 
}

# Table
if [ "$OPT_TABLE" = true ]; then
    print_table VALID_DEVICES DISK_DATA
    exit 0
fi

# Output for each device with separator
FIRST_DEVICE=true
for device in "${VALID_DEVICES[@]}"; do
    declare -A current_data
    for key in MODEL SERIAL FIRMWARE CAPACITY CAPACITY_RAW TYPE SMART_STATUS TEMP \
               PERCENTAGE_USED POWER_ON_HOURS POWER_CYCLES HEALTH_SCORE \
               AVAIL_SPARE MEDIA_ERRORS REALLOCATED_SECTORS PENDING_SECTORS \
               CRC_ERRORS VENDOR_HEALTH; do
        current_data[$key]="${DISK_DATA[${device}_${key}]:-}"
    done
    
    dev_name=$(basename "$device")
    
    # Separator between devices
    if [ "$FIRST_DEVICE" = false ]; then
        if [ "$OPT_INFO" = true ] || [ "$OPT_WORST" = true ]; then
            echo -e "${BOLD}────────────────────────────────────────────────${NC}"
            echo ""
        fi
    fi
    
    if [ "$OPT_INFO" = true ]; then
        print_info "$dev_name" current_data
    fi
    
    if [ "$OPT_STATE" = true ] && [ "$OPT_INFO" = false ]; then
        print_state current_data
    fi
    
    if [ "$OPT_PERCENT" = true ] && [ "$OPT_INFO" = false ]; then
        echo "${current_data[HEALTH_SCORE]:-100}"
    fi
    
    if [ "$OPT_WORST" = true ]; then
        [ "$OPT_INFO" = true ] && echo ""
        print_worst "$dev_name" current_data
    fi
    
    FIRST_DEVICE=false
done