#!/bin/bash

#######################################
# SMART Disk Health Monitor
# Usage: ./script.sh [-a|-s|-h] [--auto-install]
#######################################

# Color codes
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Flags
AUTO_INSTALL=false

# Function to print colored output
print_color() {
    local color=$1
    shift
    echo -e "${color}$@${NC}"
}

# Function to check if command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Function to detect package manager
detect_package_manager() {
    if command_exists apt-get; then
        echo "apt"
    elif command_exists yum; then
        echo "yum"
    elif command_exists dnf; then
        echo "dnf"
    elif command_exists zypper; then
        echo "zypper"
    elif command_exists pacman; then
        echo "pacman"
    elif command_exists apk; then
        echo "apk"
    else
        echo "unknown"
    fi
}

# Function to install smartmontools
install_smartmontools() {
    local pkg_mgr=$(detect_package_manager)
    
    print_color "$YELLOW" "Attempting to install smartmontools..."
    
    case $pkg_mgr in
        apt)
            if [ "$EUID" -eq 0 ]; then
                apt-get update -qq && apt-get install -y smartmontools
            else
                sudo apt-get update -qq && sudo apt-get install -y smartmontools
            fi
            ;;
        yum)
            if [ "$EUID" -eq 0 ]; then
                yum install -y smartmontools
            else
                sudo yum install -y smartmontools
            fi
            ;;
        dnf)
            if [ "$EUID" -eq 0 ]; then
                dnf install -y smartmontools
            else
                sudo dnf install -y smartmontools
            fi
            ;;
        zypper)
            if [ "$EUID" -eq 0 ]; then
                zypper install -y smartmontools
            else
                sudo zypper install -y smartmontools
            fi
            ;;
        pacman)
            if [ "$EUID" -eq 0 ]; then
                pacman -Sy --noconfirm smartmontools
            else
                sudo pacman -Sy --noconfirm smartmontools
            fi
            ;;
        apk)
            if [ "$EUID" -eq 0 ]; then
                apk add --no-cache smartmontools
            else
                sudo apk add --no-cache smartmontools
            fi
            ;;
        *)
            print_color "$RED" "Unknown package manager. Cannot auto-install."
            return 1
            ;;
    esac
    
    return $?
}

# Function to prompt for installation
prompt_install() {
    local pkg_mgr=$(detect_package_manager)
    
    if [ "$pkg_mgr" == "unknown" ]; then
        print_color "$RED" "ERROR: Unable to detect package manager"
        print_color "$YELLOW" "Please install smartmontools manually:"
        echo "  Debian/Ubuntu: apt-get install smartmontools"
        echo "  RHEL/CentOS:   yum install smartmontools"
        echo "  Fedora:        dnf install smartmontools"
        echo "  Arch:          pacman -S smartmontools"
        echo "  Alpine:        apk add smartmontools"
        exit 1
    fi
    
    print_color "$YELLOW" "smartmontools is not installed."
    
    if [ "$AUTO_INSTALL" = true ]; then
        print_color "$BLUE" "Auto-install enabled. Installing smartmontools..."
        if install_smartmontools; then
            print_color "$GREEN" "✓ Successfully installed smartmontools"
            return 0
        else
            print_color "$RED" "✗ Failed to install smartmontools"
            exit 1
        fi
    else
        print_color "$BLUE" "Detected package manager: $pkg_mgr"
        echo ""
        read -p "Install smartmontools now? [y/N]: " -n 1 -r
        echo ""
        
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            if install_smartmontools; then
                print_color "$GREEN" "✓ Successfully installed smartmontools"
                return 0
            else
                print_color "$RED" "✗ Failed to install smartmontools"
                exit 1
            fi
        else
            print_color "$YELLOW" "Installation cancelled. Exiting."
            exit 1
        fi
    fi
}

# Function to check required dependencies
check_dependencies() {
    local missing_deps=()
    local core_missing=()
    
    # Check core utilities (should be present on all systems)
    for cmd in lsblk awk grep; do
        if ! command_exists "$cmd"; then
            core_missing+=("$cmd")
        fi
    done
    
    # If core utilities are missing, fail immediately
    if [ ${#core_missing[@]} -ne 0 ]; then
        print_color "$RED" "ERROR: Missing core system utilities: ${core_missing[*]}"
        print_color "$YELLOW" "Install with your package manager:"
        echo "  util-linux (for lsblk)"
        echo "  gawk or mawk (for awk)"
        echo "  grep (usually pre-installed)"
        exit 1
    fi
    
    # Check if smartctl is available
    if ! command_exists smartctl; then
        prompt_install
    fi
}

# Function to check root privileges
check_privileges() {
    if [ "$EUID" -ne 0 ] && ! sudo -n true 2>/dev/null; then
        print_color "$RED" "ERROR: This script requires root privileges"
        print_color "$YELLOW" "Run with: sudo $0 $*"
        exit 1
    fi
}

# Function to safely get smartctl data
get_smart_data() {
    local drive=$1
    local flags=$2
    local output
    
    output=$(smartctl $flags "/dev/$drive" 2>/dev/null)
    local exit_code=$?
    
    # Exit codes: Bit 1 (2): Device open failed or doesn't support SMART
    if [ $exit_code -eq 2 ] || [ $exit_code -eq 1 ]; then
        return 1
    fi
    
    echo "$output"
    return 0
}

# Function to display usage
show_usage() {
    cat << EOF
Usage: $(basename "$0") [OPTIONS]

Monitor disk health using SMART data.

OPTIONS:
    -a                Full SMART output for all drives
    -s                Short health summary for all drives
    -h                Display this help message
    --auto-install    Auto-install smartmontools without prompting
    (default)         Table summary of all drives

EXAMPLES:
    $(basename "$0")                    # Table view (default)
    $(basename "$0") -a                 # Full SMART data
    $(basename "$0") -s                 # Quick health check
    $(basename "$0") --auto-install     # Install dependencies automatically

REQUIREMENTS:
    - Root/sudo privileges
    - smartmontools package (will offer to install if missing)
    - util-linux package

SECURITY NOTE:
    Use --auto-install only on trusted systems. Interactive install
    is recommended for production environments.

EOF
    exit 0
}

# Parse options
mode="table"
for arg in "$@"; do
    case $arg in
        -a) mode="full" ;;
        -s) mode="short" ;;
        -h) show_usage ;;
        --auto-install) AUTO_INSTALL=true ;;
        -*)
            print_color "$RED" "Invalid option: $arg"
            show_usage
            ;;
    esac
done

# Run dependency and privilege checks
check_dependencies
check_privileges

# Auto-detect all disk devices
mapfile -t drives < <(lsblk -dn -o NAME,TYPE 2>/dev/null | awk '$2=="disk"{print $1}')

if [ ${#drives[@]} -eq 0 ]; then
    print_color "$RED" "ERROR: No disk drives detected"
    exit 1
fi

# Full mode
if [[ "$mode" == "full" ]]; then
    for drive in "${drives[@]}"; do
        echo ""
        print_color "$BLUE" "=== /dev/$drive ==="
        smart_data=$(get_smart_data "$drive" "-a")
        if [ $? -eq 0 ]; then
            echo "$smart_data"
        else
            print_color "$RED" "  SMART not supported or device unavailable"
        fi
        echo ""
    done

# Short mode
elif [[ "$mode" == "short" ]]; then
    for drive in "${drives[@]}"; do
        print_color "$BLUE" "Drive: /dev/$drive"
        smart_data=$(get_smart_data "$drive" "-H")
        if [ $? -eq 0 ]; then
            health=$(echo "$smart_data" | grep -i "SMART overall-health" | awk -F: '{print $2}' | xargs)
            if [[ "$health" == "PASSED" ]]; then
                print_color "$GREEN" "  Health: $health"
            else
                print_color "$RED" "  Health: ${health:-UNKNOWN}"
            fi
        else
            print_color "$RED" "  SMART not supported or device unavailable"
        fi
        echo ""
    done

# Table mode
else
    printf "| %-12s | %-22s | %-13s | %-9s | %-18s | %-15s |\n" "Drive" "Model" "Power-On Hrs" "Temp (°C)" "Realloc/MediaErr" "Health"
    echo "|--------------|------------------------|---------------|-----------|--------------------|--------------------|"
    
    for drive in "${drives[@]}"; do
        model="N/A"
        power="N/A"
        temp="N/A"
        errors="N/A"
        health="UNKNOWN"
        
        info=$(get_smart_data "$drive" "-i")
        if [ $? -ne 0 ]; then
            printf "| %-12s | %-22s | %-13s | %-9s | %-18s | %-15s |\n" "/dev/$drive" "SMART N/A" "-" "-" "-" "N/A"
            continue
        fi
        
        healthinfo=$(get_smart_data "$drive" "-A -H")
        
        if [[ "$drive" == nvme* ]]; then
            model=$(echo "$info" | awk -F': ' '/Model Number/{gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}')
            power=$(echo "$healthinfo" | awk -F': ' '/Power On Hours/{gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}')
            temp=$(echo "$healthinfo" | awk -F': ' '/^Temperature:/{gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}' | awk '{print $1}')
            media_err=$(echo "$healthinfo" | awk -F': ' '/Media and Data Integrity Errors/{gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}')
            health=$(echo "$healthinfo" | grep -i "SMART overall-health self-assessment test result" | awk -F': ' '{gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}')
            errors="${media_err:-0}"
        else
            model=$(echo "$info" | awk -F': ' '/Device Model:|Product:/{gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}' | head -n1)
            power=$(echo "$healthinfo" | awk '/Power_On_Hours/{print $10}' | head -n1)
            temp=$(echo "$healthinfo" | awk '/Temperature_Celsius|Airflow_Temperature_Cel/{print $10}' | head -n1)
            realloc=$(echo "$healthinfo" | awk '/Reallocated_Sector_Ct/{print $10}' | head -n1)
            health=$(echo "$healthinfo" | grep -i "SMART overall-health self-assessment test result" | awk -F': ' '{gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}')
            errors="${realloc:-0}"
        fi
        
        if [ ${#model} -gt 22 ]; then
            model="${model:0:19}..."
        fi
        
        model="${model:-N/A}"
        power="${power:-N/A}"
        temp="${temp:-N/A}"
        errors="${errors:-N/A}"
        health="${health:-UNKNOWN}"
        
        printf "| %-12s | %-22s | %-13s | %-9s | %-18s | %-15s |\n" \
            "/dev/$drive" "$model" "$power" "$temp" "$errors" "$health"
    done
    
    echo ""
    print_color "$YELLOW" "Note: PASSED = Healthy, FAILED = Replace immediately, N/A = SMART unavailable"
fi

exit 0