#!/bin/bash
set -e

RED='\e[31m'
GREEN='\e[32m'
YELLOW='\e[33m'
RESET='\e[0m'

MOUNT_POINTS=("/dev" "/dev/pts" "/proc" "/sys" "/sys/firmware/efi/efivars" "/run")
IMG_FILE="PBX.img3.lz4"
NO_WRITE=false
NO_RETABLE=false
NO_USB=false
NO_CLEANUP=false

for arg in "$@"; do
  case $arg in
  --nowrite)
    NO_WRITE=true
    shift
    ;;
  --noretable)
    NO_RETABLE=true
    shift
    ;;
  --nousb)
    NO_USB=true
    shift
    ;;
  --nocleanup)
    NO_CLEANUP=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

main() {
  root_check

  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

  trap 'echo "Status Code: $?"; cleanup;' EXIT

  if [[ $NO_RETABLE == true ]]; then
    echo -e "${YELLOW}--noretable is specified. Tables will not be erased and no partitions will be created${RESET}"
  else
    rewrite_gpt_table
    make_filesystem
    lvm_init
  fi

  if [[ $NO_WRITE == true ]]; then
    echo -e "${YELLOW}--nowrite is specified. No need for Img no disks will be written to${RESET}"
  fi

  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}"
}

rewrite_gpt_table() {
  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
}

write_img() {
  img_file_dir="null"

  if [[ -f "./${IMG_FILE}" ]]; then
    img_file_dir="./${IMG_FILE}"
  elif [[ ! -f "./Image/${IMG_FILE}" ]]; then
    img_file_dir="./Image/${IMG_FILE}"
  else
    img_file_dir="null"
  fi

  echo "Writing image to $(get_partition $disk_name "3")..."

  lz4 -dc ${img_file_dir} | partclone.ext4 -r -o $(get_partition $disk_name "3")
  partprobe || true
}

lvm_init() {
  pvcreate $(get_partition $disk_name "3")

  vgcreate PBX-vg $(get_partition $disk_name "3")

  vgchange -ay PBX-vg

  lvcreate -L 4G PBX-vg -n swap_1
  lvcreate -l 100%FREE PBX-vg -n root
}

make_filesystem() {
  mkfs.vfat -F 32 $(get_partition $disk_name "1")
  mkfs.ext4 $(get_partition $disk_name "2")
}

root_check() {
  if [[ $EUID -ne 0 ]]; then
    echo -e "${RED}You must be root to do this." 1>&2
    exit 100
  fi
}

cleanup() {
  if [[ $NO_CLEANUP = true ]]; then
    echo -e "${YELLOW}No cleanup specified. Make sure to manually clean mounts." 1>&2
    exit 100
  fi

  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
}

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"
}

mount_chroot() {
  vgchange -ay PBX-vg

  mount_path="/mnt/debian"

  echo "Mounting restored system..."

  mkdir -p $mount_path/boot
  mount /dev/mapper/PBX--vg-root $mount_path
  mount $(get_partition $disk_name "2") ${mount_path}/boot

  mkdir -p $mount_path/boot/efi
  mount $(get_partition $disk_name "1") ${mount_path}/boot/efi

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

main
cleanup