#!/bin/bash

set -eu

UKI_ROOT='/boot/efi/EFI/puavo'
UKI_LOCK='/run/puavo-uki-image.lock'

logmsg() {
  logger -s -t puavo-cleanup-uki-installations "$@"
}

panic() {
  logmsg -p user.err "$1"
  exit 1
}

images_dir=/images

while [ $# -gt 0 ]; do
  case "$1" in
    --images-dir)
      images_dir="$2"; shift 2 ;;
    *)
      shift ;;
  esac
done

[ -d "$images_dir" ] || panic "images directory '$images_dir' does not exist"

if [ -z "${UKI_MODE:-}" ]; then
  PUAVO_INSTALL_LIBDIR="/usr/lib/puavo-ltsp-install"
  export UKI_MODE
  UKI_MODE=$("$PUAVO_INSTALL_LIBDIR"/is-uki-install "$images_dir")
fi

[ "$UKI_MODE" = "1" ] || exit 0

exec {lockfd}<> "$UKI_LOCK"
flock -nx "$lockfd" || panic "another UKI image operation is already running"

find_efi_device_path() {
  local diskdev="$1"

  lsblk -n -l -o PATH,PARTTYPE "$diskdev" \
    | awk '$2 == "c12a7328-f81f-11d2-ba4b-00a0c93ec93b" { print $1; exit(0) }'
}

mount_disk_efi_partition() {
  local diskdev="$1"

  efi_devpath=$(find_efi_device_path "$diskdev")
  if [ -z "$efi_devpath" ]; then
    logmsg -p user.err "error: failed to find EFI partition in ${diskdev}"
    return 1
  fi

  mount "$efi_devpath" /boot/efi 1>/dev/null
}

cleanup_orphaned_uki_directories() {
  [ -d "$UKI_ROOT" ] || return 0

  local status=0

  for uki_directory in "$UKI_ROOT"/*/; do
    [ -d "$uki_directory" ] || continue

    local image_name
    image_name=$(basename "$uki_directory")

    # Skip directories that are not UKI installation directories
    if ! find "$uki_directory" -maxdepth 1 -name '*.uki.efi' -type f | grep -q .; then
      continue
    fi

    local image_path_no_suffix="${images_dir}/${image_name}"
    if [ -f "${image_path_no_suffix}.img" ]; then
      # we have an image corresponding to UKI, but as bootservers may have
      # several images, we check if the images are something to boot from
      if [ -e "${image_path_no_suffix}.default" ] \
             || [ -e "${image_path_no_suffix}.backup" ]; then
        # we keep this one, because it has been marked as either a default
        # or a backup image
        continue
      fi
      if [ "$(stat -c %h "${image_path_no_suffix}.img")" -gt 1 ]; then
        # if the image file is hardlinked, it is likely hard linked to
        # either ltsp.img or ltsp-backup.img and it also qualifies
        # (do not remove the UKI directory)
        continue
      fi
    fi

    # remove the UKI directory
    logmsg -p user.info "removing orphaned UKI directory '${uki_directory}'"
    if ! rm --one-file-system -rf "$uki_directory"; then
      logmsg -p user.err "error: failed to remove '${uki_directory}'"
      status=1
    fi
  done

  return "$status"
}

boot_disks=$(/usr/lib/puavo-core/puavo-get-boot-disks "$images_dir")
[ -n "$boot_disks" ] || panic "failed to find boot disks"

status=0

for diskdev in $boot_disks; do
  mount_disk_efi_partition "$diskdev" || {
    logmsg -p user.err "error: failed to mount EFI partition on ${diskdev}"
    status=1
    continue
  }

  cleanup_orphaned_uki_directories || status=1

  # Remove fsck recovery files at the ESP root so they do not fill up the partition.
  find /boot/efi -maxdepth 1 -type f -name 'FSCK*.REC' -delete || true

  sync || true # Ensure ESP writes are fully transferred before unmount.
  umount /boot/efi || true
done

exit $status
