#!/bin/sh

set -eu

images_dir=/images

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

grubenv_path="${images_dir}/boot/grub/grubenv"

PUAVO_INSTALL_LIBDIR='/usr/lib/puavo-ltsp-install'
UKI_ROOT='/boot/efi/EFI/puavo'

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

# TODO: Duplicated code
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 "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 >/dev/null
}

grubedit() {
  if ! grub-editenv "$grubenv_path" "$@"; then
    echo 'creating new grubenv because of an error' >&2
    grub-editenv "$grubenv_path" create
    grub-editenv "$grubenv_path" "$@"
  fi
}

setup_default_and_backup_images() {
  local backup_image_tag default_image_tag

  default_image_tag=$(find "$images_dir" -maxdepth 1 -type f -name '*.default' \
                        | head -n 1)
  backup_image_tag=$(find "$images_dir" -maxdepth 1 -type f -name '*.backup' \
                        | head -n 1)

  if [ -n "$default_image_tag" ]; then
    default_image_name=$(basename "$default_image_tag" .default)
    grubedit set "puavo_grub_puavo_os_default_image=${default_image_name}.img"
  else
    grubedit unset puavo_grub_puavo_os_default_image
  fi

  if [ -n "$backup_image_tag" ]; then
    backup_image_name=$(basename "$backup_image_tag" .backup)
    grubedit set "puavo_grub_puavo_os_backup_image=${backup_image_name}.img"
  else
    grubedit unset puavo_grub_puavo_os_backup_image
  fi
}

setup_developer_mode() {
  if [ "$(puavo-conf puavo.grub.developer_mode.enabled)" = 'true' ]; then
    grubedit set "puavo_show_imageoverlays=true"
  else
    grubedit unset puavo_show_imageoverlays
  fi
}

setup_grub_theme() {
  local grub_theme
  grub_theme=$(puavo-conf puavo.grub.theme)
  if [ -n "$grub_theme" ]; then
    grubedit set "puavo_grub_theme=${grub_theme}"
  else
    grubedit unset puavo_grub_theme
  fi
}

setup_kernel_arguments() {
  local kernel_arguments puavo_kernel_arguments selected_kernel_arguments

  if ! puavo_kernel_arguments="$(puavo-conf puavo.kernel.arguments)"; then
    echo 'failed to load puavo kernel arguments' >&2
    return 1
  fi

  if ! selected_kernel_arguments="$(selected_puavo_conf_arguments_as_kernel_args)"; then
    echo 'failed to lookup additional kernel arguments from puavo-conf' >&2
    return 1
  fi

  kernel_arguments="${puavo_kernel_arguments}${puavo_kernel_arguments:+ }${selected_kernel_arguments}"

  if [ -n "$kernel_arguments" ]; then
    grubedit set "puavo_kernel_arguments=${kernel_arguments}"
  else
    grubedit unset puavo_kernel_arguments
  fi

  if [ "$UKI_MODE" = 1 ]; then
    /usr/lib/puavo-core/puavo-install-uki-commandline-addon "$kernel_arguments" \
      || return 1
  fi
}

setup_kernel_version() {
  local kernel_version
  kernel_version=$(puavo-conf puavo.kernel.version)
  if [ -n "$kernel_version" ]; then
    grubedit set "puavo_kernel_version=${kernel_version}"
  else
    grubedit unset puavo_kernel_version
  fi
}

setup_locale() {
  local puavo_l10n_locale
  puavo_l10n_locale=$(puavo-conf puavo.l10n.locale)
  if [ -n "$puavo_l10n_locale" ]; then
    grubedit set "lang=${puavo_l10n_locale}"
  else
    grubedit unset lang
  fi
}

setup_locked() {
  local puavo_grub_locked
  puavo_grub_locked=$(puavo-conf puavo.grub.locked)
  if [ "$puavo_grub_locked" = 'true' ]; then
    grubedit set 'locked=true'
  else
    grubedit unset locked
  fi
}

setup_optional_grub_config() {
  local puavo_optional_grub_conf
  puavo_optional_grub_conf=$(puavo-conf puavo.grub.config)
  if [ -n "$puavo_optional_grub_conf" ]; then
    grubedit set "puavo_optional_grub_conf=${puavo_optional_grub_conf}"
  else
    grubedit unset puavo_optional_grub_conf
  fi
}

setup_timeout() {
  local grub_timeout
  grub_timeout=$(puavo-conf puavo.grub.timeout)
  if [ -n "$grub_timeout" ]; then
    grubedit set "puavo_grub_timeout=${grub_timeout}"
  else
    grubedit unset puavo_grub_timeout
  fi
}

setup_windows_enabled() {
  local grubdisk

  if [ "$(puavo-conf puavo.hosttype)" = 'bootserver' ]; then
    # Bootservers should not have Windows at all and usually use RAID
    # which has multiple boot disks (making this fail).
    return 0
  fi

  if ! grubdisk=$(/usr/lib/puavo-core/puavo-get-boot-disks); then
    echo 'failed to get the Grub disk(s)' >&2
    return 1
  fi

  if [ "$(printf %s "$grubdisk" | wc -w)" -ne 1 ]; then
    echo 'failed to find exactly one Grub disk' >&2
    return 1
  fi

  if [ "$(puavo-conf puavo.grub.windows.enabled)" = 'true' ]; then
    puavo-manage-efi --no-efi-is-ok enable-windows "$grubdisk" || return 1
    grubedit set "puavo_grub_windows_enabled=true" || return 1
  else
    grubedit unset puavo_grub_windows_enabled || return 1
    puavo-manage-efi --no-efi-is-ok disable-windows "$grubdisk" || return 1
  fi

  return 0
}

get_puavoconf_value() {
  # Do not use puavo-conf to look this up.  We only pass this to kernel
  # arguments only if it has been changed to non-default from Puavo.
  # At least puavo.* values from kernel arguments should be ignored here.
  jq -r --arg puavo_conf_key "$1" '.puavoconf[$puavo_conf_key]' \
    /state/etc/puavo/device.json 2>/dev/null || echo null
}

selected_puavo_conf_arguments_as_kernel_args() {
  local maybe_space puavo_conf_to_kernel_args puavo_conf_value is_encrypted
  maybe_space=''

  # Specify the root device for encrypted installations.
  # If the boot process requires user interaction (e.g. PIN),
  # the auto-discovery (systemd) for the root device may timeout.
  # By specifying the root device, we can bypass this timeout issue.
  if ! is_encrypted=$("$PUAVO_INSTALL_LIBDIR"/is-encrypted-install "$images_dir"); then
    echo "failed to detect installation encryption type" >&2
    return 1
  fi

  if [ "$is_encrypted" = "true" ]; then
    printf "%s" "root=/dev/mapper/root"
    maybe_space=' '
  fi

  puavo_conf_to_kernel_args='
    puavo.admin.persistent_homedirs
    puavo.boot.plymouth.theme
    puavo.graphics.driver
    puavo.image.overlay
  '

  for puavo_conf_key in $puavo_conf_to_kernel_args; do
    # "null" can not be a valid value for the above puavo-conf variables.
    puavo_conf_value=$(get_puavoconf_value "$puavo_conf_key")
    if [ "$puavo_conf_value" != 'null' ]; then
      printf "%s%s" "$maybe_space" "${puavo_conf_key}=${puavo_conf_value}"
      maybe_space=' '
    fi
  done
}

update_boot_device_environments() {
  local boot_disks status

  if [ "$UKI_MODE" != '1' ]; then
    return 0
  fi

  boot_disks=$(/usr/lib/puavo-core/puavo-get-boot-disks "$images_dir")
  status=0

  for diskdev in $boot_disks; do
    echo "installing GRUB environment to boot device ${diskdev}"
    if ! mount_disk_efi_partition "$diskdev"; then
      echo "failed to mount ${diskdev}" >&2
      status=1
      continue
    fi
    if ! cmp "$grubenv_path" "${UKI_ROOT}/grub/grubenv" >/dev/null 2>&1; then
      if ! cp -p "$grubenv_path" "${UKI_ROOT}/grub/grubenv.tmp"; then
        echo "failed to copy Grub environment to ${diskdev} (temporary file)" >&2
        status=1
      elif ! mv "${UKI_ROOT}/grub/grubenv.tmp" "${UKI_ROOT}/grub/grubenv"; then
        echo "failed to move the Grub environment file in ${diskdev}" >&2
        status=1
      fi
    fi
    umount /boot/efi || true
  done

  return $status
}

status=0

setup_default_and_backup_images  || status=1
setup_developer_mode             || status=1
setup_grub_theme                 || status=1
setup_kernel_arguments           || status=1
setup_kernel_version             || status=1
setup_locale                     || status=1
setup_locked                     || status=1
setup_optional_grub_config       || status=1
setup_timeout                    || status=1
setup_windows_enabled            || status=1
update_boot_device_environments  || status=1

exit $status
