#!/bin/bash

set -euo pipefail

# ==============================================================
#  ZFSBootMenu Deployment Script (Legacy Boot / Devuan)
#  Safe Unified Kernel Image Extractor & GRUB Auto Integration
# ==============================================================

# Boot Flow Visualization
# 1.    BIOS/Legacy: Execution starts at the MBR/VBR.
# 2.    GRUB Stage 1 & 2: GRUB loads its modules and reads your custom script from /etc/grub.d/.
# 3.    The "Direct" Load: GRUB uses the linux and initrd commands to pull the extracted ZBM components into RAM.
# 4.    ZBM Environment: ZBM initializes, finds your ZFS datasets, and provides the UI.
# 5.    kexec: Once you select a kernel in ZBM, it uses kexec to replace itself with your actual Devuan kernel.

# --- Configuration ---
REPO="zbm-dev/zfsbootmenu"
ZBM_DIR="/boot/zfsbootmenu"
KEYRING="/usr/share/keyrings/zfsbootmenu.gpg"
PUBKEY_URL="https://raw.githubusercontent.com/${REPO}/master/releng/keys/zfsbootmenu.pub.gpg"
HASH_FILE="sha256.txt"
SIG_FILE="sha256.sig"

# --- 1. Dependency Check ---
echo "[*] Checking dependencies..."
for cmd in objcopy curl jq gpg findmnt zfs; do
  if ! command -v "$cmd" >/dev/null 2>&1; then
    echo "Installing missing dependency: $cmd..."
    apt-get update -qq && apt-get install -y binutils curl jq gnupg zfsutils-linux >/dev/null
  fi
done

# --- 2. Failure Recovery Logic ---
restore_backups() {
  echo "!!! Error encountered. Restoring previous version from .old backups..."
  for file in vmlinuz-zbm initramfs-zbm.img; do
    if [ -f "$ZBM_DIR/$file.old" ]; then
      mv -f "$ZBM_DIR/$file.old" "$ZBM_DIR/$file"
    fi
  done
}

# --- 3. Version & Hash Idempotency ---
echo "[*] Checking for latest ZFSBootMenu release..."
LATEST_TAG=$(curl -sfSL "https://api.github.com/repos/$REPO/releases/latest" | jq -r .tag_name)
# LATEST_TAG=$(curl -sfSL "https://api.github.com/repos/$REPO/releases/latest" | jq -r '.tag_name' | tr -d '"')
echo "Using clean tag: '$LATEST_TAG'" # v3.1.0 ✓

MACHINE="$(uname -m)"                            # x86_64
KERNEL="$(uname -r)"                             # 6.12.6-amd64
KERNEL_VERSION=$(echo "$KERNEL" | cut -d. -f1,2) # 6.12
KERNEL_NAME="$(uname -s)"                        # Linux
KERNEL_NAME="${KERNEL_NAME,,}"                   # linux lowercase
EFI_NAME="zfsbootmenu-recovery-${MACHINE}-${LATEST_TAG}-${KERNEL_NAME}${KERNEL_VERSION}.EFI"
HASH_URL="https://github.com/$REPO/releases/download/$LATEST_TAG/${HASH_FILE}"
SIG_URL="https://github.com/$REPO/releases/download/$LATEST_TAG/${SIG_FILE}"
REMOTE_HASH_DATA=$(curl -sfSL --connect-timeout 10 "$HASH_URL")
# REMOTE_HASH=$(echo "$REMOTE_HASH_DATA" | grep "$EFI_NAME" | awk '{print $1}')
# REMOTE_HASH=$(echo "$REMOTE_HASH_DATA" | grep "$EFI_NAME" | sed -n "s/.*= \([0-9a-f]*\).*/\1/p")
REMOTE_HASH=$(echo "$REMOTE_HASH_DATA" | grep "$EFI_NAME" | sed -n "s/.*= *\([0-9a-f]\{64\}\).*/\1/p")

mkdir -p "$ZBM_DIR"
if [[ -f "$ZBM_DIR/.hash" && "$(cat "$ZBM_DIR/.hash")" == "$REMOTE_HASH" ]]; then
  echo "✔ ZFSBootMenu $LATEST_TAG already installed and verified. Exiting."
  exit 0
fi

# --- 4. Workspace & Error Handling ---
TMP_DIR=$(mktemp -d)

cleanup_on_error() {
  local exit_code=$?
  if [[ $exit_code -ne 0 ]]; then
    echo "ERROR: Script failed (line $BASH_LINENO). Rolling back..."
    restore_backups
  fi
  rm -rf "$TMP_DIR"
  exit "$exit_code"
}
trap 'cleanup_on_error' EXIT ERR

# --- 5. Secure Download & Verification ---
echo "[*] Ensuring GPG keyring exists..."
[ ! -f "$KEYRING" ] && sudo curl -sfSL "$PUBKEY_URL" -o "$KEYRING"

cd "$TMP_DIR"
echo "[*] Downloading latest UKI: $EFI_NAME..."
curl -sfSLO "https://github.com/$REPO/releases/download/$LATEST_TAG/$EFI_NAME"
curl -sfSLO "${HASH_URL}"

# echo "[*] Verifying GPG signature..."
# gpg --no-default-keyring --keyring "$KEYRING" --verify "${SIG_FILE}" "$EFI_NAME" 2>/dev/null || {
#   echo "ERROR: GPG signature verification failed!"
#   exit 1
# }

# CORRECT SHA256:
curl -sfSLO "$HASH_URL"
echo "$REMOTE_HASH  $EFI_NAME" | sha256sum -c || {
  echo "ERROR: SHA256 checksum failed!"
  exit 1
}

# --- 6. Extraction & Validation ---
echo "[*] Extracting components from UKI..."
objcopy --dump-section .linux="vmlinuz-zbm.tmp" "$EFI_NAME"
objcopy --dump-section .initrd="initramfs-zbm.img.tmp" "$EFI_NAME"

# Validate outputs
if [[ ! -s "vmlinuz-zbm.tmp" || ! -s "initramfs-zbm.img.tmp" ]]; then
  echo "ERROR: Extracted components are empty or missing! Abort."
  exit 1
fi

echo "[*] Replacing ZFSBootMenu components..."
sudo bash -c "
  [ -f "$ZBM_DIR/vmlinuz-zbm" ] && cp "$ZBM_DIR/vmlinuz-zbm" "$ZBM_DIR/vmlinuz-zbm.old"
  [ -f "$ZBM_DIR/initramfs-zbm.img" ] && cp "$ZBM_DIR/initramfs-zbm.img" "$ZBM_DIR/initramfs-zbm.img.old"

  mv vmlinuz-zbm.tmp "$ZBM_DIR/vmlinuz-zbm"
  mv initramfs-zbm.img.tmp "$ZBM_DIR/initramfs-zbm.img"
  chmod 600 "$ZBM_DIR/vmlinuz-zbm" "$ZBM_DIR/initramfs-zbm.img"
  echo "$REMOTE_HASH" >"$ZBM_DIR/.hash"
"

# --- 7. GRUB Integration ---
ZBM_POOL=$(findmnt -n -o SOURCE -T "$ZBM_DIR" | cut -d'/' -f1)
ZBM_DATASET_MNT=$(findmnt -n -o TARGET -T "$ZBM_DIR")
REL_PATH="${ZBM_DIR#$ZBM_DATASET_MNT}"
[[ -z "$REL_PATH" ]] && REL_PATH="/"
REL_PATH="${REL_PATH//|/}"

echo "[*] Generating GRUB configuration for pool: $ZBM_POOL..."

conf_print_grub_menu_zbm() {
  cat <<EOF
#!/bin/sh
exec tail -n +3 \$0
# Dynamically generated by ZBM Deployment Script

submenu 'ZFSBootMenu>Recovery Options' {
  menuentry 'ZFSBootMenu (Direct Boot)' --class zfs --class gnu-linux {
      insmod part_gpt
      insmod zfs
      search --no-floppy --set=root --label $ZBM_POOL
      linux $REL_PATH/vmlinuz-zbm loglevel=4 zbm.import_delay=5
      initrd $REL_PATH/initramfs-zbm.img
  }

  menuentry 'ZFSBootMenu (Previous/Backup)' --class zfs --class gnu-linux {
      insmod part_gpt
      insmod zfs
      search --no-floppy --set=root --label $ZBM_POOL
      linux $REL_PATH/vmlinuz-zbm.old loglevel=4 zbm.import_delay=5
      initrd $REL_PATH/initramfs-zbm.img.old
  }

  menuentry 'ZFSBootMenu (Recovery/Read-Only)' --class recovery --class zfs {
      insmod part_gpt
      insmod zfs
      search --no-floppy --set=root --label $ZBM_POOL
      linux $REL_PATH/vmlinuz-zbm loglevel=4 zbm.import_delay=5 zbm.readonly=1
      initrd $REL_PATH/initramfs-zbm.img
  }

  menuentry 'ZFSBootMenu (Recovery/Force Import)' --class recovery --class zfs {
      insmod part_gpt
      insmod zfs
      search --no-floppy --set=root --label $ZBM_POOL
      linux $REL_PATH/vmlinuz-zbm loglevel=4 zbm.import_delay=5 zbm.readonly=1 zbm.prefer=$ZBM_POOL!
      initrd $REL_PATH/initramfs-zbm.img
  }
}
EOF
}
conf_print_grub_menu_zbm | sudo tee /etc/grub.d/40_zfsbootmenu >/dev/null
sudo chmod +x /etc/grub.d/40_zfsbootmenu

if command -v update-grub >/dev/null 2>&1; then
  sudo update-grub
else
  sudo grub-mkconfig -o /boot/grub/grub.cfg
fi

echo "✔ Deployment complete and verified."