#!/bin/bash
set -e

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

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/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


if [[ $DEBUG == "true" ]]; then
    echo "DEBUG is true"
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
if [[ $DEBUG != "true" ]]; then
    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