#!/bin/bash --norc
kernel_ver_dir='/var/lib/initramfs/kernel'

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

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

usage () {
    [[ $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]"
    $cmd "       [--no-hostonly]"
    $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
    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 photon
default_kernel_images() {
    local kernel_version

    for kernel_version in $(ls $kernel_ver_dir) ; do

        # 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;;
        --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*) ;;
        --bootchart) ;;
	-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"

for ((i=0 ; $i<${#targets[@]} ; i++)); do

    if [[ $img_vers ]];then
	target="${targets[$i]}-${kernels[$i]}"
    else
	target="${targets[$i]}"
    fi
    kernel="${kernels[$i]}"

    if [[ -s "$kernel_ver_dir/$kernel" ]]; then
        readarray -t kernel_cfg < <(xargs -n1 -a $kernel_ver_dir/$kernel)
    else
        kernel_cfg=()
    fi

    # Duplicate code: No way found how to redirect output based on $quiet
    if [[ $quiet == 1 ]];then
	echo "Creating $target"
	if [[ $basicmodules ]]; then
            dracut $dracut_args --add-drivers "$basicmodules" "${kernel_cfg[@]}" "$target" \
		"$kernel" &>/dev/null
	else
            dracut $dracut_args "${kernel_cfg[@]}" "$target" "$kernel" &>/dev/null
	fi
    else
	if [[ $basicmodules ]]; then
            dracut $dracut_args --add-drivers "$basicmodules" "${kernel_cfg[@]}" "$target" \
		"$kernel"
	else
            dracut $dracut_args "${kernel_cfg[@]}" "$target" "$kernel"
	fi
    fi
done
