#!/bin/bash

trap graceful_exit INT TERM

# Inofity events which signal something changed.
MODIFICATION_EVENTS="CREATE,CLOSE_WRITE,MOVED_TO,MOVED_FROM,MOVE"

# Signals if the watcher should stop.
STOP=false

graceful_exit () {
  STOP=true

  if [[ ! -z "$KPID" ]]; then
    kill -TERM $KPID >/dev/null 2>/dev/null
  fi
  if [[ ! -z "$MPID" ]]; then
    kill -TERM $MPID >/dev/null 2>/dev/null
  fi
}

load_config() {
  KERNEL=/boot/vmlinuz
  EFI_KERNEL=/boot/efi/EFI/Slackware/vmlinuz
  INITRD=/boot/initrd.gz
  EFI_INITRD=/boot/efi/EFI/Slackware/initrd.gz

  if [[ -f /etc/efi-sync.conf ]]; then
    source /etc/efi-sync.conf || exit 1
  fi

  INSTALL_PROGRAMS="slackpkg,upgradepkg,installpkg,removepkg"
  WAIT_TIME=1
}

log () {
  echo "$(date --rfc-3339=seconds) $@"
}

require_efi () {
  if ! mount | grep vfat >/dev/null; then
    log "the efi partition must be mounted"
    exit 1
  fi
}

require_root () {
  if [[ "$(id --user --name)" != "root" ]]; then
    log "command must be run as root"
    exit 1
  fi
}


efi_install () {
  require_root
  require_efi
  load_config

  KERNEL_VERSION="$(file $(realpath "$KERNEL") | sed -E 's/.* version (\S*) .*/\1/')"

  if ! cp -H "$KERNEL" "${EFI_KERNEL}"; then
    log "Failed to copy '${KERNEL}' into your EFI at '${EFI_KERNEL}'."
    return 1
  fi
  log  "Installed $(realpath "$KERNEL")."

  if [[ -f /etc/mkinitrd.conf ]]; then
    if ! mkinitrd -F -c -k "${KERNEL_VERSION}" >/dev/null 2>&1; then
      log "mkinitrd.conf detected but failed to build your initial ramdisk for '${KERNEL_VERSION}'."
      return 1
    fi
    log "mkinitrd.conf detected and successfully built your initial ramdisk for '${KERNEL_VERSION}'."
  fi

  if [[ -f "$INITRD" ]] && [[ "$INITRD" != "$EFI_INITRD" ]]; then
    if [[ "$INITRD" != "$EFI_INITRD" ]] && ! cp -H "$INITRD" "${EFI_INITRD}"; then
      log "Failed to copy '${INITRD}' into your EFI at '${EFI_INITRD}'."
      return 1
    fi
    log  "Installed the initial RAM disk '${INITRD}'."
  fi
}

efi_watch () {
  require_root
  require_efi
  load_config

  while ! $STOP; do
    log  "Watching ${KERNEL} for updates."

    >/dev/null inotifywait "$(dirname "$KERNEL")" --include "./$(basename $KERNEL)$" \
      --event "$MODIFICATION_EVENTS" \
      --quiet &
    KPID="$!"

    # Watching for directory updates is enough and helps avoid the system's watch limit.
    readarray -d '' INITRD_FILES < <(find /lib/firmware /lib/modules -type d -print0)
    >/dev/null inotifywait \
      --event "$MODIFICATION_EVENTS" \
      --quiet \
      --recursive \
      "${INITRD_FILES[@]}" &
    MPID="$!"

    if ! wait -n "$KPID" "$MPID" || $STOP; then
      kill -TERM $KPID $MPID >/dev/null 2>/dev/null
      continue
    fi
    kill -TERM $KPID $MPID >/dev/null 2>/dev/null
    unset KPID MPID

    # Be safe and consider changes to the kernel system which happen during
    # installation.
    #
    # Use file creation time to check if re-installation is needed.
    TIMESTAMP_FILE="$(mktemp -t "efi-sync.XXXXXXX")"

    while true; do
      touch "$TIMESTAMP_FILE"

      let "count = 0"
      while [[ $count -lt $WAIT_TIME ]]; do
        if ps -C "$INSTALL_PROGRAMS" >/dev/null; then
          let "count = 0"
	  sleep 1

	  continue
	fi

        let "count = $count + 1"
	sleep 1
      done

      efi_install

      NEW_FILES="$(find "$KERNEL" /lib/firmware /lib/modules -newer "$TIMESTAMP_FILE")"
      if [[ -z "$NEW_FILES" ]]; then
        break
      fi
    done

    rm "$TIMESTAMP_FILE"
  done

  log  "Stopped watching ${KERNEL} for updates."
}

case "$1" in
  'install')
    efi_install
    ;;
  'watch')
    efi_watch
    ;;
  *)
    echo "Usage: $0 [install|watch]"
    echo ""
    echo -e "\t- install: install the kernel and initrd into your EFI."
    echo -e "\t- watch: install on updates to your kernel"
    echo ""
esac
