#!/bin/bash -e

# erase.sh - Disk Erasure Script
# WARNING: This script will PERMANENTLY DESTROY ALL DATA on selected disks!
# Use with extreme caution and only on systems you intend to completely wipe.

set -euo pipefail  # Exit on error, undefined vars, pipe failures
IFS=$'\n\t'       # Secure Internal Field Separator

# Jump to script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Colors for output
RED='\033[0;31m'
YELLOW='\033[1;33m'
GREEN='\033[0;32m'
NC='\033[0m' # No Color

# Logging function
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1" >&2
}

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        error "This script must be run as root (use sudo)"
        exit 1
    fi
}

# List all SSD disks
list_ssd_disks() {
    log "Scanning for SSD disks..."
    
    # Get all block devices that are disks (not partitions)
    local disks=()
    
    # Use lsblk to find disks and check if they're SSDs
    while IFS= read -r line; do
        local device=$(echo "$line" | awk '{print $1}')
        local type=$(echo "$line" | awk '{print $2}')
        
        # Only process actual disks (not partitions)
        if [[ "$type" == "disk" ]]; then
            # Check if it's an SSD by looking at rotational status
            local rotational_file="/sys/block/${device}/queue/rotational"
            if [[ -f "$rotational_file" ]]; then
                local rotational=$(cat "$rotational_file" 2>/dev/null || echo "1")
                if [[ "$rotational" == "0" ]]; then
                    disks+=("/dev/$device")
                fi
            fi
        fi
    done < <(lsblk -dn -o NAME,TYPE 2>/dev/null)
    
    if [[ ${#disks[@]} -eq 0 ]]; then
        warn "No SSD disks found on this system"
        return 1
    fi
    
    echo
    log "Found SSD disks:"
    for i in "${!disks[@]}"; do
        local disk="${disks[$i]}"
        local size=$(lsblk -dn -o SIZE "$disk" 2>/dev/null || echo "Unknown")
        local model=$(lsblk -dn -o MODEL "$disk" 2>/dev/null || echo "Unknown")
        printf "%2d) %-12s %-10s %s\n" $((i+1)) "$disk" "$size" "$model"
    done
    echo
    
    # Store disks in global array for use by erase function
    AVAILABLE_DISKS=("${disks[@]}")
}

# Function to erase a single disk
erase_disk() {
    local disk="$1"
    
    if [[ ! -b "$disk" ]]; then
        error "Device $disk is not a valid block device"
        return 1
    fi
    
    log "Starting erasure of $disk..."
    
    # Unmount any mounted partitions on this disk
    log "Unmounting any mounted partitions on $disk..."
    for partition in $(lsblk -ln -o NAME "$disk" | tail -n +2); do
        local part_path="/dev/$partition"
        if mountpoint -q "$part_path" 2>/dev/null; then
            log "Unmounting $part_path"
            umount "$part_path" 2>/dev/null || warn "Failed to unmount $part_path"
        fi
    done
    
    # Get disk size for progress indication
    local disk_size_bytes=$(blockdev --getsize64 "$disk" 2>/dev/null || echo "0")
    local disk_size_gb=$((disk_size_bytes / 1024 / 1024 / 1024))
    
    log "Disk size: ${disk_size_gb}GB"
    log "Erasing first 1GB of $disk (this will destroy the partition table and filesystem headers)..."
    
    # Use dd to zero out the first 1GB of the disk
    # This is usually sufficient to make data unrecoverable for most purposes
    if dd if=/dev/zero of="$disk" bs=1M count=1024 status=progress 2>/dev/null; then
        log "Successfully erased first 1GB of $disk"
        
        # Also erase the last 1MB (where backup partition tables might be)
        log "Erasing backup partition table area at end of disk..."
        local last_mb_seek=$((disk_size_bytes / 1024 / 1024 - 1))
        if [[ $last_mb_seek -gt 0 ]]; then
            dd if=/dev/zero of="$disk" bs=1M count=1 seek="$last_mb_seek" 2>/dev/null || warn "Failed to erase end of disk"
        fi
        
        # Force kernel to re-read partition table
        partprobe "$disk" 2>/dev/null || true
        
        log "Disk $disk has been successfully erased"
    else
        error "Failed to erase $disk"
        return 1
    fi
}

# Function to erase all detected SSD disks
erase_all_disks() {
    if [[ ${#AVAILABLE_DISKS[@]} -eq 0 ]]; then
        error "No disks available for erasure"
        return 1
    fi
    
    echo
    error "WARNING: This will PERMANENTLY DESTROY ALL DATA on the following disks:"
    for disk in "${AVAILABLE_DISKS[@]}"; do
        echo "  - $disk"
    done
    echo
    
    # Multiple confirmations for safety
    warn "This action is IRREVERSIBLE and will make data recovery extremely difficult!"
    echo
    read -p "Type 'Y' to continue: " confirm1
    if [[ "$confirm1" != "Y" ]]; then
        log "Operation cancelled by user"
        return 0
    fi
    
    log "Starting disk erasure process..."
    
    local success_count=0
    local total_count=${#AVAILABLE_DISKS[@]}
    
    for disk in "${AVAILABLE_DISKS[@]}"; do
        echo
        log "Processing disk $disk..."
        if erase_disk "$disk"; then
            ((success_count++))
        else
            error "Failed to erase $disk"
        fi
    done
    
    echo
    log "Erasure complete: $success_count/$total_count disks processed successfully"
    
    if [[ $success_count -eq $total_count ]]; then
        log "All disks have been successfully erased"
    else
        warn "Some disks failed to erase completely"
    fi
}

# Main function
main() {
    log "Disk Erasure Script Starting..."
    log "Current directory: $(pwd)"
    
    # Safety checks
    check_root
    
    # List available SSD disks
    if ! list_ssd_disks; then
        exit 1
    fi
    
    # Interactive menu
    while true; do
        echo
        echo "Options:"
        echo "1) Erase ALL detected SSD disks"
        echo "2) Erase specific disk"
        echo "3) Refresh disk list"
        echo "4) Exit"
        echo
        read -p "Select option [1-4]: " choice
        
        case $choice in
            1)
                erase_all_disks
                ;;
            2)
                echo
                read -p "Enter disk number to erase [1-${#AVAILABLE_DISKS[@]}]: " disk_num
                if [[ "$disk_num" =~ ^[0-9]+$ ]] && [[ $disk_num -ge 1 ]] && [[ $disk_num -le ${#AVAILABLE_DISKS[@]} ]]; then
                    local selected_disk="${AVAILABLE_DISKS[$((disk_num-1))]}"
                    echo
                    error "WARNING: This will PERMANENTLY DESTROY ALL DATA on $selected_disk"
                    read -p "Type 'ERASE' to confirm: " confirm
                    if [[ "$confirm" == "ERASE" ]]; then
                        erase_disk "$selected_disk"
                    else
                        log "Operation cancelled"
                    fi
                else
                    error "Invalid disk number"
                fi
                ;;
            3)
                list_ssd_disks
                ;;
            4)
                log "Exiting..."
                exit 0
                ;;
            *)
                error "Invalid option"
                ;;
        esac
    done
}

# Global array to store available disks
declare -a AVAILABLE_DISKS

# Run main function
main "$@"