#!/bin/bash

#############################################################
#### Expand Drive                                        ####
#### Auto-expand partitions and filesystems              ####
####                                                     ####
#### Author: Phil Connor                                 ####
#### Contact: contact@mylinux.work                       ####
#### License: MIT                                        ####
#### Version: 2.4                                        ####
####                                                     ####
#### Usage: sudo ./expand-drive.sh                       ####
#############################################################
# v2.4 changes:
#   - Fixed: grep in pipeline crashes under set -euo pipefail when no matches found. Added || true guard
#############################################################

# Set strict error handling:
# -e: Exit immediately if a command exits with a non-zero status
# -u: Treat unset variables as an error when substituting
# -o pipefail: The return value of a pipeline is the status of the last command to exit with a non-zero status
set -euo pipefail

# Constants - Define paths to required system binaries (use command names, let PATH resolve)
readonly BLKID_PATH="blkid"           # Tool to locate/print block device attributes
readonly LSBLK_PATH="lsblk"            # Tool to list block devices
readonly LOG_FILE="/var/log/expand_drive.log"   # Location for script log output

# Configuration - Runtime behavior settings
readonly DRY_RUN=${DRY_RUN:-false}                                  # If true, show what would be done without making changes
readonly REQUIRED_COMMANDS=("growpart" "xfs_growfs" "resize2fs")    # Commands that must be available
readonly SUPPORTED_FILESYSTEMS=("xfs" "ext2" "ext3" "ext4")         # Filesystem types we can expand

# Exit codes - Standardized exit status values
readonly EXIT_SUCCESS=0          # Script completed successfully
readonly EXIT_ERROR=1            # General error occurred
readonly EXIT_ROOT_REQUIRED=2    # Script must be run as root user
readonly EXIT_MISSING_DEPS=3     # Required dependencies are missing

# Function to log messages with timestamp to both console and log file
log_message() {
    echo "$(date): $1" | tee -a "$LOG_FILE"
}

# Function to log error messages with timestamp to both console, log file, and stderr
log_error() {
    echo "$(date): ERROR: $1" | tee -a "$LOG_FILE" >&2
}

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

# Function to handle script interruption (SIGINT/SIGTERM) and perform cleanup
cleanup() {
    # shellcheck disable=SC2317  # Suppress warning about unreachable code
    log_message "Script interrupted, cleaning up..."
    # shellcheck disable=SC2317  # Suppress warning about unreachable code
    exit "$EXIT_ERROR"
}

# Function to validate prerequisites before script execution
validate_prerequisites() {
    # Check if script is run as root (required for partition/filesystem operations)
    if [ "$(id -u)" -ne 0 ]; then
        echo "Error: This script must be run as root"
        exit "$EXIT_ROOT_REQUIRED"
    fi

    # Ensure log directory exists and is writable
    local log_dir
    log_dir=$(dirname "$LOG_FILE")
    if [ ! -d "$log_dir" ]; then
        mkdir -p "$log_dir" || {
            echo "Error: Cannot create log directory $log_dir"
            exit "$EXIT_ERROR"
        }
    fi

    # Verify all required system commands are available
    for cmd in "${REQUIRED_COMMANDS[@]}"; do
        if ! command_exists "$cmd"; then
            log_error "Required command '$cmd' not found. Please install it."
            exit "$EXIT_MISSING_DEPS"
        fi
    done
}

# Function to check if filesystem type is supported by this script
is_supported_filesystem() {
    local fs_type="$1"
    # Loop through supported filesystem types array
    for supported in "${SUPPORTED_FILESYSTEMS[@]}"; do
        if [[ "$fs_type" == "$supported" ]]; then
            return 0  # Filesystem type is supported
        fi
    done
    return 1  # Filesystem type is not supported
}

# Function to expand filesystem based on type (XFS or EXT variants)
expand_filesystem() {
    local partition="$1"     # Block device path (e.g., /dev/sda1)
    local fs_type="$2"       # Filesystem type (xfs, ext2, ext3, ext4)
    local mount_point="$3"   # Where the filesystem is mounted

    # Validate filesystem type is one we support
    if ! is_supported_filesystem "$fs_type"; then
        log_error "Unsupported filesystem type $fs_type on $partition"
        return 1
    fi

    # Handle different filesystem types with appropriate expansion commands
    case $fs_type in
    "xfs")
        log_message "Expanding XFS filesystem on $partition"
        if [ "$DRY_RUN" = "true" ]; then
            log_message "DRY RUN: Would expand XFS filesystem on $partition"
            return 0
        # XFS uses xfs_growfs and requires the mount point as argument
        elif xfs_growfs "$mount_point" >/dev/null 2>&1; then
            log_message "Successfully expanded XFS filesystem on $partition"
            return 0
        else
            log_error "Failed to expand XFS filesystem on $partition"
            return 1
        fi
        ;;
    "ext2" | "ext3" | "ext4")
        log_message "Expanding EXT filesystem on $partition"
        if [ "$DRY_RUN" = "true" ]; then
            log_message "DRY RUN: Would expand EXT filesystem on $partition"
            return 0
        # EXT filesystems use resize2fs and require the device path as argument
        elif resize2fs "$partition" >/dev/null 2>&1; then
            log_message "Successfully expanded EXT filesystem on $partition"
            return 0
        else
            log_error "Failed to expand EXT filesystem on $partition"
            return 1
        fi
        ;;
    esac
}

# Function to expand partition to use available disk space
expand_partition() {
    local disk="$1"        # Parent disk device (e.g., /dev/sda)
    local partition="$2"   # Partition device (e.g., /dev/sda1)
    local part_num="$3"    # Partition number (e.g., 1)

    # Check if partition can be expanded using growpart dry-run
    if ! growpart "$disk" "$part_num" --dry-run 2>/dev/null; then
        log_message "Partition $partition doesn't need expansion or cannot be expanded, skipping..."
        return 1  # Not an error, just nothing to do
    fi

    # Perform the actual partition expansion
    if [ "$DRY_RUN" = "true" ]; then
        log_message "DRY RUN: Would expand partition $partition"
        return 0
    elif growpart "$disk" "$part_num" >/dev/null 2>&1; then
        log_message "Successfully expanded partition $partition"
        return 0
    else
        log_error "Failed to expand partition $partition"
        return 1
    fi
}

# Set up signal trap to handle interruptions gracefully
trap cleanup INT TERM

# Initialize script by validating prerequisites
validate_prerequisites

# Function to process a single partition (expand partition and filesystem)
process_partition() {
    local partition="$1"  # Partition device path (e.g., /dev/sda1)
    local disk="$2"       # Parent disk device path (e.g., /dev/sda)

    log_message "Processing partition $partition"

    # Check if the filesystem is currently mounted (required for filesystem expansion)
    local mount_point
    mount_point=$(findmnt -n -o TARGET "$partition" 2>/dev/null)
    if [ -z "$mount_point" ]; then
        log_message "Warning: $partition is not mounted, skipping filesystem resize"
        return 0
    fi

    # Extract partition number from device path (e.g., extract "1" from "/dev/sda1")
    local part_num
    part_num=$(echo "$partition" | { grep -o '[0-9]\+$' || true; } | tail -1)
    if [ -z "$part_num" ]; then
        log_error "Could not extract partition number from $partition"
        return 1
    fi

    # First expand the partition to use available disk space
    if ! expand_partition "$disk" "$partition" "$part_num"; then
        return 0 # Not an error if partition doesn't need expansion
    fi

    # Detect the filesystem type using blkid
    local fs_type
    fs_type=$($BLKID_PATH -s TYPE -o value "$partition")
    if [ -z "$fs_type" ]; then
        log_message "Warning: Could not detect filesystem type for $partition, skipping..."
        return 0
    fi

    # Get current filesystem size before expansion
    local current_size
    current_size=$(df -h "$mount_point" | awk 'NR==2 {print $2}')
    log_message "Current filesystem size on $partition: $current_size"
    
    # Expand the filesystem to use the newly available partition space
    expand_filesystem "$partition" "$fs_type" "$mount_point"
    
    # Show new size after expansion
    local new_size
    new_size=$(df -h "$mount_point" | awk 'NR==2 {print $2}')
    log_message "New filesystem size on $partition: $new_size"
}

# Function to process a disk with direct filesystem (no partitions)
process_direct_filesystem() {
    local disk="$1"         # Disk device path (e.g., /dev/nvme3n1)
    local mount_point="$2"  # Where the filesystem is mounted
    
    log_message "Processing direct filesystem on $disk mounted at $mount_point"
    
    # Detect the filesystem type using blkid
    local fs_type
    fs_type=$($BLKID_PATH -s TYPE -o value "$disk")
    if [ -z "$fs_type" ]; then
        log_message "Warning: Could not detect filesystem type for $disk, skipping..."
        return 0
    fi
    
    # Get current filesystem size before expansion
    local current_size
    current_size=$(df -h "$mount_point" | awk 'NR==2 {print $2}')
    log_message "Current filesystem size on $disk: $current_size"
    
    # Expand the filesystem to use the full disk space
    expand_filesystem "$disk" "$fs_type" "$mount_point"
    
    # Show new size after expansion
    local new_size
    new_size=$(df -h "$mount_point" | awk 'NR==2 {print $2}')
    log_message "New filesystem size on $disk: $new_size"
}

# Function to process all partitions on a single disk
process_disk() {
    local disk="$1"  # Disk device path (e.g., /dev/sda)

    log_message "Checking partitions on $disk..."

    # Get list of partitions for the current disk using lsblk
    # Filter for partition type and extract device names
    local partitions
    local lsblk_output
    lsblk_output=$($LSBLK_PATH -pln -o NAME,TYPE "$disk" 2>&1) || {
        log_error "lsblk command failed for $disk: $lsblk_output"
        return 1
    }
    partitions=$(echo "$lsblk_output" | grep "part" | cut -d' ' -f1 || true)
    
    if [ -z "$partitions" ]; then
        # Check if the disk itself has a filesystem (no partition table)
        local mount_point
        mount_point=$(findmnt -n -o TARGET "$disk" 2>/dev/null)
        if [ -n "$mount_point" ]; then
            log_message "No partitions found on $disk, but disk has direct filesystem. Processing disk directly..."
            process_direct_filesystem "$disk" "$mount_point"
        else
            log_message "No partitions found on $disk, skipping..."
        fi
        return 0
    fi

    # Process each partition found on this disk
    for partition in $partitions; do
        process_partition "$partition" "$disk"
    done
}

# Main execution function - orchestrates the entire drive expansion process
main() {
    log_message "Starting drive expansion process..."

    # Get list of all disk devices in the system using lsblk
    # Filter for disk type and extract device names
    local devices
    devices=$($LSBLK_PATH -pln -o NAME,TYPE | { grep "disk" || true; } | cut -d' ' -f1)

    # Verify we found at least one disk device
    if [ -z "$devices" ]; then
        log_error "No disk devices found"
        exit "$EXIT_ERROR"
    fi

    # Process each disk device found
    for disk in $devices; do
        # Verify device is actually a block device before processing
        if [ ! -b "$disk" ]; then
            log_error "Device $disk is not a block device, skipping..."
            continue
        fi
        process_disk "$disk"
    done

    log_message "Drive expansion completed"
    exit "$EXIT_SUCCESS"
}

# Execute the main function to start the script
main