#!/bin/bash -e

IMG_NAME="${IMG_FILENAME}${IMG_SUFFIX}"
IMG_FILE="${STAGE_WORK_DIR}/${IMG_NAME}.img"
MOUNT_DIR="${ROOTFS_DIR}"
BOOT_SIZE=$((512 * 1024 * 1024))
EFI_SIZE=$((128 * 1024 * 1024)) # 128MB EFI partition
RECOVERY_SIZE=$((256 * 1024 * 1024))
ALIGN=$((4 * 1024 * 1024))

unmount_image "${IMG_FILE}"
rm -f "${IMG_FILE}"
rm -rf "${MOUNT_DIR}"
mkdir -p "${MOUNT_DIR}"

# Calculate rootfs size
ROOT_SIZE=$(du -x --apparent-size -s "${EXPORT_ROOTFS_DIR}" --exclude var/cache/apt/archives --exclude boot/firmware --block-size=1 | cut -f 1)
ROOT_MARGIN="$(echo "($ROOT_SIZE * 0.2 + 200 * 1024 * 1024) / 1" | bc)"
BOOT_PART_START=$((ALIGN))
BOOT_PART_SIZE=$(((BOOT_SIZE + ALIGN - 1) / ALIGN * ALIGN))
EFI_PART_START=$((BOOT_PART_START + BOOT_PART_SIZE))
EFI_PART_SIZE=$(((EFI_SIZE + ALIGN - 1) / ALIGN * ALIGN))
ROOT_PART_START=$((EFI_PART_START + EFI_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=$((BOOT_PART_START + BOOT_PART_SIZE + EFI_PART_SIZE + ROOT_PART_SIZE + RECOVERY_PART_SIZE))

# Create raw image
truncate -s "${IMG_SIZE}" "${IMG_FILE}"
parted --script "${IMG_FILE}" mklabel gpt
parted --script "${IMG_FILE}" unit B mkpart primary fat32 "${BOOT_PART_START}" "$((BOOT_PART_START + BOOT_PART_SIZE - 1))"
parted --script "${IMG_FILE}" set 1 boot on
parted --script "${IMG_FILE}" unit B mkpart primary fat32 "${EFI_PART_START}" "$((EFI_PART_START + EFI_PART_SIZE - 1))"
parted --script "${IMG_FILE}" set 2 esp on
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))"

# Create loop device
cnt=0
until ensure_next_loopdev && LOOP_DEV="$(losetup --show --find --partscan "$IMG_FILE")"; do
    ((cnt++))
    if [ $cnt -ge 5 ]; then echo "ERROR: losetup failed"; exit 1; fi
    echo "Retrying losetup..."; sleep 5
done
ensure_loopdev_partitions "$LOOP_DEV"

BOOT_DEV="${LOOP_DEV}p1"
EFI_DEV="${LOOP_DEV}p2"
ROOT_DEV="${LOOP_DEV}p3"
RECOVERY_DEV="${LOOP_DEV}p4"

# Format partitions
FAT_SIZE=$([ "$BOOT_SIZE" -lt 134742016 ] && echo 16 || echo 32)
mkdosfs -n bootfs -F "$FAT_SIZE" -s 4 "$BOOT_DEV"
mkfs.fat -F 32 -n EFI "$EFI_DEV"
mkfs.ext4 -L rootfs -O "^64bit,^huge_file" "$ROOT_DEV"
mkfs.ext4 -L recovery "$RECOVERY_DEV"


# Mount and copy rootfs
mount "$ROOT_DEV" "$MOUNT_DIR" -t ext4
mkdir -p "${MOUNT_DIR}/boot"
mount "$BOOT_DEV" "${MOUNT_DIR}/boot" -t vfat
mkdir -p "${MOUNT_DIR}/boot/efi"
mount "$EFI_DEV" "${MOUNT_DIR}/boot/efi" -t vfat
rsync -aHAXx --exclude /var/cache/apt/archives --exclude /boot/firmware --exclude /boot/overlays "${EXPORT_ROOTFS_DIR}/" "${MOUNT_DIR}/"

# Recovery partition
mkdir -p "${MOUNT_DIR}/recovery"
mount "$RECOVERY_DEV" "${MOUNT_DIR}/recovery" -t ext4
# Optionally copy recovery tools
# cp recovery-tools/recovery.sh "${MOUNT_DIR}/recovery/"

# Install GRUB (BIOS)
mkdir -p "${MOUNT_DIR}/boot/grub"
cat > "${MOUNT_DIR}/boot/grub/grub.cfg" <<EOF
set default=0
set timeout=5

menuentry "Linux" {
    linux /boot/vmlinuz root=/dev/sda3
    initrd /boot/initrd.img
}
EOF

grub-install \
  --target=i386-pc \
  --boot-directory="${MOUNT_DIR}/boot" \
  --modules="part_msdos ext2 fat normal biosdisk" \
  --force \
  --no-floppy \
  --boot-directory="${MOUNT_DIR}/boot" \
  "$LOOP_DEV"

# Install GRUB (UEFI)
mkdir -p "${MOUNT_DIR}/boot/efi/EFI/boot"
grub-install \
  --target=x86_64-efi \
  --efi-directory="${MOUNT_DIR}/boot/efi" \
  --boot-directory="${MOUNT_DIR}/boot" \
  --removable \
  --no-nvram \
  --modules="part_gpt part_msdos ext2 fat normal efi_gop efi_uga biosdisk" \
  --bootloader-id=GRUB \
  --recheck

# Unmount all
umount -l "${MOUNT_DIR}/boot/efi"
umount -l "${MOUNT_DIR}/boot"
umount -l "${MOUNT_DIR}/recovery"
umount -l "${MOUNT_DIR}"
losetup -d "$LOOP_DEV"

# Create hybrid ISO using xorriso (make sure it's installed)
ISO_FILE="${STAGE_WORK_DIR}/${IMG_NAME}.iso"
EFI_IMG="${STAGE_WORK_DIR}/efi.img"
dd if=/dev/zero of="${EFI_IMG}" bs=1M count=128
mkfs.vfat "${EFI_IMG}"
mmd -i "${EFI_IMG}" ::/EFI ::/EFI/boot
mcopy -i "${EFI_IMG}" "${MOUNT_DIR}/boot/efi/EFI/boot/bootx64.efi" ::/EFI/boot/bootx64.efi
xorriso -as mkisofs \
  -r -J -joliet -l \
  -b boot/grub/i386-pc/eltorito.img \
  -no-emul-boot -boot-load-size 4 -boot-info-table \
  -eltorito-alt-boot \
  -e efi.img \
  -no-emul-boot \
  -isohybrid-gpt-basdat \
  -o "${ISO_FILE}" \
  "${MOUNT_DIR}" \
  --efi-boot-part --efi-boot-image --protective-msdos-label

echo "Bootable ISO with BIOS and UEFI support created at ${ISO_FILE}"