#!/bin/bash
set -euo pipefail

IMG="boot.img"
IMG_SIZE_MB=4
TARGET="riscv64gc-unknown-none-elf"
BUILD_DIR="target/${TARGET}/release"

echo "=== Creating k-os boot disk image ==="

# Create empty image
dd if=/dev/zero of="${IMG}" bs=1M count=${IMG_SIZE_MB} 2>/dev/null
echo "Created ${IMG} (${IMG_SIZE_MB}MB)"

# Format as FAT16
# On macOS: use newfs_msdos, on Linux: use mkfs.fat
if command -v newfs_msdos &>/dev/null; then
    newfs_msdos -F 16 "${IMG}" >/dev/null
elif command -v mkfs.fat &>/dev/null; then
    mkfs.fat -F 16 "${IMG}" >/dev/null
else
    echo "Error: need newfs_msdos (macOS) or mkfs.fat (Linux)"
    exit 1
fi
echo "Formatted as FAT16"

# Mount and copy files
MOUNT_DIR=$(mktemp -d)

if [[ "$(uname)" == "Darwin" ]]; then
    # macOS: use hdiutil
    DEVICE=$(hdiutil attach -nomount "${IMG}" 2>/dev/null | head -1 | awk '{print $1}')
    mount -t msdos "${DEVICE}" "${MOUNT_DIR}" 2>/dev/null || {
        echo "Warning: could not mount image. Files must be copied manually."
        echo "On macOS, try: brew install mtools"
        hdiutil detach "${DEVICE}" 2>/dev/null
        rmdir "${MOUNT_DIR}"

        # Fallback: use mtools if available
        if command -v mcopy &>/dev/null; then
            copy_with_mtools
        else
            echo "Install mtools: brew install mtools"
            exit 0
        fi
        exit 0
    }

    copy_files "${MOUNT_DIR}"
    umount "${MOUNT_DIR}"
    hdiutil detach "${DEVICE}" 2>/dev/null
else
    # Linux: use mount
    sudo mount -o loop "${IMG}" "${MOUNT_DIR}"
    copy_files "${MOUNT_DIR}"
    sudo umount "${MOUNT_DIR}"
fi

rmdir "${MOUNT_DIR}" 2>/dev/null
echo "=== Done: ${IMG} ==="

copy_files() {
    local DIR="$1"
    mkdir -p "${DIR}/servers" "${DIR}/apps"

    # Copy init
    if [ -f "${BUILD_DIR}/hello" ]; then
        cp "${BUILD_DIR}/hello" "${DIR}/init.elf"
        echo "  Copied init.elf"
    fi

    # Copy any server binaries that exist
    for server in display input log permission storage session; do
        if [ -f "${BUILD_DIR}/${server}-server" ]; then
            cp "${BUILD_DIR}/${server}-server" "${DIR}/servers/${server}.elf"
            echo "  Copied servers/${server}.elf"
        fi
    done

    # Copy any app binaries that exist
    for app in snake; do
        if [ -f "${BUILD_DIR}/${app}" ]; then
            cp "${BUILD_DIR}/${app}" "${DIR}/apps/${app}.elf"
            echo "  Copied apps/${app}.elf"
        fi
    done
}

copy_with_mtools() {
    echo "Using mtools to copy files..."
    mmd -i "${IMG}" ::servers ::apps 2>/dev/null || true

    if [ -f "${BUILD_DIR}/hello" ]; then
        mcopy -i "${IMG}" "${BUILD_DIR}/hello" "::init.elf"
        echo "  Copied init.elf"
    fi

    for server in display input log permission storage session; do
        if [ -f "${BUILD_DIR}/${server}-server" ]; then
            mcopy -i "${IMG}" "${BUILD_DIR}/${server}-server" "::servers/${server}.elf"
            echo "  Copied servers/${server}.elf"
        fi
    done

    for app in snake; do
        if [ -f "${BUILD_DIR}/${app}" ]; then
            mcopy -i "${IMG}" "${BUILD_DIR}/${app}" "::apps/${app}.elf"
            echo "  Copied apps/${app}.elf"
        fi
    done
}