#!/usr/bin/env bash
set -x

# Dependencies
sudo apt update && sudo apt install -y binutils wget binwalk # lastversion

# --- CONFIGURATION ---
REPO="zbm-dev/zfsbootmenu"
ZBM_VERSION=$(lastversion --format tag "${REPO}")
ZBM_RELEASE="recovery"
KERNEL=$(uname -r | cut -d'.' -f1,2)
ARCH=$(uname -m)

# Target directory for GRUB to find
FINAL_DEST=/boot/zfsbootmenu
sudo mkdir -p ${FINAL_DEST}

ZBM_EFI="zfsbootmenu-${ZBM_RELEASE}-${ARCH}-${ZBM_VERSION}-linux${KERNEL}.EFI"
ZBM_URL="https://github.com/zbm-dev/zfsbootmenu/releases/download/${ZBM_VERSION}/${ZBM_EFI}"

# --- TMPFS SETUP ---
WORK_DIR=$(mktemp -d -p /dev/shm zbm_extract.XXXXXX)
cleanup() {
  echo "--- Cleaning up RAM storage ---"
  rm -rf "$WORK_DIR"
  rm -f "$ZBM_EFI"
}
trap cleanup EXIT

echo "Downloading ZFSBootMenu ${ZBM_VERSION}..."
# lastversion -d \
#   --output "${ZBM_EFI}" \
#   zbm-dev/zfsbootmenu \
#   --filter "${ZBM_RELEASE}.*${ARCH}.*${ZBM_VERSION}.*${KERNEL}.*EFI"

sudo curl -L "${ZBM_URL}" -o "${ZBM_EFI}"

# Extraction paths
KERNEL_OUT="${WORK_DIR}/vmlinuz-zbm"
INITRD_OUT="${WORK_DIR}/initrd-zbm.img"

echo "--- Extracting assets from $ZBM_EFI ---"

# 1. Extract Kernel & Initrd
if objcopy --dump-section .linux="$KERNEL_OUT" "$ZBM_EFI" 2>/dev/null &&
  objcopy --dump-section .initrd="$INITRD_OUT" "$ZBM_EFI" 2>/dev/null; then

  echo "[✓] Extraction successful. Moving to $FINAL_DEST..."
  sudo cp "$KERNEL_OUT" "$FINAL_DEST/"
  sudo cp "$INITRD_OUT" "$FINAL_DEST/"
else
  echo "[!] Extraction failed. Check if binutils is installed correctly."
  exit 1
fi

# --- GRUB MENU GENERATION ---
conf_print_zbm_grub_menu() {
  cat <<EOF
#!/bin/sh
exec tail -n +3 \$0
# This file provides an easy way to add custom menu entries.
menuentry "ZFSBootMenu (Extracted)" --class zfs --class gnu-linux {
    insmod part_gpt
    insmod zfs
    insmod gzio

    # Use UUID instead of file search
    search --no-floppy --fs-uuid --set=root f383eb9e238f6712

    echo "Loading ZFSBootMenu Kernel..."
    linux /zfsbootmenu/vmlinuz-zbm zbm.import_policy=hostid zbm.prefer=rpool loglevel=4

    echo "Loading ZFSBootMenu Initrd..."
    initrd /zfsbootmenu/initrd-zbm.img
}
EOF
}

echo "Generating GRUB configuration..."
conf_print_zbm_grub_menu | sudo tee /etc/grub.d/40_zbm
sudo chmod +x /etc/grub.d/40_zbm

echo "Updating GRUB..."
sudo update-grub

echo "---------------------------------------"
echo "Done! Files are located in $FINAL_DEST"
grep -A 10 "menuentry \"ZFSBootMenu" /boot/grub/grub.cfg