#!/bin/bash

set -eu

UKI_ROOT='/boot/efi/EFI/puavo'
INSTALL_OPERATION_LOCK='/run/puavo-install.lock'

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

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

[ $# -eq 3 ] || usage
image_path="$1"
hosttype="$2"
images_dir="$3"

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

exec {lock_fd}<> "$INSTALL_OPERATION_LOCK"
flock -nx "$lock_fd" || panic "another installation operation is already running"

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
}

image_mountpoint=$(mktemp -d)
cleanup() {
  if [ -d "$UKI_ROOT" ]; then
    cleanup_temporary_files || echo "warning: failed to remove temporary files" >&2
  fi

  umount /boot/efi 2>/dev/null || true
  umount "$image_mountpoint" 2>/dev/null || true
}
trap cleanup INT TERM EXIT

mount -o ro "$image_path" "$image_mountpoint" 1>/dev/null

# Load the image name without suffix for determining the installation location
image_name_file="${image_mountpoint}/etc/puavo-image/name"

if ! image_name=$(cat "$image_name_file") || [ -z "$image_name" ]; then
  panic "failed to extract image name from '$image_name_file'"
fi

image_basename=$(basename "$image_path")
if [ "$image_basename" != "$image_name" ]; then
  # The boot process expects the image filename to match the internal image name (see UKI command-line generation)
  panic "image filename '${image_basename}' does not match internal image name '${image_name}'"
fi

image_name_no_suffix=${image_name%.img}

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
    echo "error: failed to find EFI partition in ${diskdev}" >&2
    return 1
  fi

  fsck.fat -a "$efi_devpath" >/dev/null 2>&1 || true
  mount "$efi_devpath" /boot/efi 1>/dev/null
}

estimate_uki_size() {
  local file_sizes initrd_img_path kernel_version kernel_version_with_suffix \
        uki_patch vmlinuz_path

  uki_patch=$1
  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}"

  file_sizes=$(stat -c %s "$initrd_img_path" "$uki_patch" "$vmlinuz_path")
  printf %s "$file_sizes" | awk '{ sum += $1 } END { print sum }'
}

estimate_required_disk_space() {
  local existing_size required_disk_space safety_margin uki_destination \
        uki_kernels_dir uki_patch uki_size_estimate
  required_disk_space=0

  uki_kernels_dir="${UKI_ROOT}/${image_name_no_suffix}"

  for uki_patch in "$image_mountpoint"/boot/*.uki.efi.bsdiff; do
    test -e "$uki_patch" || continue
    uki_size_estimate=$(estimate_uki_size "$uki_patch")
    uki_destination="${uki_kernels_dir}/$(basename "$uki_patch" .bsdiff)"
    if [ -e "$uki_destination" ]; then
      existing_size=$(stat -c %s "$uki_destination")
      required_disk_space=$((required_disk_space + uki_size_estimate - existing_size))
    else
      required_disk_space=$((required_disk_space + uki_size_estimate))
    fi
  done

  safety_margin=$((32 * 1024 * 1024)) # 32 MB
  required_disk_space=$((required_disk_space + safety_margin))

  echo "$required_disk_space"
}

has_enough_disk_space() {
  local diskdev="$1"
  local required_space="$2"
  local available_space
  available_space=$(df --output=avail -B1 /boot/efi | tail -n1)
  [ "$available_space" -ge "$required_space" ]
}

install_uki_to_efi() {
  if [ ! -d "$UKI_ROOT" ]; then
    echo "error: boot device does not have a valid EFI partition" >&2
    return 1
  fi

  puavo-install-uki "$image_mountpoint" "$images_dir" "$hosttype" || {
    echo "error: failed to install UKI to EFI partition" >&2
    return 1
  }
}

mount_efi_and_install_uki() {
  local diskdev="$1"
  local status=0
  mount_disk_efi_partition "$diskdev" || return 1
  install_uki_to_efi "$diskdev" || status=1
  cleanup_temporary_files || echo "warning: failed to remove temporary files" >&2
  umount /boot/efi || true
  return $status
}

# Identify boot disks and install UKI files to their EFI partitions
boot_disks=$(/usr/lib/puavo-core/puavo-get-boot-disks "$images_dir")
[ -n "$boot_disks" ] || panic "failed to find boot disks"

# Before installing, verify the boot partitions have enough space
for diskdev in $boot_disks; do
  mount_disk_efi_partition "$diskdev" || panic "failed to mount EFI partition on $diskdev"
  cleanup_temporary_files || echo "warning: failed to remove temporary files" >&2

  required_space=$(estimate_required_disk_space "$image_path")

  [ -n "$required_space" ] || panic "failed to estimate required disk space"
  has_enough_disk_space "$diskdev" "$required_space" || panic "disk $diskdev does not have enough space"

  umount /boot/efi || true
done

status=0

for diskdev in $boot_disks; do
  echo "installing UKI files to boot device $diskdev"
  mount_efi_and_install_uki "$diskdev" || status=1
done

exit $status
