#!/bin/bash

set -euo pipefail

. /usr/share/debconf/confmodule

SIZE_EXPECTED=536870912 # 512M
ESP_TYPE_GUID="c12a7328-f81f-11d2-ba4b-00a0c93ec93b"

# Functions modified from /usr/lib/grub/grub-multi-install, from grub2-unsigned
# Debian package
#
# Copyright: 2003, 2004, 2005, 2006, 2007, 2008, 2009, Robert Millan
#            2005, 2006, 2007, Otavio Salvador
#            2008, 2009, Felix Zielcke
#            2009, Jordi Mallach
# License: GPL-3+

# This only works on a Linux system with udev running.  This is probably the
# vast majority of systems where we need any of this, though, and we fall
# back reasonably gracefully if we don't have it.
cached_available_ids=
available_ids()
{
  local id path

  if [ "$cached_available_ids" ]; then
    echo "$cached_available_ids"
    return
  fi

  [ -d /dev/disk/by-id ] || return
  # LC_ALL=C: Ensure disk by ID has the expected order
  LC_ALL=C cached_available_ids="$(
    for path in /dev/disk/by-id/*; do
      [ -e "$path" ] || continue
      printf '%s %s\n' "$path" "$(readlink -f "$path")"
    done | sort -k2 -s -u | cut -d' ' -f1
  )"
  echo "$cached_available_ids"
}

# Returns non-zero and no output if no mapping can be found.
device_to_id()
{
  local id
  for id in $(available_ids); do
    if [ "$(readlink -f "$id")" = "$(readlink -f "$1")" ]; then
      echo "$id"
      return 0
    fi
  done
  # Fall back to the plain device name if there's no by-id link for it.
  if [ -e "$1" ]; then
    echo "$1"
    return 0
  fi
  return 1
}
# End of functions from /usr/lib/grub/grub-multi-install

is_protected() {
    dev="$(readlink -f "$1")"
    for mp in / /boot /boot/efi; do
        src="$(findmnt -n -o SOURCE "$mp" 2>/dev/null)" || continue
        [ "$(readlink -f "$src")" = "$dev" ] && return 0
    done
    return 1
}

derive_grub_target() {
    # /boot/grub/<target>/ is created by grub-install itself and
    # contains that platform's modules -- it's a direct artifact of
    # whatever install already succeeded on this unit.
    for d in /boot/grub/*-efi; do
        [ -d "$d" ] || continue
        basename "$d"
        return 0
    done
    return 1
}

# Step 0: detect all empty wannabe ESPs
match=()
for disk in /dev/nvme*n1; do
	part="$disk"p1

	if is_protected "$part"; then
		continue
	fi

	if [ -n "$(blkid -o value -s TYPE "$part")" ]; then
		continue
	fi

	size="$(blockdev --getsize64 "$part" 2>/dev/null)" || continue
	if [ "$size" != "$SIZE_EXPECTED" ]; then
		continue
	fi

	match+=("$disk")
done

if [ "${#match[@]}" -gt 1 ]; then
	#echo "more than 1 partitions match the criteria, backing off."
	exit 0
elif [ "${#match[@]}" -eq 0 ]; then
	#echo "no empty ESP partition."
	exit 0
fi

disk=${match[0]}
part=${match[0]}p1

part_by_id="$(device_to_id "$part")"

# Step 1: format missing ESP with FAT32 and set it as ESP
sgdisk --typecode="1:${ESP_TYPE_GUID}" "$disk"
udevadm settle --timeout=30 || true
mkfs.vfat -F 32 "$part"

# Step 2: Set up grub-efi to also install onto the new ESP
db_get grub-efi/install_devices
current="$RET"
if [ -n "$current" ]; then
    db_set grub-efi/install_devices "$current, $part_by_id"
else
    db_set grub-efi/install_devices "$part_by_id"
fi
db_fset grub-efi/install_devices seen true

# Step 3: Run grub-multi-install
grub_target="$(derive_grub_target)" || {
	echo "cannot derive grub target from current grub installation"
	exit 1
}
DEBIAN_FRONTEND=noninteractive /usr/lib/grub/grub-multi-install --target="$grub_target"

echo "Successfully installed GRUB onto secondary ESP: $part_by_id"
