#!/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

    echo "Cannot detect initramfs path, aborting..."
    exit 1
}

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}
}

# 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"
    bin="$($SKIPCPIO_BIN "$file_path" | { read -r -N 6 bin && echo "$bin"; })"
    case $bin in
        $'\x1f\x8b'*)
            CAT="zcat"
            ;;
        BZh*)
            CAT="bzcat"
            ;;
        $'\x71\xc7'* | 070701)
            CAT="cat"
            ;;
        $'\x02\x21'*)
            CAT="lz4 -d -c"
            ;;
        $'\x89'LZO$'\0'*)
            CAT="lzop -d -c"
            ;;
        $'\x28\xB5\x2F\xFD'*)
            CAT="zstd -d -c"
            ;;
        *)
            if echo "test" | xz | xzcat --single-stream > /dev/null 2>&1; then
                CAT="xzcat --single-stream"
            else
                CAT="xzcat"
            fi
            ;;
    esac

    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()

while [[ $# -gt 0 ]]; do
  case $1 in
    --kver)
      kver="$2"
      shift 2
      ;;
    -*)
      echo "Unknown option $1"
      exit 1
      ;;
    *)
      break;
      ;;
  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=$(< /etc/fstab grep "${INITRAMFS_DIR}.*ext4" | awk '{print $1}')
  printf "%s\n%s\n%s\n%s\n%s\n%s\n" \
         "bootfs $boot_partition" \
         "bootfstype ext4" \
         "initoverlayfs_builder dracut -H -f -v -M --reproducible -o \"initoverlayfs\"" \
         "initrd_builder dracut -H -f -v -M --reproducible -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

  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 $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"

detect_initramfs
extract_initrd_into_initoverlayfs

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

