#!/bin/bash
set -e


# THE RAINBOW HAHAHAH
RED='\e[31m'
GREEN='\e[32m'
YELLOW='\e[33m'
RESET='\e[0m'


# root check
if [[ $EUID -ne 0 ]]; then
    echo -e "${RED}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
            ;;
        --noretable)
            NO_RETABLE=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 2>/dev/null || true

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

    umount -lf /mnt/debian/boot/efi 2>/dev/null || true
    umount -lf /mnt/debian/boot 2>/dev/null || true
    umount -lf /mnt/debian 2>/dev/null || true
}
trap 'echo "Status Code: $?"; cleanup;' EXIT


get_partition() {
    local disk_name="$1"
    local part_num="$2"

    local part_path="$(
      lsblk -lnp -o NAME,TYPE "$disk_name" \
      | awk '$2=="part"{print $1}' \
      | sed -n "${part_num}p"
    )"

    printf "$part_path"
}


get_partuuid() {
  local disk_name="$1"
  local part_num="$2"

  local part_uuid="$(
    lsblk -ln -o PATH,PARTUUID $disk_name | awk '$1==$(get_partition $disk_name $part_num) {print $2}'
  )"

  printf "$part_uuid"
}


# 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_RETABLE == true ]]; then
    echo -e "${YELLOW}--noretable is specified. Tables will not be erased and no partitions will be created${RESET}"
  else 
    sgdisk -Z $disk_name

    sgdisk -c $disk_name

    sgdisk -n 1:1MiB:+512M -t 1:ef00 -c 1:"EFI System" $disk_name
    sgdisk -n 2:0:+1G -t 2:8300 -c 2:"boot" $disk_name
    sgdisk -n 3:0:0 -t 3:8e00 -c 3:"Linux LVM" $disk_name

    partprobe || true
fi

if [[ $NO_WRITE == true ]]; then
    echo -e "${YELLOW}--nowrite is specified. No need for Img no disks will be written to${RESET}"
  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 "$(echo -e "${RED}You are about to WRITE and ERASE disk ${disk_name} | Are you sure? y/N: ${RESET}")" 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 $(get_partition $disk_name "3")..."
    lz4 -dc PBX.img.lz4 | dd of=$(get_partition $disk_name "3") status=progress

    partprobe || true
fi


# Resize physical volume
echo "Resizing physical volume..."
pvresize $(get_partition $disk_name "3") || true


# 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
mount_path="/mnt/debian"
echo "Mounting restored system..."
mkdir -p $mount_path
mount /dev/mapper/PBX--vg-root $mount_path
mount $(get_partition $disk_name "2") ${mount_path}/boot
mount $(get_partition $disk_name "1") ${mount_path}/boot/efi

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

echo "Generating fstab..."
boot_uuid=$()

cp ./fstab.template ${mount_path}/etc/fstab

sed --in-place "s/boot_uuid/$(get_partuuid $disk_name "2")/g" ${mount_path}/etc/fstab
sed --in-place "s/efi_uuid/$(get_partuuid $disk_name "1")/g" ${mount_path}/etc/fstab


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

echo -e "${GREEN}Restore and resize complete!${RESET}"

cleanup