#!/bin/bash
set -eu

UKI_ROOT='/boot/efi/EFI/puavo'

panic() {
  echo "error: $1" >&2
  exit 1
}

usage() {
  echo "Usage: $0 <image-mountpoint> <images-dir> <host-type>" >&2
  exit 1
}

[ $# -eq 3 ] || usage

image_mountpoint="$1"
images_dir="$2"
hosttype="$3"

[ -d "$image_mountpoint" ] || panic "image mountpoint '$image_mountpoint' does not exist"
[ -d "$UKI_ROOT" ] || panic "UKI root '$UKI_ROOT' does not exist"
tmpdir=''

cleanup_temporary_files() {
  find "$UKI_ROOT" -xdev -type d -name '*.tmp' -exec rm -rf {} + 2>/dev/null || true
  find "$UKI_ROOT" -xdev -type f -name '*.tmp' -delete 2>/dev/null || true
  test -n "$tmpdir" && rm -rf "$tmpdir"
}

cleanup_temporary_files || echo "warning: failed to remove temporary files" >&2
trap 'cleanup_temporary_files' INT TERM EXIT

tmpdir=$(mktemp -d /tmp/puavo-install-uki.XXXXXX) \
  || panic 'could not create a temporary directory'

boot_directory="${image_mountpoint}/boot"
[ -d "$boot_directory" ] || panic "boot directory '$boot_directory' does not exist"

determine_selected_kernel_version() {
  local kernel_image_symbolic_link="${boot_directory}/vmlinuz"
  local kernel_image
  kernel_image=$(readlink -f "${kernel_image_symbolic_link}")

  echo "${kernel_image#"${kernel_image_symbolic_link}-"}"
}

install_boot_entry_config() {
  local extra_directory="$1"
  local class="$2"
  local image_name="$3"
  local release="$4"
  local enabled="$5"
  local environment_file="${extra_directory}/grubenv"
  local temporary_environment_file="${environment_file}.tmp"

  rm -f "$temporary_environment_file"
  grub-editenv "$temporary_environment_file" create
  grub-editenv "$temporary_environment_file" set "class=${class}"
  grub-editenv "$temporary_environment_file" set "image_name=${image_name}"
  grub-editenv "$temporary_environment_file" set "enabled=${enabled}"
  grub-editenv "$temporary_environment_file" set "release=${release}"
  mv "$temporary_environment_file" "$environment_file"
}

patch_uki() {
  local initrd_img_path kernel_version kernel_version_with_suffix \
        tmp_uki_destination tmp_ukibase_path uki_checksum_file \
        uki_destination uki_dir uki_patch vmlinuz_path
  uki_patch=$1
  uki_checksum_file=$2
  uki_destination=$3

  # If the copy is interrupted, a temporary file prevents the partial image
  # from being used.  Note that this is required when we are updating the
  # UKI, because the destination directory is not a temporary directory.
  local tmp_uki_destination="${uki_destination}.tmp"

  uki_dir=$(dirname "$uki_patch")
  kernel_version_with_suffix=${uki_patch##*/kernel-}
  kernel_version=${kernel_version_with_suffix%.uki.efi.bsdiff}

  initrd_img_path="${uki_dir}/initrd.img-${kernel_version}"
  vmlinuz_path="${uki_dir}/vmlinuz-${kernel_version}"

  tmp_ukibase_path="${tmpdir}/ukibase-${kernel_version}"
  tmp_uki_path="${tmpdir}/kernel-${kernel_version}.uki.efi"
  if ! cat "$vmlinuz_path" "$initrd_img_path" > "$tmp_ukibase_path"; then
    echo 'error: could not create a base UKI to patch from' >&2
    return 1
  fi

  if ! bspatch "$tmp_ukibase_path" "$tmp_uki_path" "$uki_patch"; then
    echo "error: bspatch failed to patch with '${uki_patch}'" >&2
    return 1
  fi

  if ! (cd "$tmpdir" \
          && sha256sum -c --status "$uki_checksum_file" 2>/dev/null); then
    echo "error: checksum check of patched UKI '${uki}' failed" >&2
    return 1
  fi

  if ! cp "$tmp_uki_path" "$tmp_uki_destination" || \
     ! mv "$tmp_uki_destination" "$uki_destination"; then
    echo "error: failed to copy/move UKI '${uki}' to EFI partition" >&2
    return 1
  fi

  return 0
}

install_ukis() {
  local destination="$1"
  local enabled_uki_patch="$2"

  local class_file="${image_mountpoint}/etc/puavo-image/class"
  local image_name_file="${image_mountpoint}/etc/puavo-image/name"
  local release_file="${image_mountpoint}/etc/puavo-image/release"

  local class image_name release

  if ! class=$(cat "$class_file") || [ -z "$class" ]; then
    echo "error: image class missing in image mountpoint" >&2
    return 1
  fi
  if ! image_name=$(cat "$image_name_file") || [ -z "$image_name" ]; then
    echo "error: image name missing in image mountpoint" >&2
    return 1
  fi
  if ! release=$(cat "$release_file") || [ -z "$release" ]; then
    echo "error: release missing in image mountpoint" >&2
    return 1
  fi
  local image_name_no_suffix=${image_name%.img}

  local uki_final_installation_directory="${destination}/${image_name_no_suffix}"
  local uki_temporary_installation_directory="${uki_final_installation_directory}.tmp"
  local uki_installation_directory="${uki_final_installation_directory}"
  local is_new_installation=false

  if [ ! -d "$uki_installation_directory" ]; then
    # If we are installing for the first time, use a temporary installation
    # directory. Temporary files and directories are cleaned before and after
    # each installation.
    mkdir -p "$uki_temporary_installation_directory"
    uki_installation_directory="$uki_temporary_installation_directory"
    is_new_installation=true
  fi

  local status=0

  local uki_patch
  for uki_patch in "$boot_directory"/*.uki.efi.bsdiff; do
    local enabled="false"

    if [ "$(realpath "$uki_patch")" = "$(realpath "$enabled_uki_patch")" ]; then
      enabled="true"
    fi

    local uki="${uki_patch%.bsdiff}"
    local uki_checksum_file="${uki}.sha256"
    local uki_destination="${uki_installation_directory}/$(basename "$uki")"

    if ! (cd "$UKI_ROOT" \
            && sha256sum -c --status "$uki_checksum_file" 2>/dev/null); then
      # If the UKI already exists, it must differ from the new one
      # (e.g. possibly broken).  Remove the existing file before copying
      # the new one, otherwise the transfer would require double the
      # disk space.
      rm -f "$uki_destination" || true

      if ! patch_uki "$uki_patch" "$uki_checksum_file" "$uki_destination"; then
        echo "error: failed to patch UKI '${uki}' to EFI partition" >&2
        status=1
        continue
      fi
    else
      echo "UKI '$uki_destination' is already installed, skipping copy"
    fi

    local extra_directory="${uki_destination}.extra.d"

    puavo-update-uki-commandline "$extra_directory" \
      "$image_mountpoint" "$images_dir" "$hosttype" || {
      echo "error: failed to update UKI command-line for '$uki'" >&2
      status=1
      continue
    }

    install_boot_entry_config "$extra_directory" "$class" \
      "$image_name" "$release" "$enabled" || {
      echo "error: failed to install boot entry config for UKI '$uki'" >&2
      status=1
      continue
    }
  done

  if [ "$status" -ne 0 ]; then
    return "$status"
  fi

  # If we were installing for the first time and everything succeeded,
  # rename the temporary directory to the final location.
  if [ "$is_new_installation" = true ]; then
    if ! mv "$uki_temporary_installation_directory" \
            "$uki_final_installation_directory"; then
      echo "error: failed to rename temporary installation directory" >&2
      return 1
    fi
  fi
}

kernel_version=$(determine_selected_kernel_version)
[ -n "$kernel_version" ] || panic "failed to determine selected kernel version"

enabled_uki_patch=$(find "$boot_directory" -type f -name "*${kernel_version}.uki.efi.bsdiff" | head -n 1)
[ -n "$enabled_uki_patch" ] || panic "failed to find UKI with kernel version ${kernel_version}"

install_ukis "$UKI_ROOT" "$enabled_uki_patch" || {
  echo "error: failed to install UKI files to EFI partition" >&2
  exit 1
}
