#!/bin/bash
set -e


cleanup() {
    echo "Cleaning Mounts"
    
    mount_points=("/dev" "/dev/pts" "/proc" "/sys" "/sys/firmware/efi/efivars" "/run")

    for i in $mount_points; do umount -lf /mnt/debian$i || true; done

    umount -lf /mnt/boot/efi 2>/dev/null || true
    umount -lf /mnt/boot 2>/dev/null || true
    umount -lf /mnt 2>/dev/null || true
}
trap cleanup EXIT

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

# 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
read -rp "Enter the name of the usb partition in /dev/mapper (e.g., sda1): " disk_suffix
usb_name="/dev/mapper/${disk_suffix}"

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

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

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

# Repartition disk using sgdisk
echo "Expanding Partition 3 on $disk_name..."
sgdisk -e $disk_name
sgdisk -v $disk_name
sgdisk --delete=3 $disk_name
sgdisk --new=3:0:0 --typecode=3:8e00 --change-name=3:"Linux LVM" $disk_name

# Inform the OS of partition table changes
partprobe $disk_name
sleep 2  # Give the system a moment to recognize changes

# 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

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

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