#!/bin/bash
set -e

clear
echo "###############################################"
echo "#     Dynamic Disk Expansion Tool (LVM+RAW)   #"
echo "###############################################"
echo

# Install required tools
echo "[INFO] Installing dependencies..."
sudo apt-get update -qq
sudo apt-get install -y parted lvm2 util-linux gdisk

# Detect root device (e.g. /dev/sda3 or /dev/mapper/...)
ROOT_DEV=$(findmnt -n -o SOURCE /)
echo "[INFO] Root device: $ROOT_DEV"

# Get base disk (e.g., /dev/sda)
BASE_DISK="/dev/$(lsblk -no pkname "$ROOT_DEV")"
echo "[INFO] Base disk: $BASE_DISK"

# Rescan disk (in case hypervisor resize happened)
echo "[INFO] Rescanning disk..."
echo 1 | sudo tee /sys/class/block/$(basename "$BASE_DISK")/device/rescan
sleep 3

# Fix GPT if needed
sudo sgdisk --move-second-header "$BASE_DISK"

# Find the last partition number on the disk
LAST_PART_NUM=$(lsblk -ln "$BASE_DISK" | awk '$6 == "part" {print $1}' | grep -o '[0-9]*$' | tail -n 1)
PARTITION="${BASE_DISK}${LAST_PART_NUM}"
echo "[INFO] Last partition on disk: $PARTITION"

# Resize the partition to fill the disk
echo "[INFO] Resizing partition $LAST_PART_NUM to 100%..."
sudo parted "$BASE_DISK" ---pretend-input-tty <<EOF
resizepart $LAST_PART_NUM 100%
Yes
EOF

# Re-read partition table
echo "[INFO] Updating partition table..."
sudo partprobe "$BASE_DISK"
sleep 3

# --- Main Detection Logic ---
if [[ "$ROOT_DEV" =~ /dev/mapper/ ]]; then
    echo "[INFO] Detected LVM setup."

    # Get underlying PV (e.g. /dev/sda3)
    PV_PART=$(sudo pvs --noheadings -o pv_name | head -n 1 | xargs)
    echo "[INFO] Physical volume: $PV_PART"

    sudo pvresize "$PV_PART"

    LV_PATH="$ROOT_DEV"
    FSTYPE=$(lsblk -no FSTYPE "$LV_PATH")
    echo "[INFO] Filesystem type: $FSTYPE"

    echo "[INFO] Expanding logical volume..."
    sudo lvextend -l +100%FREE "$LV_PATH"

    if [[ "$FSTYPE" == "ext4" ]]; then
        sudo resize2fs "$LV_PATH"
    elif [[ "$FSTYPE" == "xfs" ]]; then
        sudo xfs_growfs "$LV_PATH"
    else
        echo "[ERROR] Unsupported filesystem: $FSTYPE"
        exit 1
    fi

else
    echo "[INFO] Detected non-LVM setup."

    FSTYPE=$(lsblk -no FSTYPE "$ROOT_DEV")
    echo "[INFO] Filesystem type: $FSTYPE"

    if [[ "$FSTYPE" == "ext4" ]]; then
        echo "[INFO] Resizing ext4 filesystem..."
        sudo resize2fs "$ROOT_DEV"
    elif [[ "$FSTYPE" == "xfs" ]]; then
        echo "[INFO] Resizing XFS filesystem..."
        sudo xfs_growfs "$ROOT_DEV"
    else
        echo "[ERROR] Unsupported filesystem: $FSTYPE"
        exit 1
    fi
fi

echo
echo "✅ Disk and filesystem expansion completed."
df -h /
