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

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

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

if [[ -f /etc/machine-id ]]; then
    read -r TOKEN < /etc/machine-id || true
fi
[ -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+="systemd.journald.forward_to_console=1 "
fi

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

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"

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

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

quote_args() {
    local arg args=()
    for arg in "$@"; do
        if [[ -z $arg || $arg =~ [[:space:]\'] ]]; then
            args+=("'${arg//\'/\'\\\'\'}'")
        else
            args+=("$arg")
        fi
    done
    echo "${args[*]}"
}

# Call dracut and log its call in verbose mode
call_dracut() {
    # Uncomment this to debug failures
    #set -- --stdlog 4 "$@"
    if [[ ${V-} == "1" || ${V-} == "2" ]]; then
        echo "Calling $DRACUT $(quote_args "$@")"
    fi
    "$DRACUT" "$@"
}

# 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

    if [ "${V-}" -ge 1 ]; then
        tail -n "+$((printed_lines + 1))" "$TESTDIR"/server.log
        echo "Server has been start up."
    fi
}

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

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

determine_kernel_version() {
    lsinitrd "$1" | grep modules.dep | head -1 | sed 's;.*/\([^/]\+\)/modules.dep;\1;'
}

determine_kernel_image() {
    local kversion="$1"
    local paths=(
        "/lib/modules/${kversion}/vmlinuz"
        "/lib/modules/${kversion}/vmlinux"
        "/lib/modules/${kversion}/Image"
        "/boot/vmlinuz-${kversion}"
        "/boot/vmlinux-${kversion}"
    )

    for path in "${paths[@]}"; do
        if [ -f "$path" ]; then
            echo "$path"
            return 0
        fi
    done

    echo "Could not find a Linux kernel image for version '$kversion'!" >&2
    exit 1
}

# 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'

# build an ext4 image from the given directory
build_ext4_image() {
    local source_dir="$1"
    local image="$2"
    local label="$3"

    dd if=/dev/zero of="$image" bs=200MiB count=1 status=none
    sync "$image"
    mkfs.ext4 -q -L "$label" -d "$source_dir" "$image"
    sync "$image"
}

# override the init script from the test-root dracut module (see module-setup.sh)
inst_init() {
    local init_script="$1"
    local rootdir="$2"
    if [ -e "$rootdir/sbin/test-init" ]; then
        cp "$init_script" "$rootdir/sbin/test-init"
    else
        cp "$init_script" "$rootdir/sbin/init"
    fi
}

# 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 "${file}"
    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}++"))
}

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

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

print_test_result() {
    local ret=${1-1}
    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 [[ -f "$TESTDIR"/server.log ]]; then
            mv "$TESTDIR"/server.log ./server${TEST_RUN_ID:+-$TEST_RUN_ID}.log
        fi
        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
}

while (($# > 0)); do
    case $1 in
        --run)
            echo "TEST RUN: $TEST_DESCRIPTION"
            test_check
            test_run
            exit $?
            ;;
        --setup)
            echo "TEST SETUP: $TEST_DESCRIPTION"
            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)
            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 77
            else
                echo -e "TEST: $TEST_DESCRIPTION " "$COLOR_SUCCESS" "[STARTED]" "$COLOR_NORMAL"
            fi
            trap print_test_result ERR
            if [[ ${V-} == "1" || ${V-} == "2" ]]; then
                tee_command="tee"
                set -o pipefail
                (
                    test_setup
                    test_run
                    test_cleanup
                    rm -fr -- "$TESTDIR"
                    rm -f -- .testdir${TEST_RUN_ID:+-$TEST_RUN_ID}
                ) 2>&1 | "$tee_command" "test${TEST_RUN_ID:+-$TEST_RUN_ID}.log"
            else
                (
                    test_setup
                    test_run
                    test_cleanup
                    rm -fr -- "$TESTDIR"
                    rm -f -- .testdir${TEST_RUN_ID:+-$TEST_RUN_ID}
                ) > test${TEST_RUN_ID:+-$TEST_RUN_ID}.log 2>&1
            fi
            set +o pipefail
            print_test_result 0
            exit 0
            ;;
        *) break ;;
    esac
    shift
done