#!/usr/bin/env bash
set -euxo pipefail
GPTDISKFILE=disk.raw
DOSDISKFILE=dosdisk.raw
UKIBOOTPARTIMG=ukibootctrl.bin
ESPPARTIMG=esp.raw

declare -A UKI1_SEEK
UKI1_SEEK["gpt"]=69632
UKI1_SEEK["dos"]=69634
declare -A UKI2_SEEK
UKI2_SEEK["gpt"]=71680
UKI2_SEEK["dos"]=71684
declare -A CTL_SEEK
CTL_SEEK["gpt"]=73728
CTL_SEEK["dos"]=73734

fallocate -l 40M $GPTDISKFILE
fallocate -l 40M $DOSDISKFILE

sfdisk $GPTDISKFILE <<EOF
label: gpt
label-id: 235DAC1C-F4CC-4932-9C62-AD6FF20EA2DF
device: disk.raw
unit: sectors
first-lba: 34
last-lba: 81886
sector-size: 512

disk.raw1 : start=        2048, size=       67584, type=C12A7328-F81F-11D2-BA4B-00A0C93EC93B, uuid=49127E7B-BAEC-402A-922F-DFA9442995F6
disk.raw2 : start=       69632, size=        2048, type=DF331E4D-BE00-463F-B4A7-8B43E18FB53A, uuid=02BD6AD1-7466-4625-B2E9-9F5F2D74DCBF, name="ukiboot_a"
disk.raw3 : start=       71680, size=        2048, type=DF331E4D-BE00-463F-B4A7-8B43E18FB53A, uuid=3A8D4AFC-113C-47AF-8619-C0CDB5BBF2E3, name="ukiboot_b"
disk.raw4 : start=       73728, size=        6144, type=FEFD9070-346F-4C9A-85E6-17F07F922773, uuid=890246DB-298E-4F65-A121-CA143666854F, name="ukibootctl"
EOF

sfdisk $DOSDISKFILE <<EOF
label: dos
device: dosdisk.raw
unit: sectors

disk.raw1 : start=        2048, size=       67584, type=ef
disk.raw2 : start=       69632  size=       10240, type=5
disk.raw5 : start=       69634, size=        2048, type=46
disk.raw6 : start=       71684, size=        2048, type=47
disk.raw7 : start=       73734, size=        6138, type=48
EOF

fallocate -l $((67584 * 512)) $ESPPARTIMG
mkfs.vfat -F 32 $ESPPARTIMG
mmd   -i $ESPPARTIMG ::EFI
mmd   -i $ESPPARTIMG ::EFI/BOOT

QEMU_CMD="qemu-system-`uname -m`"
case `uname -m` in
    'x86_64')
        UEFI_ARCH_EXT=X64
        QEMU_MTYPE=q35
        QEMU_CPU=qemu64
        OVMF=/usr/share/edk2/ovmf/OVMF_CODE.fd
    ;;
    'aarch64')
        UEFI_ARCH_EXT=AA64
        QEMU_MTYPE=virt
        QEMU_CPU=cortex-a57
        OVMF=/usr/share/edk2/aarch64/QEMU_EFI.fd
    ;;
    *)
    #TODO: Add support for other UEFI relevant architectures
    echo "`uname -m` not supported" 1>&2
    exit 1;
    ;;
esac

mcopy -i $ESPPARTIMG efi/ukiboot${UEFI_ARCH_EXT,,}.efi ::EFI/BOOT/BOOT$UEFI_ARCH_EXT.EFI

# Generate ukibootctrl.bin
touch $UKIBOOTPARTIMG

ukibootctl_path_init () {
    $UKIBOOTCTL_PATH init --partition=$UKIBOOTPARTIMG
}

load_ctl_part () {
    dd of=$DISK_FILE conv=notrunc if=$UKIBOOTPARTIMG seek=${CTL_SEEK[${LABEL_ID}]}
}

copy_to_partitions () {
    dd of=$DISK_FILE conv=notrunc if=$ESPPARTIMG     seek=2048
    dd of=$DISK_FILE conv=notrunc if=test/SLOTA.EFI  seek=${UKI1_SEEK[${LABEL_ID}]}
    dd of=$DISK_FILE conv=notrunc if=test/SLOTB.EFI  seek=${UKI2_SEEK[${LABEL_ID}]}
    load_ctl_part ${LABEL_ID} ${DISK_FILE}
}


test_qemu_output () {
    SEARCH_STRING=$1
    # Run the program with a timeout, and pipe its output to grep
    if [ -c /dev/kvm ]; then
        QEMU_ACCEL_ARGS="-accel kvm"
        CPU=host
    else
        QEMU_ACCEL_ARGS="-accel tcg,thread=multi"
        CPU=$QEMU_CPU
    fi
    if $QEMU_CMD $QEMU_ACCEL_ARGS -M $QEMU_MTYPE -bios $OVMF -cpu $CPU -drive if=virtio,id=system,format=raw,file=$DISK_FILE -display none -serial stdio | grep -q "##$SEARCH_STRING" ; then
        echo "String '##$SEARCH_STRING' found! Exiting."
    else
        # Check the exit status of timeout to differentiate between string not found and actual timeout
        TIMEOUT_EXIT_STATUS=$?
        if [ $TIMEOUT_EXIT_STATUS -eq 124 ]; then
            echo "Timeout: String '$SEARCH_STRING' not found within $TIMEOUT_SECONDS seconds. Exiting with error."
            exit 1
        elif [ $TIMEOUT_EXIT_STATUS -eq 1 ]; then
            echo "String '$SEARCH_STRING' not found, and program completed before timeout. Exiting."
            exit 2
        else
            echo "An unexpected error occurred (exit status: $TIMEOUT_EXIT_STATUS)."
            exit 3
        fi
    fi
}

run_test () {
    LABEL_ID=$1
    DISK_FILE=$2
    ukibootctl_path_init
    copy_to_partitions $LABEL_ID $DISK_FILE
    test_qemu_output "SLOT A" $DISK_FILE
    $UKIBOOTCTL_PATH --partition=$UKIBOOTPARTIMG prepare-switch 1
    $UKIBOOTCTL_PATH --partition=$UKIBOOTPARTIMG finalize-switch 1
    load_ctl_part $LABEL_ID $DISK_FILE
    test_qemu_output "SLOT B" $DISK_FILE
}


echo ***Test: GPT***
run_test "gpt" $GPTDISKFILE

echo ***Test: DOS***
run_test "dos" $DOSDISKFILE
