#!/usr/bin/env bash
set -euo pipefail

repo_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
target_dir="${repo_dir}/target/x86_64-unknown-uefi/debug"
esp_dir="${repo_dir}/target/esp"
disk_image="${repo_dir}/target/zeroos.raw"
esp_image="${repo_dir}/target/zeroos-esp.fat"
root_image="${repo_dir}/target/zeroos-root.ext4"
root_dir="${repo_dir}/target/rootfs"
ovmf_vars="${repo_dir}/target/OVMF_VARS.4m.fd"

ovmf_code="/usr/share/edk2/x64/OVMF_CODE.4m.fd"
ovmf_vars_template="/usr/share/edk2/x64/OVMF_VARS.4m.fd"
qemu_display="${QEMU_DISPLAY:-gtk}"
qemu_serial="${QEMU_SERIAL:-stdio}"
qemu_extra_args="${QEMU_EXTRA_ARGS:--no-reboot}"
qemu_accel="${QEMU_ACCEL:-auto}"
qemu_debug_log="${QEMU_DEBUG_LOG:-${repo_dir}/target/qemu-debug.log}"
zeroos_sh="${ZEROOS_SH:-/home/archzero/C++/busybox/busybox}"
rebuild_disk="${ZEROOS_REBUILD_DISK:-0}"
update_rootfs="${ZEROOS_UPDATE_ROOTFS:-0}"

esp_start=2048
esp_sectors=131072
root_start=$((esp_start + esp_sectors))
root_sectors=389120
esp_size=$((esp_sectors * 512))
root_size=$((root_sectors * 512))
disk_size=$(((root_start + root_sectors + 2048) * 512))

if [[ ! -r "${ovmf_code}" || ! -r "${ovmf_vars_template}" ]]; then
    echo "OVMF firmware not found under /usr/share/edk2/x64." >&2
    echo "Install the OVMF/edk2 package for your distribution." >&2
    exit 1
fi

cargo build --manifest-path "${repo_dir}/Cargo.toml" --target x86_64-unknown-uefi --workspace

mkdir -p "${esp_dir}/EFI/BOOT" "${esp_dir}/EFI/ZeroOS"
cp "${target_dir}/zeroos-loader.efi" "${esp_dir}/EFI/BOOT/BOOTX64.EFI"
cp "${target_dir}/ZeroOS.efi" "${esp_dir}/EFI/ZeroOS/ZeroOS.EFI"

truncate -s "${esp_size}" "${esp_image}"
mformat -i "${esp_image}" -F ::
mmd -i "${esp_image}" ::/EFI ::/EFI/BOOT ::/EFI/ZeroOS
mcopy -i "${esp_image}" "${esp_dir}/EFI/BOOT/BOOTX64.EFI" ::/EFI/BOOT/BOOTX64.EFI
mcopy -i "${esp_image}" "${esp_dir}/EFI/ZeroOS/ZeroOS.EFI" ::/EFI/ZeroOS/ZeroOS.EFI

if [[ ! -f "${disk_image}" ]]; then
    rebuild_disk=1
fi

if [[ "${rebuild_disk}" == "1" || "${update_rootfs}" == "1" ]]; then
    if [[ ! -r "${zeroos_sh}" ]]; then
        echo "sh source is not readable: ${zeroos_sh}" >&2
        echo "Set ZEROOS_SH=/path/to/static/busybox-or-sh." >&2
        exit 1
    fi

    rm -rf "${root_dir}"
    mkdir -p "${root_dir}/bin" "${root_dir}/usr/bin" "${root_dir}/dev" "${root_dir}/etc" "${root_dir}/tmp"
    cp "${zeroos_sh}" "${root_dir}/usr/bin/busybox"
    while IFS= read -r applet; do
        [[ -n "${applet}" ]] || continue
        if [[ "${applet}" != "busybox" ]]; then
            ln -f "${root_dir}/usr/bin/busybox" "${root_dir}/usr/bin/${applet}"
        fi
        ln -f "${root_dir}/usr/bin/busybox" "${root_dir}/bin/${applet}"
    done < <("${zeroos_sh}" --list)
    ln -f "${root_dir}/usr/bin/busybox" "${root_dir}/usr/bin/bash"
    ln -f "${root_dir}/usr/bin/busybox" "${root_dir}/bin/bash"

    truncate -s "${root_size}" "${root_image}"
    mkfs.ext4 -q -F -b 4096 -O '^64bit,^metadata_csum,^has_journal,^dir_index' -d "${root_dir}" "${root_image}"
fi

if [[ "${rebuild_disk}" == "1" ]]; then
    truncate -s "${disk_size}" "${disk_image}"
    sgdisk --clear \
        --new=1:${esp_start}:$((esp_start + esp_sectors - 1)) \
        --typecode=1:ef00 \
        --change-name=1:ZeroOSESP \
        --new=2:${root_start}:$((root_start + root_sectors - 1)) \
        --typecode=2:8300 \
        --change-name=2:ZeroOSRoot \
        "${disk_image}" >/dev/null
    dd if="${root_image}" of="${disk_image}" bs=512 seek="${root_start}" conv=notrunc status=none
elif [[ "${update_rootfs}" == "1" ]]; then
    dd if="${root_image}" of="${disk_image}" bs=512 seek="${root_start}" conv=notrunc status=none
fi

dd if="${esp_image}" of="${disk_image}" bs=512 seek="${esp_start}" conv=notrunc status=none

if [[ ! -f "${ovmf_vars}" ]]; then
    cp "${ovmf_vars_template}" "${ovmf_vars}"
fi

extra_args=()
if [[ -n "${qemu_extra_args}" ]]; then
    read -r -a extra_args <<< "${qemu_extra_args}"
fi

debug_args=()
if [[ "${QEMU_DEBUG:-1}" != "0" ]]; then
    debug_args=(-d int,cpu_reset -D "${qemu_debug_log}")
    rm -f "${qemu_debug_log}"
fi

accel_args=()
case "${qemu_accel}" in
    auto)
        if [[ -r /dev/kvm && -w /dev/kvm ]]; then
            accel_args=(-enable-kvm -cpu host)
        fi
        ;;
    kvm)
        accel_args=(-enable-kvm -cpu host)
        ;;
    tcg|off|none)
        ;;
    *)
        echo "invalid QEMU_ACCEL: ${qemu_accel}" >&2
        echo "Use QEMU_ACCEL=auto, kvm, or tcg." >&2
        exit 1
        ;;
esac

qemu-system-x86_64 \
    -machine q35 \
    "${accel_args[@]}" \
    -m 256M \
    -drive "if=pflash,format=raw,readonly=on,file=${ovmf_code}" \
    -drive "if=pflash,format=raw,file=${ovmf_vars}" \
    -device "ich9-ahci,id=ahci0" \
    -drive "id=zeroosdisk,if=none,format=raw,file=${disk_image}" \
    -device "ide-hd,drive=zeroosdisk,bus=ahci0.0" \
    -serial "${qemu_serial}" \
    -display "${qemu_display}" \
    "${debug_args[@]}" \
    "${extra_args[@]}"
status=$?
echo "qemu exited with status ${status}" >&2
if [[ -f "${qemu_debug_log}" ]]; then
    echo "qemu debug log: ${qemu_debug_log}" >&2
fi
exit "${status}"