#!/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-Image-*.tar.zst"

MOUNT_PATH="/mnt/debian"
MOUNT_POINTS=("/dev" "/dev/pts" "/proc" "/sys" "/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
  check_dependencies

  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

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

check_dependencies() {
  local missing=()
  local required_tools=("sgdisk" "pv" "pvcreate" "vgcreate" "lvcreate" "vgchange" "mkfs.vfat" "mkfs.ext4" "mkswap" "tar" "zstd" "lsblk" "partprobe" "chroot")

  for tool in "${required_tools[@]}"; do
    if ! command -v "$tool" &>/dev/null; then
      missing+=("$tool")
    fi
  done

  if [[ ${#missing[@]} -gt 0 ]]; then
    echo -e "${RED}Missing required tools: ${missing[*]}${RESET}"
    echo
    echo "Install with:"
    echo "  apt install gdisk pv lvm2 zstd dosfstools e2fsprogs parted util-linux coreutils"
    exit 1
  fi

  echo -e "${GREEN}All dependencies found${RESET}"
}

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

  echo -e "${RED}WARNING: ALL DATA ON ${disk_name} WILL BE DESTROYED${RESET}"
  lsblk $disk_name 2>/dev/null || true
  echo
  read -rp "Are you sure you want to continue? (yes/no): " confirm
  if [[ "$confirm" != "yes" ]]; then
    echo "Aborted."
    exit 0
  fi
  echo "Chose disk $disk_name"
  echo
}

rewrite_gpt_table() {
  vgchange -an $VGROUP 2>/dev/null || true
  vgremove -f $VGROUP 2>/dev/null || true
  pvremove -ff $ROOT_PART 2>/dev/null || true

  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 -I -F 32 $EFI_PART
  mkfs.ext4 -F $BOOT_PART
  mkfs.ext4 -F "${VG_PATH}${ROOT_LV}"
  mkswap -f "${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

  local img_file_dir=$(ls -1 ./PBX-Image-*.tar.zst 2>/dev/null | head -n1)

  if [[ -z "$img_file_dir" ]]; then
    img_file_dir=$(ls -1 ./Image/PBX-Image-*.tar.zst 2>/dev/null | head -n1)
  fi
  if [[ -z "$img_file_dir" ]]; then
    echo -e "${RED}NO TAR.ZST FILE FOUND${RESET}"
    exit 100
  fi

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

  pv "${img_file_dir}" | tar -I zstd -x -C "${MOUNT_PATH}"

  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"

  blkid -s UUID -o value "$target_part"
}

mount_partitions() {
  echo "Mounting partitions..."

  vgchange -ay $VGROUP

  mkdir -p "${MOUNT_PATH}"
  mount "${VG_PATH}${ROOT_LV}" $MOUNT_PATH

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

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

mount_chroot() {
  echo "Mounting restored system..."

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

  if [[ -d /sys/firmware/efi/efivars ]]; then
    mount -t efivarfs efivarfs "${MOUNT_PATH}/sys/firmware/efi/efivars"
  fi

  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_GRUB_INSTALL="grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=debian"
  CHROOT_GRUB_UPDATE="update-grub"

  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