#!/usr/bin/env bash
PATH=/usr/sbin:/usr/bin:/sbin:/bin
export PATH

# shellcheck disable=SC1090
[[ -e .testdir${TEST_RUN_ID:+-$TEST_RUN_ID} ]] && . .testdir${TEST_RUN_ID:+-$TEST_RUN_ID}
if [[ -z $TESTDIR ]] || [[ ! -d $TESTDIR ]]; then
    TESTDIR=$(mktemp -d -p "/var/tmp" -t dracut-test.XXXXXX)
fi
echo "TESTDIR=\"$TESTDIR\"" > .testdir${TEST_RUN_ID:+-$TEST_RUN_ID}
export TESTDIR
export BOOT_ROOT="$TESTDIR"

[ -z "$USE_NETWORK" ] && USE_NETWORK="network"
[ -z "$TEST_FSTYPE" ] && TEST_FSTYPE="ext4"

if [[ -z $basedir ]]; then basedir="$(realpath ../..)"; fi

DRACUT=${DRACUT-dracut}
PKGLIBDIR=${PKGLIBDIR-$basedir}

[[ -f /etc/machine-id ]] && read -r TOKEN < /etc/machine-id
[ -z "$TOKEN" ] && . /etc/os-release && TOKEN="$ID"

TEST_KERNEL_CMDLINE+=" root=LABEL=dracut panic=1 oops=panic softlockup_panic=1 systemd.crash_reboot $DEBUGFAIL "

if [[ $V != "1" && $V != "2" ]]; then
    TEST_KERNEL_CMDLINE+="quiet "
fi

if [[ $V == "2" ]]; then
    TEST_KERNEL_CMDLINE+="rd.debug "
fi

ARCH="${ARCH-$(uname -m)}"

case "$ARCH" in
    amd64 | i?86 | x86_64)
        TEST_KERNEL_CMDLINE+="console=ttyS0 "
        ;;
esac

test_dracut() {
    # directory for test configurations and for the generated initrd (extracted)
    mkdir -p "$TESTDIR"/dracut.conf.d "$TESTDIR"/initrd

    # grab the distro configuration from the host and make it available for the tests
    if [ -d /usr/lib/dracut/dracut.conf.d ]; then
        cp -a /usr/lib/dracut/dracut.conf.d "$TESTDIR"/
    fi

    # pick up configuration from $TESTDIR/dracut.conf.d when running the tests
    TEST_DRACUT_ARGS+=" --confdir $TESTDIR/dracut.conf.d --add-confdir test --keep --tmpdir $TESTDIR/initrd --kver $KVERSION"

    # include $TESTDIR"/overlay if exists
    if [ -d "$TESTDIR"/overlay ]; then
        TEST_DRACUT_ARGS+=" --include $TESTDIR/overlay /"
    fi

    # shellcheck disable=SC2162
    IFS=' ' read -a TEST_DRACUT_ARGS_ARRAY <<< "$TEST_DRACUT_ARGS"

    "$DRACUT" \
        "${TEST_DRACUT_ARGS_ARRAY[@]}" \
        "$@" "$TESTDIR"/initramfs.testing || return 1
}

# Wait for the server QEMU has been started up.
# It should print "Serving" in the server.log in that case.
wait_for_server_startup() {
    local lines printed_lines=0 server_pid
    server_pid=$(cat "$TESTDIR"/server.pid)

    echo "Waiting for the server to startup"
    while ! grep -q Serving "$TESTDIR"/server.log; do
        if [ "$V" -ge 1 ]; then
            lines=$(wc -l "$TESTDIR"/server.log | cut -f 1 -d ' ')
            if [ "$lines" -gt "$printed_lines" ]; then
                tail -n "+$((printed_lines + 1))" "$TESTDIR"/server.log
                printed_lines=$lines
            fi
        fi
        if ! test -f "/proc/$server_pid/status"; then
            echo "Error: Server QEMU process $server_pid is gone. Please check $TESTDIR/server.log for failures." >&2
            return 1
        fi
        sleep 1
    done
}

ovmf_code() {
    for path in \
        "/usr/share/OVMF/OVMF_CODE.fd" \
        "/usr/share/OVMF/OVMF_CODE_4M.fd" \
        "/usr/share/edk2/x64/OVMF_CODE.fd" \
        "/usr/share/edk2/x64/OVMF_CODE.4m.fd" \
        "/usr/share/edk2-ovmf/OVMF_CODE.fd" \
        "/usr/share/qemu/ovmf-x86_64-4m.bin"; do
        [[ -s $path ]] && echo -n "$path" && return
    done
}

determine_kernel_version() {
    # inspired by kernel version detection in lsinitrd.sh

    # Some test containers do not include systemd-detect-virt, so let's just assume
    # we are running inside a container already

    # shellcheck disable=SC2012
    [[ $KVERSION ]] || KVERSION="$(cd /lib/modules && ls -1 | tail -1)"
    # shellcheck disable=SC2012
    [[ $KVERSION ]] || KVERSION="$(cd /usr/lib/modules && ls -1 | tail -1)"
    [[ $KVERSION ]] || KVERSION="$(uname -r)"
}

set_vmlinux_env() {
    VMLINUZ=${VMLINUZ-"/lib/modules/${KVERSION}/vmlinuz"}
    if ! [ -f "$VMLINUZ" ]; then
        VMLINUZ="/lib/modules/${KVERSION}/vmlinux"
    fi

    if ! [ -f "$VMLINUZ" ]; then
        [[ -f /etc/machine-id ]] && read -r MACHINE_ID < /etc/machine-id

        if [[ $MACHINE_ID ]] && { [[ -d /boot/${MACHINE_ID} ]] || [[ -L /boot/${MACHINE_ID} ]]; }; then
            VMLINUZ="/boot/${MACHINE_ID}/$KVERSION/linux"
        elif [ -f "/boot/vmlinuz-${KVERSION}" ]; then
            VMLINUZ="/boot/vmlinuz-${KVERSION}"
        elif [ -f "/boot/vmlinux-${KVERSION}" ]; then
            VMLINUZ="/boot/vmlinux-${KVERSION}"
        elif [ -f "/boot/kernel-${KVERSION}" ]; then
            VMLINUZ="/boot/kernel-${KVERSION}"
        elif [ -f "/boot/Image-${KVERSION}" ]; then
            VMLINUZ="/boot/Image-${KVERSION}"
        fi
    fi

    if ! [ -f "$VMLINUZ" ]; then
        VMLINUZ=$(find /boot/vmlinuz-* -type f 2> /dev/null | tail -1)
    fi

    if ! [ -f "$VMLINUZ" ]; then
        echo "Could not find a Linux kernel version $KVERSION to test with!" >&2
        echo "Please install linux." >&2
        exit 1
    fi

    export VMLINUZ
}

set_test_envonment_variables() {
    determine_kernel_version
    set_vmlinux_env
}

command -v test_check &> /dev/null || test_check() {
    :
}

command -v test_cleanup &> /dev/null || test_cleanup() {
    :
}

# terminal sequence to set color to a 'success' color (currently: green)
function SETCOLOR_SUCCESS() { echo -en '\033[0;32m'; }
# terminal sequence to set color to a 'failure' color (currently: red)
function SETCOLOR_FAILURE() { echo -en '\033[0;31m'; }
# terminal sequence to set color to a 'warning' color (currently: yellow)
function SETCOLOR_WARNING() { echo -en '\033[0;33m'; }
# terminal sequence to reset to the default color.
function SETCOLOR_NORMAL() { echo -en '\033[0;39m'; }

COLOR_SUCCESS='\033[0;32m'
COLOR_FAILURE='\033[0;31m'
COLOR_WARNING='\033[0;33m'
COLOR_NORMAL='\033[0;39m'

# generate qemu arguments for named raw disks
#
# qemu_add_drive <index> <args> <filename> <id-name> [<bootindex>]
#
# index: name of the index variable (set to 0 at start)
# args: name of the argument array variable (set to () at start)
# filename: filename of the raw disk image
# id-name: name of the disk in /dev/disk/by-id -> /dev/disk/by-id/scsi-0QEMU_QEMU_HARDDISK_$name
# size: optional boolean to initialize the drive
# bootindex: optional bootindex number
#
# to be used later with `qemu … "${args[@]}" …`
# The <index> variable will be incremented each time the function is called.
#
# can't be easier than this :-/
#
# # EXAMPLES
# ```
#   declare -a disk_args=()
#   declare -i disk_index=0
#   qemu_add_drive disk_index disk_args "$TESTDIR"/root.ext3 root 0 1
#   qemu_add_drive disk_index disk_args "$TESTDIR"/client.img client
#   qemu_add_drive disk_index disk_args "$TESTDIR"/iscsidisk2.img iscsidisk2
#   qemu_add_drive disk_index disk_args "$TESTDIR"/iscsidisk3.img iscsidisk3
#   qemu "${disk_args[@]}"
# ```
qemu_add_drive() {
    local index=${!1}
    local file=$3
    local name=${4:-$index}
    local size=${5:-0}
    local bootindex=$6

    if [ "${size}" -ne 0 ]; then
        if [[ ${name} =~ "marker" ]]; then
            size=1
        else
            size=512
        fi

        dd if=/dev/zero of="${file}" bs=1MiB count="${size}" status=none
        sync
    fi

    eval "${2}"'+=(' \
        -device "virtio-scsi-pci,id=scsi${index}" \
        -drive "if=none,format=raw,file=${file},id=drive-data${index}" \
        -device "scsi-hd,bus=scsi${index}.0,drive=drive-data${index},id=data${index},${bootindex:+bootindex=$bootindex,}serial=${name}" \
        ')'

    [ $(("${1}++")) -ne 0 ]
}

test_marker_reset() {
    dd if=/dev/zero of="$TESTDIR"/marker.img bs=1MiB count=1 status=none
    sync
}

test_marker_check() {
    local marker=${1:-dracut-root-block-success}
    local file=${2:-marker.img}

    grep -U --binary-files=binary -F -m 1 -q "$marker" "$TESTDIR/$file"
    return $?
}

while (($# > 0)); do
    case $1 in
        --run)
            echo "TEST RUN: $TEST_DESCRIPTION"
            set_test_envonment_variables
            test_check && test_run
            exit $?
            ;;
        --setup)
            echo "TEST SETUP: $TEST_DESCRIPTION"
            set_test_envonment_variables
            test_check && test_setup
            exit $?
            ;;
        --clean)
            echo "TEST CLEANUP: $TEST_DESCRIPTION"
            test_cleanup
            rm -fr -- "$TESTDIR"
            rm -f -- .testdir${TEST_RUN_ID:+-$TEST_RUN_ID}
            exit $?
            ;;
        --all)
            set_test_envonment_variables
            if ! test_check 2 &> test${TEST_RUN_ID:+-$TEST_RUN_ID}.log; then
                if [[ $V == "1" || $V == "2" ]]; then
                    cat test${TEST_RUN_ID:+-$TEST_RUN_ID}.log
                fi
                echo -e "TEST: $TEST_DESCRIPTION " "$COLOR_WARNING" "[SKIPPED]" "$COLOR_NORMAL"
                exit 0
            else
                echo -e "TEST: $TEST_DESCRIPTION " "$COLOR_SUCCESS" "[STARTED]" "$COLOR_NORMAL"
            fi
            if [[ $V == "1" || $V == "2" ]]; then
                tee_command="tee"
                set -o pipefail
                (
                    test_setup && test_run
                    ret=$?
                    test_cleanup
                    if ((ret != 0)) && [[ -f "$TESTDIR"/server.log ]]; then
                        mv "$TESTDIR"/server.log ./server${TEST_RUN_ID:+-$TEST_RUN_ID}.log
                    fi
                    rm -fr -- "$TESTDIR"
                    rm -f -- .testdir${TEST_RUN_ID:+-$TEST_RUN_ID}
                    exit $ret
                ) 2>&1 | "$tee_command" "test${TEST_RUN_ID:+-$TEST_RUN_ID}.log"
            else
                (
                    test_setup && test_run
                    ret=$?
                    test_cleanup
                    rm -fr -- "$TESTDIR"
                    rm -f -- .testdir${TEST_RUN_ID:+-$TEST_RUN_ID}
                    exit $ret
                ) > test${TEST_RUN_ID:+-$TEST_RUN_ID}.log 2>&1
            fi
            ret=$?
            set +o pipefail
            if [ $ret -eq 0 ]; then
                rm -- test${TEST_RUN_ID:+-$TEST_RUN_ID}.log
                echo -e "TEST: $TEST_DESCRIPTION " "$COLOR_SUCCESS" "[OK]" "$COLOR_NORMAL"
            else
                echo -e "TEST: $TEST_DESCRIPTION " "$COLOR_FAILURE" "[FAILED]" "$COLOR_NORMAL"
                if [ "$V" == "2" ]; then
                    tail -c 1048576 "$(pwd)/server${TEST_RUN_ID:+-$TEST_RUN_ID}.log" "$(pwd)/test${TEST_RUN_ID:+-$TEST_RUN_ID}.log"
                    echo -e "TEST: $TEST_DESCRIPTION " "$COLOR_FAILURE" "[FAILED]" "$COLOR_NORMAL"
                else
                    echo "see $(pwd)/test${TEST_RUN_ID:+-$TEST_RUN_ID}.log"
                fi
            fi
            exit $ret
            ;;
        *) break ;;
    esac
    shift
done