#!/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
            ;;
        -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

# Set your target disk here
udevadm trigger

lsblk
echo

echo "Enter disk 1 for raid (e.g., sdb): "
read disk_0_suffix

echo "Enter disk 2 for raid (e.g., sdc): "
read disk_1_suffix

disk_name_0="/dev/${disk_suffix}"
disk_name_1="/dev/${disk_suffix}"

echo "Raid0 will be configured on $disk_name_0 and $disk_name_1"
echo

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

  for disk in "$disk_name_0" "$disk_name_1"; do
      echo "Partitioning $disk..."

      sgdisk --zap-all "$disk"

      sgdisk -n 1:0:+512M -t 1:ef00 -c 1:"EFI System" "$disk"

      sgdisk -n 2:0:+1G -t 2:8300 -c 2:"Boot" "$disk"

      sgdisk -n 3:0:0 -t 3:8300 -c 3:"Root" "$disk"
  done

  mdadm --create /dev/md0 --level=1 --raid-devices=2 "${disk_name_0}2" "${disk_name_1}2"
  mdadm --create /dev/md1 --level=1 --raid-devices=2 "${disk_name_0}3" "${disk_name_1}3"

  disk_name=/dev/md1

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

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

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

mdadm --detail --scan >> /mnt/debian/etc/mdadm/mdadm.conf

# 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
update-initramfs -u
"

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

cleanup