#!/usr/bin/bash

set -ex

# Constants with default value
CAT="lz4cat"

INITRAMFS_DIR=""
INITOVERLAYFS_CONF="/etc/initoverlayfs.conf"
INITRAMFS_DUMP_DIR="/var/lib/initoverlayfs"

# Only erofs-based initoverlayfs supported, ext4, btrfs, xfs, etc. to be implemented"
SUPPORTED_FILESYSTEM=("erofs")
SKIPCPIO_BIN="/usr/lib/dracut/skipcpio"

detect_path_initramfs() {
    paths_initramfs=(
	"/boot/initramfs-$kver.img"
        "/usr/lib/modules/$kver/initramfs.img"
    )

    for path in "${paths_initramfs[@]}"; do
        if [[ -e "${path}" ]]; then
	    INITRAMFS_DIR=$(dirname "${path}")
	    return
        fi

    done

    # on first build, like in osbuild, there will be no prior initrd to detect
    INITRAMFS_DIR="/boot"
}

exec_erofs() {
    pushd "$INITRAMFS_DUMP_DIR"
        "$SKIPCPIO_BIN" "${INITRAMFS_DIR}/initramfs-$kver.img" | "$CAT" | cpio -ivd
    popd
    rm -f "${INITRAMFS_DIR}/initoverlayfs-$kver.img"
    mkfs.erofs $erofs_compression "${INITRAMFS_DIR}/initoverlayfs-$kver.img" ${INITRAMFS_DUMP_DIR}
    if false; then
        veritysetup format "${INITRAMFS_DIR}/initoverlayfs-$kver.img" "/etc/initoverlayfs-hash-$kver.img"
    fi
}

# Support for ext4 is currently under development.
exec_ext4() {
    dd if=/dev/zero of="${INITRAMFS_DIR}"/initoverlayfs-"$kver".img bs=64M count=1
    dev=$(losetup -fP --show "${INITRAMFS_DIR}"/initoverlayfs-"$kver".img)
    mkfs.ext4 "${dev}"
    mount "${dev}" "${INITRAMFS_DUMP_DIR}"

    pushd "${INITRAMFS_DUMP_DIR}" || exit
        "${SKIPCPIO_BIN}" "${INITRAMFS_DIR}"/initramfs-"$kver".img | zstd -d --stdout | cpio -ivd
        sync
    popd || exit

    while ! umount "${INITRAMFS_DUMP_DIR}"; do
      sleep 1
    done

    losetup -d "${dev}"
}

detect_initramfs() {
  mkdir -p "${INITRAMFS_DUMP_DIR}"

  echo "Extracting initrd into initoverlayfs..."

  file_path="${INITRAMFS_DIR}/initramfs-$kver.img"
  skipcpio="/usr/lib/dracut/skipcpio"
  if $skipcpio $file_path | gzip -t - >/dev/null 2>&1; then
    CAT="zcat"
  elif $skipcpio $file_path | zstd -q -c -t - >/dev/null 2>&1; then
    CAT="zstd"
  elif $skipcpio $file_path | xzcat -t - >/dev/null 2>&1; then
    CAT="xzcat"
  elif $skipcpio $file_path | lz4cat -t - >/dev/null 2>&1; then
    CAT="lz4cat"
  elif $skipcpio $file_path | bzip2 -t - >/dev/null 2>&1; then
    CAT="bzcat"
  elif $skipcpio $file_path | lzop -t - >/dev/null 2>&1; then
    CAT="lzop"
  else
    CAT="cat"
  fi

  echo "  - File path: ${file_path}"
  echo "  - Decompressor: $CAT"
}

extract_initrd_into_initoverlayfs() {
    if command -v mkfs.erofs; then
        fstype="erofs"
    elif command -v mkfs.ext4; then
	fstype="ext4"
    else
	fstype="unsupported"
    fi

    case "${fstype}" in
        # Support for ext4 is currently under development.
        # *ext4*)
        #    exec_ext4
        #    ;;
        *erofs*)
	    exec_erofs
            ;;
        *)
            echo -e "The detected filesytem: is ${fstype}." \
		    "Unfortunately it's not supported at moment."
	    echo -e "Supported filesystems: ${SUPPORTED_FILESYSTEM[*]}"
	    exit 1
            ;;
    esac

    rm -rf "$INITRAMFS_DUMP_DIR" &
}

# main()

args="$*"
initoverlayfs_init="false"
while [[ $# -gt 0 ]]; do
  echo "$1"
  case $1 in
    --kver)
      kver="$2"
      shift 2
      ;;
    --initoverlayfs-init)
      initoverlayfs_init="true"
      args=$(echo "$args" | sed "s/--initoverlayfs-init//g")
      shift 1
      ;;
    *)
      shift 1
      ;;
  esac
done

# This logic for the case where there is no kernel present
# in this directory, some microVMs and containers.
#
# no_kern=""
# if ! compgen -G /boot/vmlinu* > /dev/null; then
#  no_kern="--no-kernel"
# fi

if [ -z "$kver" ]; then
  kver="$(uname -r)"
fi

detect_path_initramfs

if ! [ -e "$INITOVERLAYFS_CONF" ] || ! grep -q '[^[:space:]]' "$INITOVERLAYFS_CONF"; then
  boot_partition=$(grep "${INITRAMFS_DIR}.*ext4" /etc/fstab | awk '{print $1}')
  boot_partition_hint=$(blkid -t $boot_partition | awk -F: '{print $1}')

  printf "%s\n%s\n%s\n%s\n" \
         "bootfs $boot_partition" \
         "bootfs_hint $boot_partition_hint" \
         "bootfstype ext4" \
         "initoverlayfs_builder dracut -M -o \"initoverlayfs fcoe\"" > $INITOVERLAYFS_CONF

  if $initoverlayfs_init; then
    printf "%s\n" "initrd_builder dracut -M -m \"initoverlayfs\" -o \"kernel-modules udev-rules systemd base bash systemd-initrd i18n kernel-modules-extra rootfs-block dracut-systemd usrmount fs-lib microcode_ctl-fw_dir_override shutdown nss-softokn\"" >> $INITOVERLAYFS_CONF
  else
    printf "%s\n" "initrd_builder dracut -M -m \"kernel-modules udev-rules initoverlayfs systemd base\" -o \"bash systemd-initrd i18n kernel-modules-extra rootfs-block dracut-systemd usrmount fs-lib microcode_ctl-fw_dir_override shutdown nss-softokn\"" >> $INITOVERLAYFS_CONF
  fi

  erofs_compression_supported="true"
  # shellcheck disable=SC2034
  . /etc/os-release
  for i in $ID_LIKE; do
    if [ "$i" == "rhel" ]; then
      if [ "$VERSION_ID" -le 9 ]; then
        erofs_compression_supported="false"
        break
      fi
    fi
  done

  if $initoverlayfs_init; then
    printf "%s\n" "initoverlayfs_init true" >> $INITOVERLAYFS_CONF
  fi

  if $erofs_compression_supported; then
    printf "%s\n" "erofs_compression -zlz4hc,11" >> $INITOVERLAYFS_CONF
  fi
fi

erofs_compression=$(sed -ne "s/^erofs_compression\s//pg" "$INITOVERLAYFS_CONF")
initoverlayfs_builder=$(sed -ne "s/^initoverlayfs_builder\s//pg" "$INITOVERLAYFS_CONF")
/bin/bash -c "$initoverlayfs_builder $args"

detect_initramfs
extract_initrd_into_initoverlayfs

initrd_builder=$(sed -ne "s/^initrd_builder\s//pg" "$INITOVERLAYFS_CONF")
/bin/bash -c "$initrd_builder $args"

