#!/usr/bin/env bash
# shellcheck disable=SC2054  # Arrays are formatted for passing args to other tools

#
# Common error handling
#
#
exit_trap_cmds=()

on_exit() {
    exit_trap_cmds+=( "$1" )
}

run_exit_trap_cmds() {
    for cmd in "${exit_trap_cmds[@]}"; do
        eval "${cmd}"
    done
}

trap run_exit_trap_cmds exit

bail() {
    >&2 echo "$@"
    exit 1
}

shopt -s nullglob

arch=${BUILDSYS_ARCH}
variant=${BUILDSYS_VARIANT}
product_name=${BUILDSYS_NAME:-bottlerocket}
host_port_forwards=tcp::2222-:22
vm_mem=4G
vm_cpus=4
force_extract=
declare -A extra_files=()

boot_image=
data_image=


if ! git_toplevel=$(git rev-parse --show-toplevel); then
    bail "Failed to get the root of the repo."
else
    readonly repo_root="${git_toplevel}"
fi

show_usage() {
    echo "\
usage: ${0##*/} [--arch BUILDSYS_ARCH] [--variant BUILDSYS_VARIANT]
                      [--host-port-forwards HOST_PORT_FWDS]
                      [--vm-memory VM_MEMORY] [--vm-cpus VM_CPUS]
                      [--inject-file LOCAL_PATH[:IMAGE_PATH]]...

Launch a local virtual machine from a Bottlerocket image.

Options:

    --arch              architecture of the Bottlerocket image (must match the
                        host architecture ($(uname -m)); may be omitted if the
                        BUILDSYS_ARCH environment variable is set)
    --variant           Bottlerocket variant to run (may be omitted if the
                        BUILDSYS_VARIANT environment variable is set)
    --product-name
                        product name used for file and directory naming used when
                        building with the "-e BUILDSYS_NAME" option; may be omitted if the
                        BUILDSYS_NAME environment variable is set. Otherwise default is bottlerocket if not defined or empty
    --host-port-forwards
                        list of host ports to forward to the VM; HOST_PORT_FWDS
                        must be a valid QEMU port forwarding specifier (default
                        is ${host_port_forwards})
    --vm-memory         amount of memory to assign to the VM; VM_MEMORY must be
                        a valid QEMU memory specifier (default is ${vm_mem})
    --vm-cpus           number of CPUs to spawn for VM (default is ${vm_cpus})
    --force-extract     force recreation of the extracted Bottlerocket image,
                        e.g. to force first boot behavior
    --inject-file       adds a local file to the private partition of the
                        Bottlerocket image before launching the virtual machine
                        (may be given multiple times); existing data on the
                        private partition will be lost
    --firmware-code     override the default firmware executable file
    --firmware-vars     override the initial firmware variable storage file
    --help              shows this usage text

By default, the virtual machine's port 22 (SSH) will be exposed via the local
port 2222, i.e. if the Bottlerocket admin container has been enabled via
user-data, it can be reached by running

    ssh -p 2222 ec2-user@localhost

from the host.

Usage example:

    ${0##*/} --arch $(uname -m) --variant metal-dev --inject-file net.toml
"
}

usage_error() {
    local error=$1

    {
        if [[ -n ${error} ]]; then
            printf "%s\n\n" "${error}"
        fi
        show_usage
    } >&2

    exit 1
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_usage; exit 0 ;;
            --arch)
                shift; arch=$1 ;;
            --variant)
                shift; variant=$1 ;;
            --product-name)
                shift; product_name=$1 ;;
            --host-port-forwards)
                shift; host_port_forwards=$1 ;;
            --vm-memory)
                shift; vm_mem=$1 ;;
            --vm-cpus)
                shift; vm_cpus=$1 ;;
            --force-extract)
                force_extract=yes ;;
            --inject-file)
                shift; local file_spec=$1
                if [[ ${file_spec} = *:* ]]; then
                    local local_file=${file_spec%%:*}
                    local image_file=${file_spec#*:}
                else
                    local local_file=${file_spec}
                    local image_file=${file_spec##*/}
                fi
                extra_files[${local_file}]=${image_file}
                ;;
            --firmware-code)
                shift; firmware_code=$1
                ;;
            --firmware-vars)
                shift; firmware_vars=$1
                ;;
            *)
                usage_error "unknown option '$1'" ;;
        esac
        shift
    done

    [[ -n ${arch} ]] || usage_error 'Architecture needs to be set via either --arch or BUILDSYS_ARCH.'
    [[ -n ${variant} ]] || usage_error 'Variant needs to be set via either --variant or BUILDSYS_VARIANT.'

    declare -l host_arch
    host_arch=$(uname -m)
    [[ ${arch} == "${host_arch}" ]] || bail "Architecture needs to match host architecture (${host_arch}) for hardware virtualization."

    for path in "${!extra_files[@]}"; do
        [[ -e ${path} ]] || bail "Cannot find local file '${path}' to inject."
    done
}

extract_image() {
    local -r compressed_image=$1
    local -r uncompressed_image=$2

    if [[ ${force_extract} = yes ]] || [[ ${compressed_image} -nt ${uncompressed_image} ]]; then
        lz4 --decompress --force --keep "${compressed_image}" "${uncompressed_image}" \
            || bail "Failed to extract '${compressed_image}'."
    fi
}

prepare_raw_images() {
    local -r image_dir=build/images/${arch}-${variant}/latest
    local -r compressed_boot_image=${image_dir}/${product_name}-${variant}-${arch}.img.lz4
    local -r compressed_data_image=${image_dir}/${product_name}-${variant}-${arch}-data.img.lz4

    if [[ -e ${compressed_boot_image} ]]; then
        readonly boot_image=${compressed_boot_image%*.lz4}
        extract_image "${compressed_boot_image}" "${boot_image}"
    else
        bail 'Boot image not found. Did the last build fail?'
    fi

    if [[ -e ${compressed_data_image} ]]; then
        readonly data_image=${compressed_data_image%*.lz4}
        extract_image "${compressed_data_image}" "${data_image}"
    else
        # Missing data image is fine. This variant may not be a split build.
        readonly data_image=
    fi
}

prepare_firmware() {
    # Create local copies of the edk2 firmware variable storage, to help with
    # faciliate Secure Boot testing where custom variables are needed for both
    # architectures, but can't safely be reused across QEMU invocations. Also
    # set reasonable defaults for both firmware files, if nothing more specific
    # was requested.
    local original_vars

    if [[ ${arch} = x86_64 ]]; then
        firmware_code=${firmware_code:-/usr/share/edk2/ovmf/OVMF_CODE.fd}
        original_vars=${firmware_vars:-/usr/share/edk2/ovmf/OVMF_VARS.fd}
        firmware_vars="$(mktemp)"
        on_exit "rm '${firmware_vars}'"
        cp "${original_vars}" "${firmware_vars}"
    fi

    if [[ ${arch} = aarch64 ]]; then
        original_code=${firmware_code:-/usr/share/edk2/aarch64/QEMU_EFI.silent.fd}
        original_vars=${firmware_vars:-/usr/share/edk2/aarch64/QEMU_VARS.fd}
        firmware_code="$(mktemp)"
        firmware_vars="$(mktemp)"
        on_exit "rm '${firmware_code}' '${firmware_vars}'"
        cat "${original_code}" /dev/zero \
            | head -c 64m > "${firmware_code}"
        cat "${original_vars}" /dev/zero \
            | head -c 64m > "${firmware_vars}"
    fi
}

create_extra_files() {
    # Explicitly instruct the kernel to send its output to the serial port on
    # x86 via a bootconfig initrd. Passing in settings via user-data would be
    # too late to get console output of the first boot.
    if [[ ${arch} = x86_64 ]]; then
        extra_files["${repo_root}/tools/bootconfig/qemu-x86-console-bootconfig.data"]=bootconfig.data
    fi

    # If the private partition needs to be recreated, ensure that any bootconfig
    # data file is present, otherwise GRUB will notice the missing file and wait
    # for a key press.
    if [[ ${#extra_files[@]} -gt 0 ]]; then
        local has_bootconfig=no
        for image_file in "${extra_files[@]}"; do
            if [[ ${image_file} = bootconfig.data ]]; then
                has_bootconfig=yes
                break
            fi
        done
        if [[ ${has_bootconfig} = no ]]; then
            extra_files["${repo_root}/tools/bootconfig/empty-bootconfig.data"]=bootconfig.data
        fi
    fi
}

inject_files() {
    if [[ ${#extra_files[@]} -eq 0 ]]; then
        return 0
    fi

    # We inject files into the boot image by replacing the private partition
    # entirely. The new partition has to perfectly fit over the original one.
    # Find the first and last sector, then calculate the partition's size. In
    # absence of actual hardware, assume a traditional sector size of 512 bytes.
    local private_first_sector private_last_sector
    read -r private_first_sector private_last_sector < <(
        fdisk --list-details "${boot_image}" \
            | awk '/BOTTLEROCKET-PRIVATE/ { print $2, $3 }')
    if [[ -z ${private_first_sector} ]] || [[ -z ${private_last_sector} ]]; then
        bail "Failed to find the private partition in '${boot_image}'."
    fi
    local private_size_mib=$(( (private_last_sector - private_first_sector + 1) * 512 / 1024 / 1024 ))

    local private_mount private_image
    private_mount=$(mktemp -d)
    private_image=$(mktemp)
    on_exit "rm -rf '${private_mount}' '${private_image}'"

    for local_file in "${!extra_files[@]}"; do
        local image_file=${extra_files[${local_file}]}
        cp "${local_file}" "${private_mount}/${image_file}"
    done

    if ! mkfs.ext4 -d "${private_mount}" "${private_image}" "${private_size_mib}M" \
    || ! dd if="${private_image}" of="${boot_image}" conv=notrunc bs=512 seek="${private_first_sector}"
    then
        rm -f "${private_image}"
        rm -rf "${private_mount}"
        bail "Failed to inject files into '${boot_image}'."
    fi
}

launch_vm() {
    local -a qemu_args=(
        -nographic
        -enable-kvm
        -cpu host
        -smp "${vm_cpus}"
        -m "${vm_mem}"
        -drive if=pflash,format=raw,unit=0,file="${firmware_code}",readonly=on
        -drive if=pflash,format=raw,unit=1,file="${firmware_vars}"
        -drive index=0,if=virtio,format=raw,file="${boot_image}"
    )

    # Plug the virtual primary NIC in as BDF 00:10.0 so udev will give it a
    # consistent name we can know ahead of time--enp0s16 or ens16.
    qemu_args+=(
        -netdev user,id=net0,hostfwd="${host_port_forwards}"
        -device virtio-net-pci,netdev=net0,addr=10.0
    )

    # Resolve the last bit of uncertainty by disabling ACPI-based PCI hot plug,
    # causing udev to use the bus location when naming the NIC (enp0s16). Since
    # QEMU does not support PCI hot plug via ACPI on Arm, turn it off for the
    # emulated x86_64 chipset only to achieve parity.
    if [[ ${arch} = x86_64 ]]; then
        qemu_args+=( -global PIIX4_PM.acpi-root-pci-hotplug=off )
        qemu_args+=( -machine q35,smm=on )
    fi

    if [[ ${arch} = aarch64 ]]; then
        qemu_args+=( -machine virt )
    fi

    if [[ -n ${data_image} ]]; then
        qemu_args+=( -drive index=1,if=virtio,format=raw,file="${data_image}" )
    fi

    qemu-system-"${arch}" "${qemu_args[@]}"
}

parse_args "$@"
prepare_raw_images
prepare_firmware
create_extra_files
inject_files
launch_vm