#!/bin/bash
set -e

# root check
if [[ $EUID -ne 0 ]]; then
    echo "You must be root to do this." 1>&2
    exit 100
fi

mount_points=("/dev" "/dev/pts" "/proc" "/sys" "/sys/firmware/efi/efivars" "/run")
NO_WRITE=false

# arguments
for arg in "$@"; do
    case $arg in
        --nowrite)
            NO_WRITE=true
            shift
            ;;
        -h)
            printf "%-20s %s\n" "Flags:" "Description:"
            echo
            printf "%-20s %s\n" "--nowrite" "Disables Writing to Disk only modifys Existing Partitions"
            exit 1
            ;;
        *)
            echo "Unknown option: $arg"
            echo "Try -h"
            exit 1
            ;;
    esac
done


cleanup() {
    echo "Cleaning Mounts"
    
    umount -lf /mnt/ventoy || true

    for i in "${mount_points[@]}"; do umount -lf "/mnt/debian$i" || true; done

    umount -lf /mnt/debian/boot/efi || true
    umount -lf /mnt/debian/boot || true
    umount -lf /mnt/debian || true
}
trap cleanup EXIT

# Set your target disk here
udevadm trigger

lsblk
echo



read -rp "Enter the target disk (e.g., sda): " disk_suffix
disk_name="/dev/${disk_suffix}"

echo "Chose disk $disk_name"
echo

if [[ $NO_WRITE == true ]]; then
    echo "--nowrite is specified no need for Img no disks will be written to"
else
    read -rp "Enter the name of the usb partition in /dev/mapper (e.g., sda1): " disk_suffix

    usb_name="/dev/mapper/${disk_suffix}"

    user_response="n"

    read -rp "You are about to WRITE and ERASE disk ${disk_name} | Are you sure? y/N: " user_response

    if [[ $user_response != y ]]; then
        echo "Exiting script"
        exit
    fi

    echo "Mounting Ventoy..."
    mkdir -p /mnt/ventoy
    mount $usb_name /mnt/ventoy
    cd /mnt/ventoy/Images

    echo "Writing image to $disk_name..."
    lz4 -dc PBX-LVM.img.lz4 | dd of=$disk_name status=progress
fi


# Repartition disk using sgdisk
echo "Repairing partition tables"
sgdisk -e $disk_name
sgdisk -v $disk_name

partprobe || true

echo "Expanding Partition 3 on $disk_name..."
sgdisk --delete=3 $disk_name
sgdisk --new=3:0:0 --typecode=3:8e00 --change-name=3:"Linux LVM" $disk_name

partprobe || true

# Resize physical volume
echo "Resizing physical volume..."
pvresize ${disk_name}3

# Resize swap and root logical volumes
echo "Resizing swap and root volumes..."
lvresize -L +7G PBX-vg/swap_1
lvresize -l +100%FREE --resizefs PBX-vg/root


# Mount restored system for chroot
echo "Mounting restored system..."
mkdir -p /mnt/debian
mount /dev/mapper/PBX--vg-root /mnt/debian/
mount ${disk_name}2 /mnt/debian/boot
mount ${disk_name}1 /mnt/debian/boot/efi


for i in "${mount_points[@]}"; do mount -B $i /mnt/debian$i; done

# Chroot and reinstall GRUB
echo "Entering chroot to reinstall GRUB..."
chroot /mnt/debian /bin/bash -c "
grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=debian
update-grub
"

echo "Restore and resize complete!"

cleanup