#!/bin/bash --norc

kver="$(uname -r)"

boot_dir="/boot"
quiet=0
host_only=0
force=0
set_hostonly=0
no_hostonly=0
dracut_dbg=0

error() {
  echo "$@" >&2
}

usage() {
  local cmd=""

  [[ $1 = '-n' ]] && cmd=echo || cmd=error

  $cmd "usage: ${0##*/} [--version] [--help] [-v] [-f] [--preload <module>]"
  $cmd "       [--image-version] [--with=<module>]"
  $cmd "       [--nocompress]"
  $cmd "       [--set-hostonly] - Generates host specific initrd.img"
  $cmd "       [--no-hostonly] - Generates generic initrd.img"
  $cmd "       [--debug]"
  $cmd "       <initrd-image> <kernel-version>"
  $cmd ""
  $cmd "       (ex: ${0##*/} /boot/initramfs-$kver.img $kver)"

  [[ $1 = '-n' ]] && exit 0
  exit 1
}

# Little helper function for reading args from the commandline.
# it automatically handles -a b and -a=b variants, and returns 1 if
# we need to shift $3.
read_arg() {
  # $1 = arg name
  # $2 = arg value
  # $3 = arg parameter
  local param="$1"
  local rematch='^[^=]*=(.*)$' result
  if [[ $2 =~ $rematch ]]; then
    read "$param" <<<"${BASH_REMATCH[1]}"
  else
    for ((i = 3; $i <= $#; i++)); do
      # Only read next arg if it not an arg itself.
      if [[ ${*:$i:1} = -* ]]; then
        break
      fi
      result="$result ${@:$i:1}"
      # There is no way to shift our callers args, so
      # return "no of args" to indicate they should do it instead.
    done
    read "$1" <<<"$result"
    return $(($i - 3))
  fi
}

# For PhotonOS
default_kernel_images() {
  local kernel_version=""
  local installed_kernels=

  installed_kernels="$(find /boot/ -mindepth 1 -maxdepth 1 -type f \
                  -iname "vmlinuz-*.ph[[:digit:]]*" -exec basename {} \; | \
                  cut -d '-' -f 2-)"

  for kernel_version in ${installed_kernels}; do
    local modules_dir="/lib/modules/${kernel_version}"
    if [ ! -d "${modules_dir}" ]; then
      echo -e "\nERROR: ${modules_dir} not found for vmlinuz-${kernel_version}\n" >&2
      continue
    fi
    # Take this directory as the source of truth
    kernels="$kernels $kernel_version"
    targets="$targets $boot_dir/initrd.img-$kernel_version"
  done
}

while (($# > 0)); do
  case ${1%%=*} in
  --with-usb)
    read_arg usbmodule "$@" || shift $?
    basicmodules="$basicmodules ${usbmodule:-usb-storage}"
    unset usbmodule
    ;;
  --with-avail)
    read_arg modname "$@" || shift $?
    basicmodules="$basicmodules $modname"
    ;;
  --with)
    read_arg modname "$@" || shift $?
    basicmodules="$basicmodules $modname"
    ;;
  --version)
    echo "mkinitrd: dracut compatibility wrapper"
    exit 0
    ;;
  -v | --verbose) dracut_args="${dracut_args} -v" ;;
  -f | --force) force=1 ;;
  --preload)
    read_arg modname "$@" || shift $?
    basicmodules="$basicmodules $modname"
    ;;
  --image-version) img_vers=yes ;;
  --rootfs | -d)
    read_arg rootfs "$@" || shift $?
    dracut_args="${dracut_args} --filesystems $rootfs"
    ;;
  --nocompress) dracut_args="$dracut_args --no-compress" ;;
  --set-hostonly) set_hostonly=1 ;;
  --no-hostonly) no_hostonly=1 ;;
  --debug) dracut_dbg=1;;
  --help) usage -n ;;
  --builtin) ;;
  --without*) ;;
  --without-usb) ;;
  --fstab*) ;;
  --ifneeded) ;;
  --omit-scsi-modules) ;;
  --omit-ide-modules) ;;
  --omit-raid-modules) ;;
  --omit-lvm-modules) ;;
  --omit-dmraid) ;;
  --allow-missing) ;;
  --net-dev*) ;;
  --noresume) ;;
  --rootdev*) ;;
  --thawdev*) ;;
  --rootopts*) ;;
  --root*) ;;
  --loopdev*) ;;
  --loopfs*) ;;
  --loopopts*) ;;
  --looppath*) ;;
  --dsdt*) ;;
  -s) ;;
  --quiet | -q) quiet=1 ;;
  -b)
    read_arg boot_dir "$@" || shift $?
    if [ ! -d $boot_dir ]; then
      error "Boot directory $boot_dir does not exist"
      exit 1
    fi
    ;;
  -k) # Would be nice to get a list of images here
    read_arg kernel_images "$@" || shift $?
    for kernel_image in $kernel_images; do
      kernels="$kernels ${kernel_image#*-}"
    done
    host_only=1
    force=1
    ;;
  -i)
    read_arg initrd_images "$@" || shift $?
    for initrd_image in $initrd_images; do
      targets="$targets $boot_dir/$initrd_image"
    done
    ;;
  *) if [[ ! $targets ]]; then
    targets=$1
  elif [[ ! $kernels ]]; then
    kernels=$1
  else
    usage
  fi ;;
  esac
  shift
done

[[ $targets && $kernels ]] || default_kernel_images
[[ $targets && $kernels ]] || error "No Kernel Registered"

# We can have several targets/kernels, transform the list to an array
targets=($targets)
[[ $kernels ]] && kernels=($kernels)

# don't set hostonly flag if running in docker env
# if set initrd will be incomplete
# https://fedoraproject.org/wiki/Features/DracutHostOnly#Detailed_Description
# https://man7.org/linux/man-pages/man5/dracut.conf.5.html
if [ ${set_hostonly} -eq 0 ]; then
  if grep -qc docker /proc/self/cgroup; then
    echo "--- Generating initrd under docker environment ---"
    host_only=0
  elif [ "$(stat -c %d:%i /)" != "$(stat -c %d:%i /proc/1/root/.)" ]; then
    # set host_only=0 if initrd is generated in chroot
    # otherwise cloud images generated in a Photon host will have bad initrd
    echo "--- Generating initrd under chroot environment ---"
    host_only=0
  else
    host_only=1
  fi
else
  host_only=1
fi
force=1

if [ ${no_hostonly} -eq 1 ]; then
  host_only=0
fi

[[ $host_only == 1 ]] && dracut_args="${dracut_args} -H"
[[ $force == 1 ]] && dracut_args="${dracut_args} -f"

initrd_sanity_check() {
  local ret=0
  local status="$1"
  local log_fn="$2"
  local target="$3"
  local errmsg=""

  # initrd issues can be fatal, so if anything goes wrong during initrd creation
  # print errors to console
  if [ ${status} -ne 0 ]; then
    errmsg=$(cat << EOF
\n\n------------------------ ERROR NOTICE ------------------------------
          DRACUT RETURNED NON-ZERO EXIT STATUS(${status})

PROBABLY ${target} IS FAULTY
SYSTEM MAY BECOME UNUSABLE POST REBOOT
---------------------- PROCEED WITH CAUTION ------------------------
EOF
)
    echo -e "${errmsg}" 1>&2 |& tee -a "${log_fn}"
    ret=1
  fi

  if ! lsinitrd ${target} 1>/dev/null; then
    errmsg=$(cat << EOF
\n\n------------------------ ERROR NOTICE ------------------------------
                  lsinitrd ${target} FAILED

PROBABLY ${target} IS FAULTY
SYSTEM MAY BECOME UNUSABLE POST REBOOT
---------------------- PROCEED WITH CAUTION -----------------------\n\n
EOF
)
    echo -e "${errmsg}" 1>&2 |& tee -a "${log_fn}"
    ret=1
  fi

  return ${ret}
}

# If running in non tty, enable verbose
# It's probably happening in some kind of build env, we need verbose.
if ! test -t 1; then
  quiet=0
fi

final_ret=0

for ((i = 0; $i < ${#targets[@]}; i++)); do
  ret=0
  if [[ $img_vers ]];then
    target="${targets[$i]}-${kernels[$i]}"
  else
    target="${targets[$i]}"
  fi
  kernel="${kernels[$i]}"

  log_fn="/var/log/mkinitrd-${kernel}.log"

  echo "Creating $target"

  if [ $dracut_dbg -ne 0 ]; then
    export DRACUT_INSTALL="/usr/lib/dracut/dracut-install --verbose"
    dracut_cmd=(dracut -L 6 $dracut_args)
  else
    dracut_cmd=(dracut $dracut_args)
  fi

  # this check is for combination of newer and older kernels
  if [ -d "/lib/modules/${kernel}/dracut.conf.d" ]; then
    dracut_cmd+=(--confdir \"/lib/modules/${kernel}/dracut.conf.d /etc/dracut.conf.d\")
  elif [ -d "/var/lib/initramfs/kernel" ]; then
    kernel_ver_dir="/var/lib/initramfs/kernel"
    if [ -s "$kernel_ver_dir/$kernel" ]; then
      dracut_cmd+=("$(cat $kernel_ver_dir/$kernel)")
    fi
    unset kernel_ver_dir
  else
    error "ERROR: need /lib/modules/${kernel}/dracut.conf.d or /var/lib/initramfs/kernel"
    exit 1
  fi

  if [[ $basicmodules ]]; then
    dracut_cmd+=(--add-drivers \"$basicmodules\")
  fi
  dracut_cmd+=("$target" "$kernel")

  if [ ${quiet} -eq 1 ]; then
    eval "${dracut_cmd[@]}" &> "${log_fn}"
  else
    eval "${dracut_cmd[@]}" |& tee "${log_fn}"
  fi

  ret=${PIPESTATUS[0]}
  if ! initrd_sanity_check "${ret}" "${log_fn}" "${target}"; then
    final_ret=1
  fi
done

[ $final_ret -ne 0 ] && error "--- ERROR: mkinitrd FAILED, SOMETHING WENT WRONG ---"

exit ${final_ret}
