#!/bin/bash
set -e

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

NO_WRITE=false
NO_RETABLE=false
NO_USB=false
NO_CLEANUP=false

IMG_FILE="PBX-root.img.lz4"

MOUNT_PATH="/mnt/debian"
MOUNT_POINTS=("/dev" "/dev/pts" "/proc" "/sys" "/sys/firmware/efi/efivars" "/run")

BOOT_PART="NULL"
EFI_PART="NULL"
ROOT_PART="NULL"

VGROUP="SangomaPBX-vg"
ROOT_LV="root"
SWAP_LV="swap_1"
VG_PATH="/dev/mapper/SangomaPBX--vg-"

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

  dialogue

  assign_part_var

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

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

  write_img

  mount_chroot

  gen_fstab

  install_grub

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

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

dialogue() {
  lsblk

  echo
  read -rp "Enter the target disk (e.g., sda): " disk_suffix

  disk_name="/dev/${disk_suffix}"

  echo "Chose disk $disk_name"
  echo
}

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

  lsblk $disk_name

  echo -e "${GREEN}GPT table created. ${RESET}"
}

assign_part_var() {
  EFI_PART=$(get_partition $disk_name "1")
  BOOT_PART=$(get_partition $disk_name "2")
  ROOT_PART=$(get_partition $disk_name "3")
}

lvm_init() {
  pvcreate $ROOT_PART -ff

  vgcreate $VGROUP $ROOT_PART -f

  vgchange -ay $VGROUP

  lvcreate -L 4G $VGROUP -n swap_1
  lvcreate -l 100%FREE $VGROUP -n root

  echo -e "${GREEN}LVM initialized ${RESET}"
}

make_filesystem() {
  mkfs.vfat -F 32 $EFI_PART
  mkfs.ext4 $BOOT_PART
  mkswap "${VG_PATH}${SWAP_LV}"
  echo -e "${GREEN}Make filesystems finished.${RESET}"
}

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

  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}"
  elif [[ ! -f $img_file_dir ]]; then
    echo -e "${RED}NO IMG FILE FOUND${RESET}"
    exit 100
  fi

  echo "Writing image to ${VG_PATH}${ROOT_LV}..."

  lz4 -dc ${img_file_dir} | partclone.ext4 -r -o "${VG_PATH}${ROOT_LV}"

  partprobe || true

  echo -e "${GREEN}Finished writing image to disk${RESET}"
}

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 target_part="$1"

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

  printf "$part_uuid"
}

mount_chroot() {
  vgchange -ay $VGROUP

  echo "Mounting restored system..."

  mkdir -p "${MOUNT_PATH}/boot"
  mount "${VG_PATH}${ROOT_LV}" $MOUNT_PATH
  mount $BOOT_PART "${MOUNT_PATH}/boot"

  mkdir -p "${MOUNT_PATH}/boot/efi"
  mount $EFI_PART "${MOUNT_PATH}/boot/efi"

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

  echo -e "${GREEN}Mounts prepared for chroot.${RESET}"
}

gen_fstab() {
  echo "Generating fstab..."

  boot_uuid=$(get_partuuid $BOOT_PART)
  efi_uuid=$(get_partuuid $EFI_PART)

  cp ./fstab_template ${MOUNT_PATH}/etc/fstab

  sed --in-place "s/BOOT_UUID/${boot_uuid}/g" ${MOUNT_PATH}/etc/fstab
  sed --in-place "s/EFI_UUID/${efi_uuid}/g" ${MOUNT_PATH}/etc/fstab

  echo -e "${GREEN}fstab write complete!${RESET}"
}

install_grub() {
  echo "Entering chroot to reinstall GRUB..."

  CHROOT_LINUX_INSTALL="dpkg -i /var/cache/apt/archives/linux-image-amd64_6.1.164-1_amd64.deb"
  CHROOT_INITRAMFS="update-initramfs -u -k all"
  CHROOT_GRUB_INSTALL="grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=debian"
  CHROOT_GRUB_UPDATE="update-grub"

  chroot_cmd "$CHROOT_LINUX_INSTALL"
  chroot_cmd "$CHROOT_INITRAMFS"
  chroot_cmd "$CHROOT_GRUB_INSTALL"
  chroot_cmd "$CHROOT_GRUB_UPDATE"

  echo -e "${GREEN}Grub installed!${RESET}"
}

chroot_cmd() {
  CHROOT_PATH="export PATH=/usr/sbin:/usr/bin:/sbin:/bin"
  chroot $MOUNT_PATH /bin/bash -c "${CHROOT_PATH} && ${1}"
}

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

  echo "Cleaning Mounts"

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

  umount -lf "${MOUNT_PATH}/boot/efi" 2>/dev/null || true
  umount -lf "${MOUNT_PATH}/boot" 2>/dev/null || true
  umount -lf "${MOUNT_PATH}" 2>/dev/null || true

  vgchange -an $VGROUP
  echo -e "${GREEN}Clean.${RESET}"
}

main