#!/bin/bash

DEFAULT="$1"
LASTBOOT="lastboot"

EFI_DIR="$(bootctl -p)"
ROOT_FS=$(findmnt -no SOURCE /)
BASE="${ROOT_FS%/*}"

# default UKI name used by kernel-install
UKI_NAME="$(cat /etc/machine-id)-$(uname -r).efi"

if [ "$ROOT_FS" = "$BASE/$DEFAULT" ]; then
  echo "sucessfully booted to $BASE/$DEFAULT, preparing lastboot"

  # destroy the old lastboot file system and children
  zfs destroy -Rv "$BASE/$DEFAULT@$LASTBOOT"

  # make a new lastboot file system
  zfs snap -r "$BASE/$DEFAULT@$LASTBOOT"

  for OLDFS in $(zfs list -rHo name -t filesystem "$BASE/$DEFAULT"); do
    NEWFS=$(sed "s|^$BASE/$DEFAULT|$BASE/$LASTBOOT-$DEFAULT|" <<< $OLDFS)
    echo "cloning $OLDFS to $NEWFS"
    zfs clone -o canmount=noauto -o mountpoint=none "$OLDFS@$LASTBOOT" "$NEWFS"
  done

  zfs set mountpoint=/ "$BASE/$LASTBOOT-$DEFAULT"

  # backup the current UKI
  echo "copying $UKI_NAME to $LASTBOOT-$DEFAULT.efi"
  cp "$EFI_DIR/EFI/Linux/$UKI_NAME" "$EFI_DIR/EFI/Linux/$LASTBOOT-$DEFAULT.efi"

  echo "remove old initramfs images from /boot"
  rm -f /boot/initramfs-*.img

  echo "unpack current initramfs to /boot/initramfs-$(uname -r).img"
  objcopy -O binary -j.initrd "$EFI_DIR/EFI/Linux/$LASTBOOT-$DEFAULT.efi" "/boot/initramfs-$(uname -r).img"
else
  echo "not booted to $BASE/$DEFAULT, taking no action"
fi