#!/bin/bash
set -e

cleanup() {
    echo "Error occurred, cleaning up."
    sudo umount -v "${mount_directory}" || true  # Unmount if mounted
    sudo rm -rv "${mount_directory}" || true     # Safely remove mount directory if it exists
    sudo losetup -d "${loop_device}" || true     # Detach loop device if attached
    exit 1
}

trap cleanup ERR

# Build the project
cargo build --release

# Check if build was successful
if [ $? -eq 0 ]; then
    # If successful, run the emulator script
    echo "Building disk.img"
else
    echo "Build failed, aborting."
    exit 1
fi

script_directory=$(dirname "${BASH_SOURCE[0]}")
script_directory=$(realpath "$script_directory")
disk_image="$HOME/tmpfs/disk.img"
block_size="1M"
count="100"
disk_size="100MiB"
bootloader="${script_directory}/../target/x86_64-unknown-uefi/release/os_loader.efi"

# Get the size of the bootloader in bytes
bootloader_size=$(stat -c%s "$bootloader")

# Calculate the required disk image size with overhead
# Adding 20% overhead for FAT32 filesystem
overhead_percentage=20
overhead_size=$((bootloader_size * overhead_percentage / 100))
total_size=$((bootloader_size + overhead_size))

# Ensure the size is at least 4MB
minimum_size=$((4 * 1024 * 1024))

# Calculate the final size, which should be the maximum of total_size or minimum_size
final_size=$(( total_size > minimum_size ? total_size : minimum_size ))

# Ensure the size is aligned to a 1MB boundary
block_size=$((1024 * 1024))
count=$(( (final_size + block_size - 1) / block_size ))  # Round up to nearest MB

echo "Creating disk image with size $((count * block_size)) bytes ($count MB)"

# Create the disk image
dd if=/dev/zero of="$disk_image" bs="$block_size" count="$count" status=progress

loop_device=$(sudo losetup -fP "$disk_image" --show)
echo "Using loop device "$loop_device"."
sudo parted --script "$loop_device" mklabel gpt
sudo parted --script "$loop_device" mkpart ESP fat32 1MiB 100%
sudo parted --script "$loop_device" set 1 esp on
sudo mkfs.fat -F 32 "${loop_device}p1" # Ensure the partition number is correct

mount_directory=$(mktemp --tmpdir="$script_directory" -d)
sudo mount -v "${loop_device}p1" "$mount_directory"
sudo mkdir -pv "${mount_directory}/EFI/BOOT/"
sudo cp -vf "$bootloader" "${mount_directory}/EFI/BOOT/BOOTX64.EFI"
sudo umount -v "${mount_directory}"

echo "Partition and formatting ✓"
sudo rm -rv "$mount_directory"
sudo losetup -d "$loop_device"

qemu-system-x86_64 -enable-kvm -nodefaults -vga std -monitor vc:1024x768 -machine q35,accel=kvm -m 1024 -cpu host -drive if=pflash,format=raw,readonly=on,file="${script_directory}/../docs/OVMF_CODE.fd" -drive if=pflash,format=raw,file="${script_directory}/../docs/OVMF_VARS.fd" -drive file="$disk_image",format=raw,index=0,media=disk -boot order=c