#!/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[@]}")
}

# Check for RAID arrays and stop them
check_and_stop_raid() {
    log "Checking for active RAID arrays..."
    
    if [[ ! -f /proc/mdstat ]]; then
        log "No RAID support detected"
        return 0
    fi
    
    # Get list of active RAID devices
    local raid_devices=()
    while IFS= read -r line; do
        if [[ "$line" =~ ^(md[0-9]+) ]]; then
            raid_devices+=("/dev/${BASH_REMATCH[1]}")
        fi
    done < /proc/mdstat
    
    if [[ ${#raid_devices[@]} -eq 0 ]]; then
        log "No active RAID arrays found"
        return 0
    fi
    
    log "Found active RAID arrays:"
    for raid in "${raid_devices[@]}"; do
        echo "  - $raid"
    done
    
    # Check if any RAID devices are mounted
    log "Checking for mounted RAID filesystems..."
    for raid in "${raid_devices[@]}"; do
        if mountpoint -q "$raid" 2>/dev/null; then
            log "Unmounting $raid"
            umount "$raid" 2>/dev/null || warn "Failed to unmount $raid"
        fi
    done
    
    # Stop RAID arrays
    log "Stopping RAID arrays..."
    for raid in "${raid_devices[@]}"; do
        log "Stopping $raid"
        if mdadm --stop "$raid" 2>/dev/null; then
            log "Successfully stopped $raid"
        else
            warn "Failed to stop $raid"
        fi
    done
    
    # Wait a moment for arrays to fully stop
    sleep 2
    
    # Verify arrays are stopped
    local remaining_arrays=()
    while IFS= read -r line; do
        if [[ "$line" =~ ^(md[0-9]+) ]]; then
            remaining_arrays+=("/dev/${BASH_REMATCH[1]}")
        fi
    done < /proc/mdstat
    
    if [[ ${#remaining_arrays[@]} -gt 0 ]]; then
        warn "Some RAID arrays are still active:"
        for raid in "${remaining_arrays[@]}"; do
            echo "  - $raid"
        done
        return 1
    else
        log "All RAID arrays successfully stopped"
        return 0
    fi
}

# Remove RAID superblocks from disk partitions
remove_raid_superblocks() {
    local disk="$1"
    
    log "Checking for RAID superblocks on $disk..."
    
    # Get all partitions on this disk
    local partitions=()
    while IFS= read -r partition; do
        if [[ -n "$partition" ]]; then
            partitions+=("/dev/$partition")
        fi
    done < <(lsblk -ln -o NAME "$disk" | tail -n +2)
    
    # Check each partition for RAID superblocks
    for partition in "${partitions[@]}"; do
        if mdadm --examine "$partition" &>/dev/null; then
            log "Found RAID superblock on $partition, removing..."
            if mdadm --zero-superblock "$partition" 2>/dev/null; then
                log "Successfully removed RAID superblock from $partition"
            else
                warn "Failed to remove RAID superblock from $partition"
            fi
        fi
    done
    
    # Also check the whole disk for superblocks
    if mdadm --examine "$disk" &>/dev/null; then
        log "Found RAID superblock on $disk, removing..."
        if mdadm --zero-superblock "$disk" 2>/dev/null; then
            log "Successfully removed RAID superblock from $disk"
        else
            warn "Failed to remove RAID superblock from $disk"
        fi
    fi
}

# 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..."
    
    # Check if disk is part of any active RAID arrays and stop them if needed
    local disk_basename=$(basename "$disk")
    if grep -q "$disk_basename" /proc/mdstat 2>/dev/null; then
        warn "Disk $disk appears to be part of an active RAID array"
        log "Stopping RAID arrays that use this disk..."
        check_and_stop_raid
    fi
    
    # 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
    
    # Remove RAID superblocks if present
    remove_raid_superblocks "$disk"
    
    # 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"
    
    # Wipe partition table first
    log "Wiping partition table on $disk..."
    wipefs -af "$disk" 2>/dev/null || warn "Failed to wipe filesystem signatures"
    
    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..."
    
    # First, stop any RAID arrays that might be using these disks
    log "Checking and stopping RAID arrays before disk erasure..."
    check_and_stop_raid
    
    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 "$@"