#!/bin/bash -e

function usage {
  echo "$(basename "$0") [OPTIONS]"
  echo "  -h            shows usage"
  echo "  -g <version>  generate UKI image for specified kernel version"
  echo "  -a            generate UKI images for all available kernels"
}

function check_root {
  [ $EUID -eq 0 ] && return
  echo "dracut-initramfs requires root privileges to work" >&2
  exit 1
}

if [[ ${#} -eq 0 ]]; then
   usage
fi

ESP_PATH=$(bootctl --print-esp-path)
if [ -z "$ESP_PATH" ]; then
  exit 1
else

  if [ ! -d "${ESP_PATH}/EFI/Linux" ]; then
    if [ -d "${ESP_PATH}/EFI/linux" ]; then
      BOOT_PATH="${ESP_PATH}/EFI/linux"
    fi
  else
    BOOT_PATH="${ESP_PATH}/EFI/Linux"
  fi

  if [ -z "$BOOT_PATH" ]; then
    mkdir -p "${ESP_PATH}/EFI/Linux"
    BOOT_PATH="${ESP_PATH}/EFI/linux"
  fi

fi

declare -A kernels
update_all=0

while getopts ":hag:xyz" arg; do
  case ${arg} in
    g)
      found=0
      for line in $(pacman -Qql "$OPTARG"); do
        if [[ $line =~ ^/usr/lib/modules/([^/]+)/pkgbase$ ]]; then
          read -r kernel_name < "/${line}"
          kernels["${kernel_name}"]="${BASH_REMATCH[1]}"
          found=1
          break
        fi
      done
      if (( ! found )); then
        echo "Error occurred during '$OPTARG' package traversal" >&2
        exit 1
      fi
      ;;
    a)
      update_all=1
      ;;
    x)
      check_root
      # Trigger some IO on ESP path to be sure it's mounted by autofs if it's the case
      # Otherwise upgrading systemd may cause ESP partition not mounted at the time dracut attempt to write new image
      stat "$ESP_PATH" >/dev/null
      ;;
    y)
      check_root
      while read -r line; do
        if [[ "$line" == 'usr/lib/modules/'+([^/])'/pkgbase' ]]; then
          read -r kernel_name < "/${line}"
          path="$(grep -lE "^${kernel_name}\$" /usr/lib/modules/*/pkgbase)"
          kernel_version=$(basename "${path%/pkgbase}")

          kernel_image="$BOOT_PATH/vmlinuz-$kernel_name"
          initrd_image="$BOOT_PATH/initramfs-$kernel_name.img"
          initrd_image_fallback="$BOOT_PATH/initramfs-$kernel_name-fallback.img"

          declare -a images=()
          images+=("$kernel_image")
          images+=("$initrd_image")
          images+=("$initrd_image_fallback")

          for image in "${images[@]}";
          do
            if [ -f "$image" ]; then
              echo "==> Removing $image..."
              rm -f "$image"
            fi
          done
        fi
      done
      exit 0
      ;;
    z)
      check_root
      while read -r line; do
        if [[ $line =~ ^usr/lib/modules/([^/]+)/pkgbase$ ]]; then
          read -r kernel_name < "/${line}"
          kernels["${kernel_name}"]="${BASH_REMATCH[1]}"
        else
          update_all=1
          break 
        fi
      done
      ;;
    h)
      usage
      ;;
    *)
      usage
      ;;
  esac
done

if (( update_all )); then
  for kernel_path in /usr/lib/modules/*; do
    [ -f "$kernel_path/pkgbase" ] || continue

    kernel_version=$(basename $kernel_path)
    kernel_dir="/usr/lib/modules/$kernel_version"
    pkgbase="$kernel_dir/pkgbase"
    kernel_name=$(sed -e 's/^[[:space:]]//g' -e 's/[[:space:]]$//g' "$pkgbase")

    if [ -n "$kernel_name" ]; then
      kernels["$kernel_name"]="$kernel_version"
    fi
  done
fi

function gen_image() {
  check_root
  kernel_name="$1"
  kernel_version="$2"

  kernel_dir="/usr/lib/modules/$kernel_version"

  if [ -f "/etc/kernel/cmdline" ]; then
    cmdline=$(sed -e 's/^[[:space:]]//g' -e 's/[[:space:]]$//g' "/etc/kernel/cmdline")

  elif [ -f "/boot/refind_linux.conf" ]; then
    cmdline=$(cat "/boot/refind_linux.conf" | head -n 1 | sed -e 's|"Boot with standard options"||g' | awk '{$1=$1;print}' | sed -e 's/"//g')

  elif [ -f "/efi/refind_linux.conf" ]; then
    cmdline=$(cat "/efi/refind_linux.conf" | head -n 1 | sed -e 's|"Boot with standard options"||g' | awk '{$1=$1;print}' | sed -e 's/"//g')

  elif [ -f "/boot/efi/refind_linux.conf" ]; then
    cmdline=$(cat "/boot/efi/refind_linux.conf" | head -n 1 | sed -e 's|"Boot with standard options"||g' | awk '{$1=$1;print}' | sed -e 's/"//g')

  else
    cmdline=$(sed -e 's/^[[:space:]]//g' -e 's/[[:space:]]$//g' -e 's/initrd.*$//g' "/proc/cmdline")
  fi

  vmlinuz="$kernel_dir/vmlinuz"
  kernel_image="$BOOT_PATH/vmlinuz-$kernel_name"

  initrd_image="$BOOT_PATH/initramfs-$kernel_name.img"
  initrd_image_fallback="$BOOT_PATH/initramfs-$kernel_name-fallback.img"

  echo "==> Building initrd image $kernel_name $initrd_image"
  #dracut --force --hostonly --kver "$kernel_version" --kernel-cmdline="$cmdline" "$BOOT_PATH/initramfs-$kernel_name.img"
  #dracut --force --hostonly --no-hostonly-cmdline --kver "$kernel_version" "$BOOT_PATH/initramfs-$kernel_name.img"
  dracut --force --hostonly --no-hostonly-i18n --early-microcode --kernel-cmdline="$cmdline" --kver "$kernel_version" "$BOOT_PATH/initramfs-$kernel_name.img"

  echo "==> Building initrd image $kernel_name $initrd_image_fallback"
  #dracut --force --no-hostonly --kver "$kernel_version" --kernel-cmdline="$cmdline" "$BOOT_PATH/initramfs-$kernel_name-fallback.img"
  #dracut --force --kver "$kernel_version" "$BOOT_PATH/initramfs-$kernel_name-fallback.img"
  dracut --force --no-hostonly --no-hostonly-i18n --early-microcode --kernel-cmdline="$cmdline" --kver "$kernel_version" "$BOOT_PATH/initramfs-$kernel_name-fallback.img"

  cp -vrf $vmlinuz $kernel_image
}

for kernel_name in "${!kernels[@]}"; do
  kernel_version="${kernels[$kernel_name]}"
  gen_image "$kernel_name" "$kernel_version"
done