# Define file/directory names
IMG_NAME="${IMG_FILENAME}${IMG_SUFFIX}"
IMG_FILE="${STAGE_WORK_DIR}/${IMG_NAME}.img"
MOUNT_DIR="${ROOTFS_DIR}" # Consistent name for the main rootfs mount point

# Define partition sizes
BOOT_SIZE=$((512 * 1024 * 1024))  # 512MB
RECOVERY_SIZE=$((256 * 1024 * 1024)) # 256MB
ALIGN=$((4 * 1024 * 1024))      # 4MB alignment

# --- Cleanup ---
echo "--- Cleaning up previous artifacts ---"
# Assume unmount_image handles unmounting potential mounts and detaching loop device
unmount_image "${IMG_FILE}"
rm -f "${IMG_FILE}"
rm -rf "${MOUNT_DIR}"
mkdir -p "${MOUNT_DIR}"

# --- Calculate Image Size ---
echo "--- Calculating required image size ---"
# Calculate rootfs size excluding caches and firmware placeholder
ROOT_SIZE=$(du -x --apparent-size -s "${EXPORT_ROOTFS_DIR}" --exclude var/cache/apt/archives --exclude boot/firmware --block-size=1 | cut -f 1)
# Add a margin for overhead and future growth (using mk-img.sh margin)
ROOT_MARGIN="$(echo "($ROOT_SIZE * 0.2 + 400 * 1024 * 1024) / 1" | bc)"

# Calculate aligned partition sizes and starts
BOOT_PART_START=$((ALIGN))
BOOT_PART_SIZE=$(((BOOT_SIZE + ALIGN - 1) / ALIGN * ALIGN))
ROOT_PART_START=$((BOOT_PART_START + BOOT_PART_SIZE))
ROOT_PART_SIZE=$(((ROOT_SIZE + ROOT_MARGIN + ALIGN - 1) / ALIGN * ALIGN))
RECOVERY_PART_START=$((ROOT_PART_START + ROOT_PART_SIZE))
RECOVERY_PART_SIZE=$(((RECOVERY_SIZE + ALIGN - 1) / ALIGN * ALIGN))
IMG_SIZE=$((RECOVERY_PART_START + RECOVERY_PART_SIZE))

echo "Rootfs Size: ${ROOT_SIZE} B"
echo "Root Margin: ${ROOT_MARGIN} B"
echo "Calculated Image Size: ${IMG_SIZE} B"

# --- Create Raw Image File ---
echo "--- Creating raw image file: ${IMG_FILE} ---"
truncate -s "${IMG_SIZE}" "${IMG_FILE}"

# --- Partition the Image ---
echo "--- Partitioning the image ---"
parted --script "${IMG_FILE}" mklabel msdos
parted --script "${IMG_FILE}" unit B set 1 boot on # Set boot flag on partition 1
parted --script "${IMG_FILE}" unit B mkpart primary fat32 "${BOOT_PART_START}" "$((BOOT_PART_START + BOOT_PART_SIZE - 1))"
parted --script "${IMG_FILE}" unit B mkpart primary ext4 "${ROOT_PART_START}" "$((ROOT_PART_START + ROOT_PART_SIZE - 1))"
parted --script "${IMG_FILE}" unit B mkpart primary ext4 "${RECOVERY_PART_START}" "$((RECOVERY_PART_START + RECOVERY_PART_SIZE - 1))"

# --- Setup Loop Device ---
echo "--- Setting up loop device ---"
cnt=0
# Assume ensure_next_loopdev() ensures a free loop device is available
until ensure_next_loopdev && LOOP_DEV="$(losetup --show --find --partscan "$IMG_FILE")"; do
    ((cnt++))
    if [ $cnt -ge 5 ]; then
        echo "ERROR: losetup failed after 5 attempts."
        exit 1
    fi
    echo "Error setting up loop device. Retrying (${cnt}/5)..."
    sleep 5
done
echo "Loop device created: ${LOOP_DEV}"
# Assume ensure_loopdev_partitions() waits for partition nodes (e.g., /dev/loopXp1) to appear
ensure_loopdev_partitions "$LOOP_DEV"
BOOT_DEV="${LOOP_DEV}p1"
ROOT_DEV="${LOOP_DEV}p2"
RECOVERY_DEV="${LOOP_DEV}p3"

# --- Format Partitions ---
echo "--- Formatting partitions ---"
# Determine EXT4 features (disable 64bit if not supported/needed, keep ^huge_file)
ROOT_FEATURES="^huge_file"
if grep -q "64bit" /etc/mke2fs.conf; then
     ROOT_FEATURES="^64bit,$ROOT_FEATURES"
     echo "Enabling 64bit feature for ext4."
fi

# Determine FAT size for boot partition
if [ "$BOOT_SIZE" -lt 134742016 ]; then # Approx 128MiB
    FAT_SIZE=16
    echo "Using FAT16 for boot partition."
else
    FAT_SIZE=32
    echo "Using FAT32 for boot partition."
fi

mkdosfs -n BOOT -F "$FAT_SIZE" -s 4 -v "$BOOT_DEV" > /dev/null
mkfs.ext4 -L rootfs -O "$ROOT_FEATURES" "$ROOT_DEV" > /dev/null
mkfs.ext4 -L recovery -O "$ROOT_FEATURES" "$RECOVERY_DEV" > /dev/null # Format recovery for all archs

# --- Architecture-Specific Steps ---

if [ "$ARCH" == "amd64" ]; then
    # --- AMD64: Create IMG, Install GRUB, Create ISO ---
    echo "--- Running AMD64 specific steps (IMG + GRUB + ISO) ---"

    ISO_STAGING_DIR="${STAGE_WORK_DIR}/iso_staging" # Temp dir for ISO contents

    # Mount partitions for AMD64 (boot on /boot)
    echo "Mounting partitions for AMD64..."
    mount -v "$ROOT_DEV" "${MOUNT_DIR}" -t ext4
    mkdir -p "${MOUNT_DIR}/boot"
    mount -v "$BOOT_DEV" "${MOUNT_DIR}/boot" -t vfat
    # Optional: Mount recovery if you need to add specific recovery tools to the ISO
    # mkdir -p "${MOUNT_DIR}/recovery"
    # mount -v "$RECOVERY_DEV" "${MOUNT_DIR}/recovery" -t ext4

    # Copy root filesystem content (excluding caches, firmware, overlays which are typically ARM specific)
    echo "Copying root filesystem..."
    rsync -aHAXx --delete --exclude /var/cache/apt/archives --exclude /boot/firmware --exclude /boot/overlays "${EXPORT_ROOTFS_DIR}/" "${MOUNT_DIR}/"

    # Install GRUB bootloader for BIOS boot
    echo "Installing GRUB (BIOS)..."
    mkdir -p "${MOUNT_DIR}/boot/grub" # Ensure grub dir exists
    # Basic GRUB config - **IMPORTANT**: Adjust kernel/initrd paths and root= parameter if needed!
    # Using PARTUUID is generally more reliable than /dev/sdxN
    ROOT_PARTUUID=$(blkid -s PARTUUID -o value "$ROOT_DEV")
    cat > "${MOUNT_DIR}/boot/grub/grub.cfg" <<EOF
set default=0
set timeout=5
set gfxpayload=keep

menuentry "Start OS" {
    search --no-floppy --fs-uuid --set=root \$(blkid -s UUID -o value "$ROOT_DEV") || search --no-floppy --part-uuid --set=root ${ROOT_PARTUUID}
    linux /boot/vmlinuz root=PARTUUID=${ROOT_PARTUUID} quiet splash # Adjust kernel name and boot params as needed
    initrd /boot/initrd.img # Adjust initrd name as needed
}
EOF

    grub-install \
        --target=i386-pc \
        --boot-directory="${MOUNT_DIR}/boot" \
        --modules="part_msdos ext2 fat normal biosdisk search search_fs_uuid search_partuuid" \
        --force \
        --no-floppy \
        "$LOOP_DEV" # Install to the base loop device MBR

    # Optional: Add UEFI support here if needed
    # Requires creating a FAT filesystem image (e.g., efi.img) containing the EFI bootloader
    # mkdir -p "${MOUNT_DIR}/boot/efi/EFI/boot"
    # cp path/to/bootx64.efi "${MOUNT_DIR}/boot/efi/EFI/boot/" # Copy your UEFI bootloader

    # Prepare for ISO creation: Copy necessary files to a staging area
    echo "Preparing ISO staging area..."
    rm -rf "${ISO_STAGING_DIR}"
    mkdir -p "${ISO_STAGING_DIR}"

    # Check if El Torito boot image exists
    ELTORITO_IMG="${MOUNT_DIR}/boot/grub/i386-pc/eltorito.img"
    if [ ! -f "$ELTORITO_IMG" ]; then
        echo "WARNING: GRUB El Torito image not found at $ELTORITO_IMG. ISO may not boot."
        # Attempt to create a basic one if missing? Might require grub-mkrescue logic.
    fi

    # Copy contents from mounted image to the staging directory
    # Ensure the structure inside ISO_STAGING_DIR matches what xorriso expects
    echo "Copying files to ISO staging directory..."
    rsync -aH "${MOUNT_DIR}/" "${ISO_STAGING_DIR}/"

    # Unmount partitions before creating ISO
    echo "Unmounting partitions..."
    # if [ -d "${MOUNT_DIR}/recovery" ]; then umount -l "${MOUNT_DIR}/recovery"; fi
    umount -l "${MOUNT_DIR}/boot"
    umount -l "${MOUNT_DIR}"

    # Create Hybrid ISO from the staging directory
    echo "Creating Hybrid ISO file..."
    ISO_FILE="${STAGE_WORK_DIR}/${IMG_NAME}.iso"
    # Basic hybrid ISO command (BIOS boot via GRUB El Torito)
    # Add -isohybrid-mbr, -c, -eltorito-alt-boot, -e, -isohybrid-gpt-basdat for advanced features (UEFI, etc.)
    xorriso -as mkisofs \
      -r -J -l \
      -iso-level 3 \
      -appid "My Custom OS" \
      -pubset "My Publisher" \
      -volid "MY_OS_INSTALL" \
      -b boot/grub/i386-pc/eltorito.img `# Path relative to ISO_STAGING_DIR` \
      -no-emul-boot \
      -boot-load-size 4 \
      -boot-info-table \
      -isohybrid-mbr /usr/lib/ISOLINUX/isohdpfx.bin `# Use syslinux MBR for hybrid` \
      -output "${ISO_FILE}" \
      "${ISO_STAGING_DIR}"

    echo "AMD64 IMG created: ${IMG_FILE}"
    echo "AMD64 Bootable ISO created: ${ISO_FILE}"
    echo "Cleaning up ISO staging directory..."
    rm -rf "${ISO_STAGING_DIR}"

elif [[ "$ARCH" == "armhf" || "$ARCH" == "arm64" ]]; then
    # --- ARM: Create IMG, Copy Firmware/Overlays ---
    echo "--- Running ARM specific steps (IMG + Firmware/Overlays) ---"

    # Mount partitions for ARM (boot on /boot/firmware)
    echo "Mounting partitions for ARM..."
    mount -v "$ROOT_DEV" "${MOUNT_DIR}" -t ext4
    mkdir -p "${MOUNT_DIR}/boot/firmware" # Specific mount point for boot partition
    mount -v "$BOOT_DEV" "${MOUNT_DIR}/boot/firmware" -t vfat

    # Copy root filesystem content (excluding caches and the target firmware dir itself)
    echo "Copying root filesystem..."
    rsync -aHAXx --delete --exclude /var/cache/apt/archives --exclude /boot/firmware "${EXPORT_ROOTFS_DIR}/" "${MOUNT_DIR}/"

    # Specifically copy firmware and overlays content
    echo "Copying boot firmware and overlays..."
    # Copy everything from EXPORT_ROOTFS_DIR/boot/firmware (which might contain kernel, overlays, config.txt etc.)
    if [ -d "${EXPORT_ROOTFS_DIR}/boot/firmware" ]; then
         rsync -rtx --delete "${EXPORT_ROOTFS_DIR}/boot/firmware/" "${MOUNT_DIR}/boot/firmware/"
    else
         echo "Warning: Source directory ${EXPORT_ROOTFS_DIR}/boot/firmware does not exist. No firmware copied."
    fi
    # Handle overlays if they are outside /boot/firmware in the source
    if [ -d "${EXPORT_ROOTFS_DIR}/boot/overlays" ]; then
        echo "Copying overlays from ${EXPORT_ROOTFS_DIR}/boot/overlays..."
        mkdir -p "${MOUNT_DIR}/boot/overlays"
        rsync -rtx --delete "${EXPORT_ROOTFS_DIR}/boot/overlays/" "${MOUNT_DIR}/boot/overlays/"
    fi

    # Unmount partitions
    echo "Unmounting partitions..."
    umount -l "${MOUNT_DIR}/boot/firmware"
    umount -l "${MOUNT_DIR}"

    echo "ARM Image created: ${IMG_FILE}"
    echo "Includes boot (/boot/firmware), rootfs, recovery partition, firmware, and overlays."

else
    # --- Unsupported Architecture ---
    echo "ERROR: Unsupported architecture: $ARCH"
    echo "Cannot proceed. Cleaning up..."
    # Attempt cleanup even on error
    losetup -d "$LOOP_DEV" || echo "Warning: Failed to detach loop device $LOOP_DEV on error."
    rm -rf "${MOUNT_DIR}"
    exit 1
fi

# --- Final Cleanup ---
echo "--- Finalizing and cleaning up ---"
losetup -d "$LOOP_DEV" || echo "Loop device ${LOOP_DEV} already detached."
# MOUNT_DIR should be empty/gone now, but double-check
if [ -d "${MOUNT_DIR}" ]; then
    # If it still exists and is a mount point, something went wrong
    if findmnt -rno TARGET "${MOUNT_DIR}"; then
      echo "WARNING: ${MOUNT_DIR} appears to still be mounted."
    else
      rmdir "${MOUNT_DIR}" || echo "Warning: Could not remove mount directory ${MOUNT_DIR}."
    fi
fi