#!/bin/sh

set -eu

INSTALL_OPERATION_LOCK='/run/puavo-install.lock'
exec 9<> "$INSTALL_OPERATION_LOCK"
flock -n 9 || {
  echo "error: another installation operation is already running" >&2
  exit 1
}

# This script is run on image updates from puavo-image-preinst so that only
# /images is writable!  (and some other pseudo-filesystems such as /dev and
# /proc are accessible as well).

log() {
  priority=$1
  message=$2

  printf "> %s\n" "$message" >&2

  logger -p "$priority" -t puavo-install-grub "$message" || true
}

flush_caches() {
  # We have got some weird spurious grub errors of the sort (at first boot):
  # "error: ELF header smaller than expected", so we play it safe
  # and try weird tricks to flush all caches to actual, physical disk.
  # But we are not quite sure yet this will actually work (the problem is
  # difficult to reproduce).
  flush_status=0

  sync || flush_status=1
  echo 3 > /proc/sys/vm/drop_caches || flush_status=1

  for grubdev in "$@"; do
    blockdev --flushbufs "$grubdev" || flush_status=1
  done

  for grubdev in "$@"; do
    case "$grubdev" in
      /dev/[sv]d*)
	hdparm -qF "$grubdev" || flush_status=1
	;;
    esac
  done

  return $flush_status
}

get_grub_efi_arguments() {
  local vgname
  vgname=$1

  echo '--target=x86_64-efi --efi-directory=/boot/efi'

  if [ "$UKI_MODE" = '1' ]; then
    echo '--bootloader-id=puavo --uefi-secure-boot'
  else
    echo '--bootloader-id=debian --no-uefi-secure-boot'
  fi

  if [ "$vgname" = 'puavoinstaller' ]; then
    echo '--removable'
  elif [ "$(puavo-conf puavo.grub.efi.force_extra_removable)" = 'true' ]; then
    echo '--force-extra-removable'
  fi
}

grub_m4_template() {
  local hosttype
  hosttype=$1

  if [ "${UKI_MODE:-}" = "1" ]; then
    uki_grub_m4_template "$hosttype"
  else
    default_grub_m4_template "$hosttype"
  fi
}

default_grub_m4_template() {
  local hosttype
  hosttype=$1

  if [ "$hosttype" = 'exam' ]; then
    cat <<'EOF_GRUB_EXAM_CFG'
set default="0"
load_env

function load_video {
  if [ x$feature_all_video_module = xy ]; then
    insmod all_video
  else
    insmod efi_gop
    insmod efi_uga
    insmod ieee1275_fb
    insmod vbe
    insmod vga
    insmod video_bochs
    insmod video_cirrus
  fi
}

if loadfont unicode; then
  set gfxmode=auto
  load_video
  insmod gfxterm
  set locale_dir=$prefix/locale
  insmod gettext
fi
terminal_output gfxterm

set timeout=0

function get_pathname { regexp -s 2:"$2" '^(\(.*\))?(/.*)$' "$1"; }

function ltsp_image_entry {
  set imagefile="$1"
  set description="$2"

  get_pathname $imagefile imagepath

  # ${puavo_kernel_arguments} comes from a grub environment file
  # (if it is defined there)
  set kernelparameters="puavo.hosttype=exam"
  set boot_label="${description}"

  menuentry "${boot_label}" "$imagefile" "$imagepath" "$kernelparameters" {
    set imagefile="$2"
    set imagepath="$3"
    set kernelparameters="$4"

    set kernel_version_suffix=''

    loopback loop "${imagefile}"
    set root='(loop)'

    linux   /boot/vmlinuz${kernel_version_suffix} ro init=/sbin/init-puavo __PUAVO_KERNEL_ARGUMENTS__ ${kernelparameters}
    initrd  /boot/initrd.img${kernel_version_suffix}
  }
}

insmod gzio
insmod part_msdos
insmod ext2
insmod squash4
insmod regexp
insmod loopback
insmod usb_keyboard
insmod jpeg
insmod png
insmod btrfs

search --no-floppy --label --set puavobtrfsdev __PUAVO_IMAGES_BTRFS_LABEL__

for dev in __PUAVO_PARTITION_IMAGES__; do
  for examimage in ${dev}/ltsp.img; do
    if test -f "${examimage}"; then
      ltsp_image_entry "${examimage}" "Puavo exam environment"
    fi
  done
done
EOF_GRUB_EXAM_CFG
  else
    cat <<'EOF_GRUB_DEFAULT_CFG'
set default="0"
load_env

function load_video {
  if [ x$feature_all_video_module = xy ]; then
    insmod all_video
  else
    insmod efi_gop
    insmod efi_uga
    insmod ieee1275_fb
    insmod vbe
    insmod vga
    insmod video_bochs
    insmod video_cirrus
  fi
}

if loadfont unicode ; then
  set gfxmode=auto
  load_video
  insmod gfxterm
  set locale_dir=$prefix/locale
  insmod gettext
fi
terminal_output gfxterm

if keystatus --shift; then
  set timeout=-1
else
  if test -n "$puavo_grub_timeout"; then
    set timeout="$puavo_grub_timeout"
  else
    set timeout=5
  fi
fi

function get_imagename { regexp -s 2:"$2" '^(\(.*\)/)?(.*)$' "$1"; }
function get_pathname { regexp -s 2:"$2" '^(\(.*\))?(/.*)$' "$1"; }

function get_puavoimage_release_name {
  set imagefile="$1"

  set __puavoimage_release_name=""
  set puavoimage_release_name=""

  loopback image "$imagefile"
  load_env --file (image)/boot/grub/puavoimage_grubenv --skip-sig __puavoimage_release_name
  set puavoimage_release_name="${__puavoimage_release_name}"
  loopback -d image
}

function ltsp_image_entry {
  set imagefile="$1"
  set description="$2"
  set kernelparameters="$3"
  set boottype="$4"
  set image_alias="$5"
  set entry_os_icon="$6"

  set abitti=false

  get_imagename $image_alias image_alias_path
  get_pathname $imagefile imagepath
  get_puavoimage_release_name "$imagefile"

  if [ "$2" = "${abitti_boot_label}" ]; then
    set abitti=true
    set boot_label="${description} ($3)"
    set kernelparameters=""
    if test "${puavo_grub_boot_default}" = "abitti"; then
      set default="${boot_label}"
    fi
    set entry_os_icon="abitti"
  elif test -n "$puavoimage_release_name"; then
    set boot_label="${description} (${puavoimage_release_name})"
  else
    set boot_label="${description} (${image_alias_path})"
  fi

  # ${puavo_kernel_arguments} comes from a grub environment file
  # (if it is defined there)
  set kernelparameters="${kernelparameters} ${puavo_kernel_arguments}"
  if test "$boottype" = "exam"; then
    set kernelparameters="${kernelparameters} puavo.hosttype=exam"
  elif test "$boottype" = "liveboot"; then
    set kernelparameters="${kernelparameters} puavo.hosttype=laptop puavo.guestlogin.enabled=true"
  else
    set kernelparameters="${kernelparameters} puavo.hosttype=__PUAVO_HOSTTYPE__ puavo.guestlogin.enabled=true"
  fi

  if test "$imagepath" = "/${puavo_grub_puavo_os_default_image}" \
       -a "$boottype" = "$puavo_grub_puavo_os_default_mode"; then
    set default="${boot_label}"
  fi

  menuentry "$boot_label" --unrestricted --class="$entry_os_icon" "$imagefile" "$imagepath" "$kernelparameters" "$abitti" {
    set kernel_version=''
    # --- remove "#" from the following if you need another kernel ---
    # set kernel_version='default'
    # set kernel_version='crisp'

    set imagefile="$2"
    set imagepath="$3"

    loopback loop "${imagefile}"
    set root='(loop)'

    #
    # this mirrors code in puavo-ltspboot-config
    #

    if test "$kernel_version" = ""; then
      # ${puavo_kernel_version} comes from a grub environment file
      # (if it is defined there)
      set kernel_version="$puavo_kernel_version"
    elif test "$kernel_version" = "default"; then
      set kernel_version=""
    fi

    set kernel_version_suffix=''
    if test -n "${kernel_version}"; then
      if test -e "(loop)/boot/vmlinuz-${kernel_version}-amd64"; then
        set kernel_version_suffix="-${kernel_version}-amd64"
      elif test -e "(loop)/boot/vmlinuz-${kernel_version}"; then
        set kernel_version_suffix="-${kernel_version}"
      else
        # requested kernel not found, use default
        true
      fi

      # one more sanity test (also checks the corresponding initrd.img)
      if test -e "(loop)/boot/initrd.img${kernel_version_suffix}" \
           -a -e "(loop)/boot/vmlinuz${kernel_version_suffix}"; then
        true      # all is okay
      else
        # something is wrong, use the default kernel
        set kernel_version_suffix=''
      fi
    fi

    set kernelparameters="$4"
    set abitti="$5"

    if [ "$abitti" = "true" ]; then
      if test "$abitti_kernel_arguments" = ""; then
        set abitti_kernel_arguments="config digabidata digabi=syslinux union=overlay"
      fi
      linux /boot/vmlinuz${kernel_version_suffix} ro boot=live live-media=/dev/mapper/puavo-images live-media-path=/abitti puavo.abitti.boot=true puavo.boot.plymouth.theme=joy ${abitti_kernel_arguments} ${kernelparameters}
    else
      linux /boot/vmlinuz${kernel_version_suffix} ro init=/sbin/init-puavo __PUAVO_KERNEL_ARGUMENTS__ ${kernelparameters}
    fi

    initrd /boot/initrd.img${kernel_version_suffix}
  }
}

## Superlaptop mode = persistent image changes + local admins
function ltsp_image_entry_superlaptop {
  set imagefile="$1"
  set description="$2"
  set kernelparameters="$3"
  set image_alias="$4"
  set entry_superlaptop_os_icon="$5"

  if test "${puavo_show_imageoverlays}" = "true" \
            -a -d __PUAVO_PARTITION_OVERLAYS__; then
    set kernelparameters="$kernelparameters puavo.image.overlay=personal"

    ltsp_image_entry "${imagefile}" "${description}" "${kernelparameters}" "developer" "${image_alias}" "${entry_superlaptop_os_icon}"
  fi
}

insmod gzio
insmod part_msdos
insmod ext2
insmod squash4
insmod biosdisk
insmod regexp
insmod loopback
insmod usb_keyboard
insmod lvm
insmod jpeg
insmod png
insmod btrfs

ifelse(__PUAVO_IMAGES_BTRFS_LABEL__, `',
       `# using lvm',
       `# using btrfs
search --no-floppy --label --set puavobtrfsdev __PUAVO_IMAGES_BTRFS_LABEL__')

for dev in __PUAVO_PARTITION_IMAGES__; do
  for file in ${dev}/*.default; do
    regexp -s default_alias_no_suffix '^(.*)\.default$' "${file}"
    break
  done
  for file in ${dev}/*.backup; do
    regexp -s backup_alias_no_suffix '^(.*)\.backup$' "${file}"
    break
  done

  set default_image="${dev}/ltsp.img"
  set backup_image="${dev}/ltsp-backup.img"
  set default_image_alias="${default_alias_no_suffix}.img"
  set backup_image_alias="${backup_alias_no_suffix}.img"

  # Include optional grub configurations from the default image.
  loopback image "$default_image"
  if test -z "${puavo_optional_grub_conf}"; then
    set puavo_optional_grub_conf="default"
  fi
  if test -e "(image)/boot/grub/puavo/${puavo_optional_grub_conf}.cfg"; then
    source "(image)/boot/grub/puavo/${puavo_optional_grub_conf}.cfg"
  fi
  if [ "$locked" = "true" -a -e "(image)/boot/grub/puavo/superusers.cfg" ]; then
    source "(image)/boot/grub/puavo/superusers.cfg"
  fi
  loopback -d image

  # These might come from optional Grub configurations above,
  # so check if text_normal_startup is already set before setting these.
  # Note that some day we might drop the translations from here,
  # but should keep some fallback strings (in English) just in case
  # the optional configurations are missing from the default image.
  if test -z "$text_normal_startup"; then
    if test "$lang" = "de_CH.UTF-8" -o "$lang" = "de_DE.UTF-8"; then
      set text_normal_startup="Normaler Start"
      set text_liveboot_startup="Liveboot-Start"
      set text_developer_mode="Entwicklermodus"
      set text_backup_normal_startup="Backup-System, normaler Start"
      set text_backup_developer_mode="Backup-System, Entwicklermodus"
      set text_other_system_normal_startup="Anderes System, normaler Start"
      set text_other_system_developer_mode="Anderes System, Entwicklermodus"
      set text_abitti2_exam_application="Abitti-2-Prüflingsanwendung"
    elif test "$lang" = "fi_FI.UTF-8"; then
      set text_normal_startup="Normaali käynnistys"
      set text_liveboot_startup="Liveboot-käynnistys"
      set text_developer_mode="Kehittäjätila"
      set text_backup_normal_startup="Varajärjestelmä, normaali käynnistys"
      set text_backup_developer_mode="Varajärjestelmä, kehittäjätila"
      set text_other_system_normal_startup="Muu järjestelmä, normaali käynnistys"
      set text_other_system_developer_mode="Muu järjestelmä, kehittäjätila"
      set text_abitti2_exam_application="Abitti 2 -kokelassovellus"
    elif test "$lang" = "fr_CH.UTF-8"; then
      set text_normal_startup="Démarrage normal"
      set text_liveboot_startup="Démarrage live"
      set text_developer_mode="Mode développeur"
      set text_backup_normal_startup="Système de secours, démarrage normal"
      set text_backup_developer_mode="Système de secours, mode développeur"
      set text_other_system_normal_startup="Autre système, démarrage normal"
      set text_other_system_developer_mode="Autre système, mode développeur"
      set text_abitti2_exam_application="Application candidat Abitti 2"
    elif test "$lang" = "sv_FI.UTF-8"; then
      set text_normal_startup="Normal start"
      set text_liveboot_startup="Start av liveboot"
      set text_developer_mode="Utvecklarläge"
      set text_backup_normal_startup="Backup-system, normal start"
      set text_backup_developer_mode="Backup-system, utvecklarläge"
      set text_other_system_normal_startup="Annat system, normal start"
      set text_other_system_developer_mode="Annat system, utvecklarläge"
      set text_abitti2_exam_application="Abitti 2-examinandapplikationen"
    else
      set text_normal_startup="Normal startup"
      set text_liveboot_startup="Liveboot startup"
      set text_developer_mode="Developer mode"
      set text_backup_normal_startup="Backup system, normal startup"
      set text_backup_developer_mode="Backup-system, developer mode"
      set text_other_system_normal_startup="Other system, normal startup"
      set text_other_system_developer_mode="Other system, developer mode"
      set text_abitti2_exam_application="Abitti 2 Examinee Application"
    fi
  fi

  if test -z "$windows_boot_label"; then
    set windows_boot_label="Windows"
  fi

  if test -z "$abitti_boot_label"; then
    set abitti_boot_label="Koejärjestelmä"
  fi

  set os_icon="opinsys-os"

  if test -f "${default_image}"; then
    ltsp_image_entry "${default_image}" "${text_normal_startup}" "" "normal" "${default_image_alias}" "${os_icon}"
    if test "__PUAVO_LIVEBOOT_OPTION__" = "true"; then
      ltsp_image_entry "${default_image}" "${text_liveboot_startup}" "" "liveboot" "${default_image_alias}" "${os_icon}"
    fi
    ltsp_image_entry_superlaptop "${default_image}" "${text_developer_mode}" "" "${default_image_alias}" "${os_icon}-devel"

    loopback default_image_loop "${default_image}"
    set possible_puavo_grub_theme_path="(default_image_loop)/boot/grub/themes/${puavo_grub_theme}/theme.txt"
    if test -f "${possible_puavo_grub_theme_path}"; then
      set theme="${possible_puavo_grub_theme_path}"
    fi
  fi
  if test -f "${backup_image}"; then
    ltsp_image_entry "${backup_image}" "${text_backup_normal_startup}" "" "normal" "${backup_image_alias}" "${os_icon}-backup"
    ltsp_image_entry_superlaptop "${backup_image}" "${text_backup_developer_mode}" "" "${backup_image_alias}" "${os_icon}-devel"
  fi

  for systemimage in ${dev}/*.img; do
    if test "${systemimage}" = "${default_image}";       then continue; fi
    if test "${systemimage}" = "${backup_image}";        then continue; fi
    if test "${systemimage}" = "${default_image_alias}"; then continue; fi
    if test "${systemimage}" = "${backup_image_alias}";  then continue; fi

    ltsp_image_entry "${systemimage}" "${text_other_system_normal_startup}" "" "normal" "${systemimage}" "${os_icon}-backup"
    ltsp_image_entry_superlaptop "${systemimage}" "${text_other_system_developer_mode}" "" "${systemimage}" "${os_icon}-devel"
  done

  for examimage in ${dev}/exam/*.img; do
    if test -f "${examimage}"; then
      ltsp_image_entry "${examimage}" "Exam environment" "" "exam"
    fi
  done

  for file in ${dev}/abitti/abitti-v*.squashfs; do
    # If there are two abitti-images, do not show the menuentry because we
    # do not know what we get in the live-boot environment in this situation
    # (because we have no mechanism to choose a particular image).
    # This is mostly to stop admins doing unsupported things.
    if test -n "$abitti_image"; then
      set abitti_image=""
      break
    fi
    set abitti_image="${file}"
    if test -z "${abitti_version_number}"; then
      regexp -s abitti_version_number 'abitti-v(.*)\.squashfs$' "${file}"
    fi
  done
  if test -f "${abitti_image}"; then
    ltsp_image_entry "${default_image}" \
                     "${abitti_boot_label}"   \
                     "Abitti-yhteensopiva, versio ${abitti_version_number}"
  fi

  set is_default_abitti="true"
  for abitti_iso_path_marker in (lvm/puavo-images)/abitti2/*.default \
                                (lvm/puavo-images)/abitti2/*.backup; do
    regexp -s abitti_iso_base '^(.*)\.(default|backup)$' $abitti_iso_path_marker
    set abitti_iso_path="${abitti_iso_base}.iso"
    if [ ! -e "$abitti_iso_path" ]; then
      continue
    fi

    regexp -s abitti_iso_file '^.*/(.*\.iso)$' "${abitti_iso_path}"
    regexp -s abitti_iso_path_no_suffix '^(.*)\.iso$' "${abitti_iso_path}"
    set abitti_iso_cfg="${abitti_iso_path_no_suffix}.cfg"
    if [ ! -e "$abitti_iso_cfg" ]; then
      continue
    fi

    # sourcing the "${abitti_iso_base}.label" file should set abitti_label,
    # but fallback to $abitti_iso_file
    set abitti_label=""
    if [ -e "${abitti_iso_base}.label" ]; then
      source "${abitti_iso_base}.label"
    fi
    if [ "$abitti_label" = "" ]; then
      set abitti_label="$abitti_iso_file"
    fi

    set boot_label="${text_abitti2_exam_application} (${abitti_label})"
    if test "${is_default_abitti}" = "true" -a "${puavo_grub_boot_default}" = "abitti2"; then
      set default="$boot_label"
    fi
    set is_default_abitti="false"
    menuentry "$boot_label" --unrestricted --class="abitti2" "$abitti_iso_path" "$abitti_iso_file" "$abitti_iso_cfg" {
      set abitti_iso_path=$2
      set abitti_iso_file=$3
      set abitti_iso_cfg=$4
      export abitti_iso_path
      export abitti_iso_file
      configfile "$abitti_iso_cfg"
    }
  done
done
EOF_GRUB_DEFAULT_CFG
  fi
}

uki_grub_m4_template() {
  # XXX hosttype not used yet
  local hosttype
  hosttype=$1

  cat <<'EOF'
set default="0"

# Make sure the prefix is set correctly
if test -e "${prefix}/grub/grubenv"; then
  set prefix="${prefix}/grub"
fi

load_env

# Load the Puavo configuration
set text_normal_startup="Normal startup"
set text_backup_normal_startup="Backup system, normal startup"
set text_other_system_normal_startup="Other system, normal startup"

if test -z "${puavo_optional_grub_conf}"; then
  set puavo_optional_grub_conf="default"
fi
if test -e "${prefix}/puavo/${puavo_optional_grub_conf}.cfg"; then
  source "${prefix}/puavo/${puavo_optional_grub_conf}.cfg"
fi

function load_video {
  if [ x$feature_all_video_module = xy ]; then
    insmod all_video
  else
    insmod efi_gop
    insmod efi_uga
    insmod ieee1275_fb
    insmod vbe
    insmod vga
    insmod video_bochs
    insmod video_cirrus
  fi
}

if loadfont unicode ; then
  set gfxmode=auto
  load_video
  insmod gfxterm
  set locale_dir=$prefix/locale
  insmod gettext
fi
terminal_output gfxterm

if keystatus --shift; then
  set timeout=-1
else
  if test -n "$puavo_grub_timeout"; then
    set timeout="$puavo_grub_timeout"
  else
    set timeout=5
  fi
fi

function load_puavo_config {
  if test -z "$config_name"; then
    set config_name="default"
  fi

  set config_path="$prefix/puavo/${config_name}.cfg"

  if test -e "$config_path"; then
    source "$config_path"
  fi

  set theme_path="$prefix/themes/${puavo_grub_theme}/theme.txt"
  if test -f "$theme_path"; then
    set theme="${theme_path}"
  fi
}

function add_boot_entry {
  set image_path="$1"
  set label="$2"
  set icon="$3"

  menuentry "$label" --class "$icon" "$image_path" {
    set image_path="$2"

    search --no-floppy --fs-uuid --set=root __PUAVO_EFI_PARTITION_UUID__
    chainloader "$image_path"
  }
}

function generate_boot_entry {
  set image_path="$1"

  # Extract the dir and the basename of the image
  regexp -s 1:dir  '^(.*/)[^/]+$' "$image_path"
  regexp -s 1:name '/([^/]+)$'    "$image_path"

  # Look for image specific metadata directory next to the image file
  set metadata_dir="${dir}${name}.extra.d"

  # If such directory exists, load the image specific config
  set image_environment_file="${metadata_dir}/grubenv"

  set enabled="true"
  set image_name=""
  set release=""

  if test -f "$image_environment_file"; then
    load_env --file "$image_environment_file" --skip-sig enabled image_name release
  fi

  set label="${text_other_system_normal_startup} (${name})"
  set icon="puavo-os"

  # If the release is set, generate the label based on it
  if test -n "$release"; then
    if test "$image_name" = "${puavo_grub_puavo_os_default_image}"; then
      set label="${text_normal_startup} (${release})"
      set icon="puavo-os"
    elif test "$image_name" = "${puavo_grub_puavo_os_backup_image}"; then
      set label="${text_backup_normal_startup} (${release})"
      set icon="puavo-os-backup"
    else
      set label="${text_other_system_normal_startup} (${release})"
      set icon="puavo-os"
    fi
  fi

  if test "$enabled" = "false"; then
    return
  fi

  add_boot_entry "$image_path" "$label" "$icon"
}

function add_boot_options {
  search --no-floppy --fs-uuid --set=root __PUAVO_EFI_PARTITION_UUID__

  set default_image_no_suffix=''
  if test -n "$puavo_grub_puavo_os_default_image"; then
    regexp -s default_image_no_suffix '^(.*)\.img$' \
              "$puavo_grub_puavo_os_default_image"
    for image_uki_path in /EFI/puavo/${default_image_no_suffix}/*.uki.efi; do
      if ! test -f "$image_uki_path"; then continue; fi
      generate_boot_entry "$image_uki_path"
    done
  fi

  set backup_image_no_suffix=''
  if test -n "$puavo_grub_puavo_os_backup_image"; then
    regexp -s backup_image_no_suffix '^(.*)\.img$' \
              "$puavo_grub_puavo_os_backup_image"
    for image_uki_path in /EFI/puavo/${backup_image_no_suffix}/*.uki.efi; do
      if ! test -f "$image_uki_path"; then continue; fi
      generate_boot_entry "$image_uki_path"
    done
  fi

  for image_uki_path in /EFI/puavo/*/*.uki.efi; do
    if ! test -f "$image_uki_path"; then continue; fi
    regexp -s 1:image_uki_dir '^(.*)/[^/]+$' "$image_uki_path"
    if test "$image_uki_dir" = "/EFI/puavo/${default_image_no_suffix}" \
         -o "$image_uki_dir" = "/EFI/puavo/${backup_image_no_suffix}"; then
      continue  # already handled
    fi
    generate_boot_entry "$image_uki_path"
  done

  # Add entry for Windows if it exists
  set windows_image_path="(efiroot)/EFI/Microsoft/Boot/bootmgfw.efi"

  if test -z "$windows_boot_label"; then
    set windows_boot_label="Windows"
  fi

  if [ -f "$windows_image_path" ]; then
    add_boot_entry "$windows_image_path" "$windows_boot_label" "windows"
  fi
}

insmod fat
insmod search_fs_uuid
insmod chain
insmod part_gpt
insmod jpeg
insmod png
insmod regexp

load_puavo_config
add_boot_options
EOF
}

find_grub_target_devices_on_raid() {
  local diskdev

  diskdev=$1
  mdadm --detail "$diskdev" \
    | awk '   $1 ~ /^[0-9]+$/ \
           && $2 ~ /^[0-9]+$/ \
           && $3 ~ /^[0-9]+$/ {
             split($0, a)
             for (i in a) {
               if (a[i] ~ /^\/dev\/[a-z]+[0-9]+$/) {
                  sub(/[0-9]+$/, "", a[i])
                  print a[i]
               }
             }
           }'
}

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

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

has_bios_boot_partition() {
  local diskdev
  diskdev=$1
  # 21686148-6449-6e6f-744e-656564454649 == BIOS boot partition type
  lsblk -n -l -o PARTTYPE "$diskdev" \
    | awk '
        BEGIN { status = 1 }
        $1 == "21686148-6449-6e6f-744e-656564454649" { status = 0 }
        END { exit(status) }
      '
}

lookup_puavo_btrfs_diskdev() {
  # XXX this should make sure it is returning the correct one in case
  # XXX two btrfs filesystems are both with a "puavo"-label (or can we do
  # XXX anything about this?)
  local label
  label=$1
  lsblk -n -l -o FSTYPE,LABEL,PKNAME \
    | awk -v label="$label" '
        $1 == "btrfs" && $2 == label && $3 != "" { print "/dev/" $3; exit(0) }
      '
}

lookup_puavo_btrfs_partdev() {
  local btrfs_diskdev btrfs_label
  btrfs_diskdev=$1
  btrfs_label=$2
  lsblk -n -o FSTYPE,LABEL,UUID "$btrfs_diskdev" | awk -v label="$btrfs_label" '
    BEGIN { status = 1 }
    $1 == "btrfs" && $2 == label { print "UUID=" $3; status = 0; exit }
    END { exit(status) }
  '
}

lookup_puavo_vg_diskdev() {
  local vgname
  vgname=$1
  # pvs and assorted utilities complain if we leak file descriptors, that is
  # why we do "3>&- 4>&- 5>&-" (same with grub-install)
  pvs 3>&- 4>&- 5>&- | awk -v vgname="$vgname" '$2 == vgname { print $1; exit }'
}

lookup_windows_boot_partition() {
  local diskdev
  diskdev=$1

  # "ebd0a0a2-b9e5-4433-87c0-68b6b72699c7" is Windows partition type.
  # Restrict this check to $diskdev only, so that possible matching
  # partitions in external drives do not mess up Windows boot.
  # It might be nice to support dual boot with a separate disk for Windows,
  # but this does do not do that.
  lsblk -n -l -o PARTTYPE,UUID "$diskdev" \
    | awk '
        BEGIN { efi_uuid = ""; legacy_uuid = ""; has_windows = 0 }
        $1 == "0x7" {
          has_windows = 1
          if (!legacy_uuid) { legacy_uuid = $2 }
          }
          $1 == "c12a7328-f81f-11d2-ba4b-00a0c93ec93b" { efi_uuid = $2 }
          $1 == "ebd0a0a2-b9e5-4433-87c0-68b6b72699c7" { has_windows = 1 }
          END {
            if (has_windows) {
              if (efi_uuid && efi_uuid ~ /^[a-zA-Z0-9-]+$/) {
                print efi_uuid
              } else if (legacy_uuid && legacy_uuid ~ /^[a-zA-Z0-9-]+$/) {
                print legacy_uuid
              }
            }
          }
        '
}

lookup_windows_chainloader() {
  local diskdev efi_devpath
  diskdev=$1
  efi_devpath=$(get_efi_devpath "$diskdev")
  if [ -n "$efi_devpath" ]; then
    echo '/EFI/Microsoft/Boot/bootmgfw.efi'
  else
    echo '+1'
  fi
}

mount_efi_partition() {
  local efi_devpath
  efi_devpath=$1

  mountpoint -q /boot/efi && return 0

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

install_legacy_grub() {
  local boot_dir grubdev

  boot_dir=$1
  grubdev=$2

  if ! grub_msg=$(grub-install --boot-directory="$boot_dir" \
                    --target=i386-pc "$grubdev" 3>&- 4>&- 5>&- 2>&1); then
    log err "Grub (MBR) installation to ${grubdev} failed: ${grub_msg}"
    return 1
  fi

  log info "Grub (MBR) installation to ${grubdev} was OK."

  return 0
}

copy_grub_environment() {
  # Scripts such as "setup_grub_environment" write the
  # environment variables into /images/boot/grub/grubenv.
  # Here we sync this file to the boot directory in the EFI partition.
  if [ ! -f "${images_dir}/boot/grub/grubenv" ]; then
    # Grub will create the default environment in the EFI partition
    return 0
  fi

  cp "${images_dir}/boot/grub/grubenv" "${boot_dir}/grub/grubenv" || return 1
}

write_grub_version() {
  local output_file
  output_file=$1
  new_grub_version=$(grub-install --version | awk '{ print $3 }')
  if [ -z "$new_grub_version" ]; then
    log err 'could not determine new grub version'
    exitstatus=1
  else
    printf "%s\n" "$new_grub_version" > "${output_file}.tmp"
    mv "${output_file}.tmp" "$output_file"
    log notice "Grub version ${new_grub_version} has been installed"
  fi
}

install_efi_grub() {
  local boot_dir efi_devpath grubdev grub_efi_arguments vgname

  boot_dir=$1
  efi_devpath=$2
  grubdev=$3
  vgname=$4

  if ! mount_efi_partition "$efi_devpath"; then
    log err "could not mount /boot/efi from ${efi_devpath}"
    return 1
  fi

  # We might actually have leftover nvram dumps from previous installations
  # that prevent a successful bootloader installation. Clean them up first.
  rm -f /sys/firmware/efi/efivars/dump-*

  grub_efi_arguments=$(get_grub_efi_arguments "$vgname")
  if ! grub_msg=$(grub-install --boot-directory="$boot_dir" \
                    $grub_efi_arguments "$grubdev" \
                    3>&- 4>&- 5>&- 2>&1); then
    log err "Grub (UEFI) installation to ${grubdev} failed: ${grub_msg}"
    umount /boot/efi || true
    return 1
  fi

  if [ "${UKI_MODE:-}" = "1" ]; then
    install_puavo_grub_files_to_efi_partition "$boot_dir" && \
    copy_grub_environment && \
    write_grub_version "${boot_dir}/grub/version" || {
      log err "failed to install UKI files to EFI partition"
      umount /boot/efi || true
      return 1
    }
  fi

  umount /boot/efi || true

  log info "Grub (UEFI) installation to ${grubdev} was OK."

  return 0
}

usage() {
  cat <<EOF
$(basename $0) [optional arguments]

  $(basename $0) takes the following optional arguments:

    --hosttype           hosttype
    --images-dir         imagesdir
    --only-update-config
    --vgname             vgname ("puavo", "puavoexam", "puavoinstaller")
EOF
  exit 1
}

log notice 'starting to update Grub'

exitstatus=0

hosttype=$(cat /etc/puavo/hosttype 2>/dev/null || true)
images_dir=/images
only_update_config=false
vgname=puavo

if ! args=$(getopt -n "$0" -o +                                          \
                   -l 'hosttype:,images-dir:,only-update-config,vgname:' \
                   -- "$@"); then
  usage
fi

eval "set -- $args"
while [ $# -ne 0 ]; do
  case "$1" in
    --hosttype)           hosttype=$2;             shift; shift ;;
    --images-dir)         images_dir=$2;           shift; shift ;;
    --only-update-config) only_update_config=true; shift        ;;
    --vgname)             vgname=$2;               shift; shift ;;
    --) shift; break ;;
    *)  usage ;;
  esac
done

[ $# -eq 0 ] || usage

if [ "$vgname" != 'puavo' ] && \
   [ "$vgname" != 'puavoexam' ] && \
   [ "$vgname" != 'puavoinstaller' ]; then
  echo "vgname should be 'puavo', 'puavoexam' or 'puavoinstaller'" >&2
  usage
  exit 1
fi

if [ ! -d "$images_dir" ]; then
  log err "'${images_dir}' is not a directory"
  exit 1
fi

if [ -z "$hosttype" ]; then
  log err 'hosttype is not set'
  exit 1
fi

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

boot_dir="${images_dir}/boot"

btrfs_diskdev=$(lookup_puavo_btrfs_diskdev "$vgname")
if [ -n "$btrfs_diskdev" ]; then
  diskdev="$btrfs_diskdev"
else
  diskdev=$(lookup_puavo_vg_diskdev "$vgname")
fi
if [ -z "$diskdev" ]; then
  log err "could not find the disk device where volume group '${vgname}' is"
  exit 1
fi

# if the disk device is a luks device, we need to select the containing
# partition device instead. we assume no nesting.
if command -v cryptsetup > /dev/null 2>&1; then
  parentdev=$(lsblk -n -l -p -o NAME,PKNAME \
                | awk -v disk="$diskdev" '$1 == disk { print $2; exit }')
  if cryptsetup isLuks "$parentdev" > /dev/null 2>&1; then
    diskdev=$parentdev
  fi
fi

# remove the partition number
diskdev=$(echo $diskdev \
    | sed -E '
        /\/dev\/loop[0-9]+p/   { s|p[0-9]+$|| }
        /\/dev\/mmcblk/        { s|p[0-9]+$|| }
        /\/dev\/md/            { s|p[0-9]+$|| }
        /\/dev\/nvme/          { s|p[0-9]+$|| }
        /\/dev\/[sv]d/         { s|[0-9]+$||  }
        /\/dev\/xvd/           { s|[0-9]+$||  }
    ')

possible_windows_disk_dev=''
case "$diskdev" in
  /dev/md*)
    grub_target_devices=$(find_grub_target_devices_on_raid "$diskdev")
    if [ -z "$grub_target_devices" ]; then
      log err 'could not find Grub target devices on RAID configuration'
      exit 1
    fi
    ;;
  *)
    grub_target_devices="$diskdev"
    possible_windows_disk_dev="$diskdev"
    ;;
esac

if [ "$hosttype" = 'diskinstaller' ]; then
  puavo_liveboot_option='true'
else
  puavo_liveboot_option='false'
fi

if [ -n "$btrfs_diskdev" ]; then
  puavo_images_btrfs_label="$vgname"
  if ! puavo_partition_dev="$(lookup_puavo_btrfs_partdev "$btrfs_diskdev" \
                                "$puavo_images_btrfs_label")"; then
    log err 'could not find puavo btrfs partition device UUID'
    exit 1
  fi
  puavo_partition_grub_overlays='(${puavobtrfsdev})/imageoverlays'
  case "$hosttype" in
    diskinstaller) puavo_partition_grub_images='(${puavobtrfsdev})/installimages' ;;
    exam)          puavo_partition_grub_images='(${puavobtrfsdev})/examimages'    ;;
    *)             puavo_partition_grub_images='(${puavobtrfsdev})/images'        ;;
  esac
else
  puavo_images_btrfs_label=''
  case "$hosttype" in
    diskinstaller)
      puavo_partition_dev='/dev/mapper/puavoinstaller-installimages'
      puavo_partition_grub_images='(lvm/puavoinstaller-installimages)'
      ;;
    exam)
      puavo_partition_dev='/dev/mapper/puavoexam-examimages'
      puavo_partition_grub_images='(lvm/puavoexam-examimages)'
      ;;
    *)
      puavo_partition_dev='/dev/mapper/puavo-images'
      puavo_partition_grub_images='(lvm/puavo-images)'
      ;;
  esac
  puavo_partition_grub_overlays='(lvm/puavo-imageoverlays)'
fi

# XXX ${puavo_partition_dev} should use UUID at least in the btrfs case:
# XXX root=UUID=${uuid}
common_kernel_arguments="root=${puavo_partition_dev} puavo.image.path=\"\${imagepath}\""

case "$hosttype" in
  bootserver|diskinstaller|preinstalled|wirelessaccesspoint)
    puavo_kernel_arguments="nosplash ${common_kernel_arguments}"
    ;;
  exam)
    puavo_kernel_arguments="quiet loglevel=3 splash ${common_kernel_arguments} puavo.image.load_to_ram=true preempt=full"
    ;;
  laptop)
    puavo_kernel_arguments="quiet loglevel=3 splash ${common_kernel_arguments} preempt=full"
    ;;
  *)
    log err "hosttype '${hosttype}' is not supported"
    exit 1
    ;;
esac

install_puavo_grub_files_to_efi_partition() {
  boot_dir=$1

  cp -r /boot/grub "$boot_dir" || (
    log err "failed to copy Grub files to boot directory ${boot_dir}"
    return 1
  )
}

if [ "${UKI_MODE:-}" = "1" ]; then
  # Many scripts assume there's a GRUB boot bolder under images.
  # For example, the GRUB environment is written there and then
  # copied to all EFI partitions.
  mkdir -p "$boot_dir"/grub || (
    log err "failed to create boot directory at ${boot_dir}/grub"
    exit 1
  )

  boot_dir='/boot/efi/EFI/puavo'
fi

grub_cfg_path="${boot_dir}/grub/grub.cfg"

grub_install_done=false

if $only_update_config; then
  log info 'only updating Grub configuration'
else
  # UKI-based Grub installation doesn't support legacy GRUB,
  # because we'll install the GRUB boot files on EFI partition
  if [ "${UKI_MODE:-}" = "1" ]; then
    for grubdev in $grub_target_devices; do
      log info "installing Grub to ${grubdev}"

      efi_devpath=$(get_efi_devpath "$grubdev")
      if [ -z "$efi_devpath" ]; then
        log err "failed to find EFI partition path for Grub device"
        exit 1
      fi

      if ! install_efi_grub "$boot_dir" "$efi_devpath" "$grubdev" "$vgname"; then
        exit 1
      fi

      log info "Grub installation to ${grubdev} was successful"
    done
  else
    for grubdev in $grub_target_devices; do
      log info "installing Grub to ${grubdev}"

      grub_install_attempted=false

      if has_bios_boot_partition "$grubdev"; then
        grub_install_attempted=true
        if ! install_legacy_grub "$boot_dir" "$grubdev"; then
          exit 1
        fi
      fi

      efi_devpath=$(get_efi_devpath "$grubdev")
      if [ -n "$efi_devpath" ]; then
        grub_install_attempted=true
        if ! install_efi_grub "$boot_dir" "$efi_devpath" "$grubdev" "$vgname"; then
          exit 1
        fi
      fi

      if ! $grub_install_attempted; then
        # Must be a legacy Grub installation in case neither bios boot
        # partition nor EFI-partition was found.  Thus, we must try this:
        if ! install_legacy_grub "$boot_dir" "$grubdev"; then
          exit 1
        fi
      fi

      log info "Grub installation to ${grubdev} was successful"
    done

    write_grub_version "${boot_dir}/grub/version"
  fi

  grub_install_done=true

  if ! flush_caches; then
    log warn 'error in flushing caches after grub installation'
  fi
fi

# Next, we'll install the GRUB configuration
tmp_grub_cfg_path="${grub_cfg_path}.tmp"

generate_grub_configuration() {
  local puavo_efi_partition_uuid
  puavo_efi_partition_uuid=$1

  grub_m4_template "$hosttype"                                           \
    | m4 -D__PUAVO_EFI_PARTITION_UUID__="$puavo_efi_partition_uuid"      \
         -D__PUAVO_HOSTTYPE__="$hosttype"                                \
         -D__PUAVO_IMAGES_BTRFS_LABEL__="$puavo_images_btrfs_label"      \
         -D__PUAVO_KERNEL_ARGUMENTS__="$puavo_kernel_arguments"          \
         -D__PUAVO_LIVEBOOT_OPTION__="$puavo_liveboot_option"            \
         -D__PUAVO_PARTITION_IMAGES__="$puavo_partition_grub_images"     \
         -D__PUAVO_PARTITION_OVERLAYS__="$puavo_partition_grub_overlays" \
    > "$tmp_grub_cfg_path"

  windows_partition_path=/images/boot/.puavo_windows_partition
  if [ -n "$possible_windows_disk_dev" ]; then
    windows_boot_partition=$(
      lookup_windows_boot_partition "$possible_windows_disk_dev")
    windows_chainloader=$(lookup_windows_chainloader "$possible_windows_disk_dev")
    if [ -n "$windows_boot_partition" -a -n "$windows_chainloader" ]; then
      printf "%s %s\n" "$windows_boot_partition" "$windows_chainloader" \
        > "${windows_partition_path}.tmp"
      mv "${windows_partition_path}.tmp" "$windows_partition_path"

      cat <<EOF >> "$tmp_grub_cfg_path"

  if test "\${puavo_grub_windows_enabled}" = "true"; then
    menuentry "\${windows_boot_label}" --unrestricted --class="windows" {
      insmod ntfs
      search --no-floppy --fs-uuid --set=root ${windows_boot_partition}
      chainloader ${windows_chainloader}
    }
    if test "\${puavo_grub_boot_default}" = "windows"; then
      set default="\${windows_boot_label}"
    fi
  fi
EOF
    else
      rm -f "$windows_partition_path"
    fi
  else
    rm -f "$windows_partition_path"
  fi
}

install_grub_configuration() {
  if cmp "$grub_cfg_path" "$tmp_grub_cfg_path" >/dev/null 2>&1; then
    log info 'no changes to Grub configuration'
  else
    mv "$tmp_grub_cfg_path" "$grub_cfg_path"
    log notice 'new Grub configuration has been installed'
    if ! flush_caches; then
      log warn 'error in flushing caches after updating Grub configuration'
    fi
  fi

  # We no longer need the temporary config file
  rm -f "$tmp_grub_cfg_path"
}

if [ "${UKI_MODE:-}" = "1" ]; then
  # Install the grub config to all EFI partitions
  for grubdev in $grub_target_devices; do
    log info "updating the Grub configuration on ${grubdev}"

    efi_devpath=$(get_efi_devpath "$grubdev")
    if [ -z "$efi_devpath" ]; then
      log err "failed to find EFI partition path for Grub device"
      exit 1
    fi

    # We need to mount the EFI partition, because the
    # GRUB config will be installed there
    if ! mount_efi_partition "$efi_devpath"; then
      log err "could not mount /boot/efi from ${efi_devpath}"
      return 1
    fi

    # We need separate GRUB configs for each EFI partition,
    # because each config is dependent on the partition's UUID
    efi_partition_uuid=$(get_efi_uuid "$grubdev")
    generate_grub_configuration "$efi_partition_uuid"
    install_grub_configuration

    umount /boot/efi || true
  done
else
  generate_grub_configuration ""
  install_grub_configuration
fi

if $grub_install_done; then
  # Do an extra sleep in case we really installed grub (and have not only
  # updated the configuration), in case cache flushing was not adequate.
  sleep 5
fi

exit $exitstatus
