#!/usr/bin/bash

set -x
set -e

if [ ! -f slot_a.addon.efi ]; then
    echo Run this script from the build dir
    exit 1
fi

TEST_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )

UKIFY=ukify
FIRMWARE=ovmf
UBOOT=u-boot.rom
DUMP=disk.dump
DISK_LABEL=gpt
declare -A UKI1_SEEK
UKI1_SEEK["gpt"]=22528
UKI1_SEEK["dos"]=22530
declare -A UKI2_SEEK
UKI2_SEEK["gpt"]=327328
UKI2_SEEK["dos"]=327332
declare -A CTL_SEEK
CTL_SEEK["gpt"]=632128
CTL_SEEK["dos"]=632134

while [[ $# -gt 0 ]]; do
    case "$1" in
        --firmware=*)
            FIRMWARE="${1#*=}"
            shift
            ;;
        --firmware)
            if [[ $# -gt 1 ]]; then
                FIRMWARE="$2"
                shift 2
            else
                echo "Error: --firmware requires an argument (e.g., --firmware=u-boot)"
                exit 1
            fi
            ;;
        --disk_label=*)
            DUMP="${1#*=}.dump"
            DISK_LABEL="${1#*=}"
            shift
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

generate_uki() {
    local OUT=$1
    local CMDLINE=$2

    local KVER=$(uname --kernel-release)

    if [ $(sudo bootctl is-installed) == "yes" ]; then
        BOOT=$(bootctl --print-esp-path)
        MACHINE_ID=$(cat /etc/machine-id)
        LINUX=${BOOT}/${MACHINE_ID}/${KVER}/linux
        INITRD=${BOOT}/${MACHINE_ID}/${KVER}/initrd
    else
	BOOT="/boot"
        LINUX=${BOOT}/vmlinuz-${KVER}
        INITRD=${BOOT}/initramfs-${KVER}.img
    fi

    sudo $UKIFY build --output=$OUT --linux ${LINUX} --initrd ${INITRD} \
         --uname `uname --kernel-release` \
         --cmdline "console=ttyS0 $CMDLINE rd.shell rd.break=pre-mount"
}

create_disk() {
    if [ ! -f uki1.efi ]; then
        generate_uki uki1.efi example=1
    fi
    if [ ! -f uki2.efi ]; then
        generate_uki uki2.efi example=2
    fi

    fallocate -l 400M disk.img
    sfdisk disk.img < ${TEST_DIR}/${DUMP}
    dd if=uki1.efi of=disk.img seek=${UKI1_SEEK[${DISK_LABEL}]} conv=notrunc
    dd if=uki2.efi of=disk.img seek=${UKI2_SEEK[${DISK_LABEL}]} conv=notrunc
}

update_ukiboot() {
    fallocate -l 10M esp.img
    mkfs.vfat esp.img
    mmd -i esp.img ::/EFI
    mmd -i esp.img ::/EFI/BOOT
    mmd -i esp.img ::/EFI/BOOT/ukiboot_a.efi.extra.d
    mmd -i esp.img ::/EFI/BOOT/ukiboot_b.efi.extra.d
    mcopy -i esp.img slot_a.addon.efi ::/EFI/BOOT/ukiboot_a.efi.extra.d/slot_a.addon.efi
    mcopy -i esp.img slot_b.addon.efi ::/EFI/BOOT/ukiboot_b.efi.extra.d/slot_b.addon.efi
    if test -f efi/ukibootx64.efi; then
	mcopy -i esp.img efi/ukibootx64.efi  ::/EFI/BOOT/BOOTX64.EFI
    fi
    if test -f efi/ukibootaa64.efi; then
	mcopy -i esp.img efi/ukibootaa64.efi ::/EFI/BOOT/BOOTAA64.EFI
    fi
    dd if=esp.img of=disk.img seek=2048 conv=notrunc
}

reset_ukibootctl() {
        dd if=/dev/zero of=disk.img seek=${CTL_SEEK[${DISK_LABEL}]} bs=512 count=4 conv=notrunc # Reset boot counters, etc
}

if [ ! -f disk.img ]; then
    create_disk
fi

update_ukiboot
reset_ukibootctl

ARCH=$(arch)
QEMU_ARGS="\
    -m 1G \
    -device virtio-blk-pci,drive=virtiodisk \
    -drive file=disk.img,if=none,id=virtiodisk,format=raw \
    -nographic \
    -serial mon:stdio
"

if [ "$ARCH" == "aarch64" ]; then

    if [ "$FIRMWARE" == "u-boot" ]; then
        QEMU_ARGS+=" -bios u-boot.rom"
    else
        QEMU_ARGS+=" -bios /usr/share/edk2/aarch64/QEMU_EFI.fd"
    fi

    QEMU_ARGS+="\
      -cpu cortex-a57 \
      -machine virt \
      -boot efi \
      "
fi
if [ "$ARCH" == "x86_64" ]; then

    if [ "$FIRMWARE" == "u-boot" ]; then
        QEMU_ARGS+=" -bios u-boot.rom"
    else
        cp /usr/share/OVMF/OVMF_VARS.fd .
        QEMU_ARGS+="\
          -drive if=pflash,format=raw,readonly,file=/usr/share/OVMF/OVMF_CODE.fd \
          -drive if=pflash,format=raw,file=OVMF_VARS.fd \
          "
    fi

    QEMU_ARGS+="\
       -accel kvm \
       -cpu host \
       -debugcon file:log.txt \
       -vga std \
       -nodefaults \
       "
fi

qemu-system-$(arch) ${QEMU_ARGS}
