#!/bin/bash
#
# Create Raspberry Pi bootable image from live-build output
# This script creates a proper Pi-bootable image with:
#   - Partition 1: FAT32 boot partition (256MB) with Pi firmware
#   - Partition 2: ext4 root filesystem
#
# Usage: ./create-rpi-image.sh <squashfs_or_chroot_path> <output_image>
#

set -e

INPUT_PATH="$1"
OUTPUT_IMG="$2"
IMG_SIZE="${3:-8G}"

if [ -z "$INPUT_PATH" ] || [ -z "$OUTPUT_IMG" ]; then
    echo "Usage: $0 <squashfs_or_chroot_path> <output_image> [size]"
    echo ""
    echo "Arguments:"
    echo "  squashfs_or_chroot_path  Path to filesystem.squashfs or chroot directory"
    echo "  output_image             Output .img file path"
    echo "  size                     Image size (default: 8G)"
    exit 1
fi

echo "=== Creating Raspberry Pi Image ==="
echo "Input: $INPUT_PATH"
echo "Output: $OUTPUT_IMG"
echo "Size: $IMG_SIZE"

# Create empty image
echo "Creating empty image..."
truncate -s $IMG_SIZE "$OUTPUT_IMG"

# Create partition table (MBR for Pi compatibility)
echo "Creating partition table..."
parted -s "$OUTPUT_IMG" mklabel msdos
parted -s "$OUTPUT_IMG" mkpart primary fat32 1MiB 257MiB
parted -s "$OUTPUT_IMG" mkpart primary ext4 257MiB 100%
parted -s "$OUTPUT_IMG" set 1 boot on

# Setup loop device
echo "Setting up loop device..."
LOOP_DEV=$(losetup --find --show --partscan "$OUTPUT_IMG")
echo "Loop device: $LOOP_DEV"

# Wait for partitions to appear
sleep 2

BOOT_PART="${LOOP_DEV}p1"
ROOT_PART="${LOOP_DEV}p2"

# Verify partitions exist
if [ ! -b "$BOOT_PART" ] || [ ! -b "$ROOT_PART" ]; then
    echo "ERROR: Partitions not found. Trying partx..."
    partx -a "$LOOP_DEV" 2>/dev/null || true
    sleep 2
fi

echo "Boot partition: $BOOT_PART"
echo "Root partition: $ROOT_PART"

# Format partitions
echo "Formatting partitions..."
mkfs.vfat -F 32 -n "boot" "$BOOT_PART"
mkfs.ext4 -L "EcoOS" "$ROOT_PART"

# Create mount points
BOOT_MNT=$(mktemp -d)
ROOT_MNT=$(mktemp -d)

# Mount partitions
echo "Mounting partitions..."
mount "$BOOT_PART" "$BOOT_MNT"
mount "$ROOT_PART" "$ROOT_MNT"

# Extract or copy rootfs
echo "Copying root filesystem..."
if [ -f "$INPUT_PATH" ] && file "$INPUT_PATH" | grep -q "Squashfs"; then
    # It's a squashfs file - extract it
    echo "Extracting squashfs..."
    unsquashfs -f -d "$ROOT_MNT" "$INPUT_PATH"
elif [ -d "$INPUT_PATH" ]; then
    # It's a directory (chroot) - copy it
    echo "Copying chroot directory..."
    cp -a "$INPUT_PATH"/* "$ROOT_MNT"/
else
    echo "ERROR: Input path is neither a squashfs file nor a directory"
    umount "$BOOT_MNT" "$ROOT_MNT"
    losetup -d "$LOOP_DEV"
    rm -rf "$BOOT_MNT" "$ROOT_MNT"
    exit 1
fi

# Copy kernel and initrd to boot partition
echo "Setting up boot partition..."

# Find kernel and initrd
KERNEL=$(ls "$ROOT_MNT"/boot/vmlinuz-* 2>/dev/null | sort -V | tail -1)
INITRD=$(ls "$ROOT_MNT"/boot/initrd.img-* 2>/dev/null | sort -V | tail -1)

if [ -n "$KERNEL" ]; then
    cp "$KERNEL" "$BOOT_MNT/vmlinuz"
    echo "Copied kernel: $(basename $KERNEL)"
fi

if [ -n "$INITRD" ]; then
    cp "$INITRD" "$BOOT_MNT/initrd.img"
    echo "Copied initrd: $(basename $INITRD)"
fi

# Copy device tree blobs if present
if [ -d "$ROOT_MNT/usr/lib/linux-image-"*"-raspi" ]; then
    DTB_DIR=$(ls -d "$ROOT_MNT/usr/lib/linux-image-"*"-raspi" 2>/dev/null | tail -1)
    if [ -d "$DTB_DIR/broadcom" ]; then
        cp -r "$DTB_DIR/broadcom"/*.dtb "$BOOT_MNT/" 2>/dev/null || true
        echo "Copied device tree blobs"
    fi
    if [ -d "$DTB_DIR/overlays" ]; then
        mkdir -p "$BOOT_MNT/overlays"
        cp -r "$DTB_DIR/overlays"/* "$BOOT_MNT/overlays/" 2>/dev/null || true
        echo "Copied device tree overlays"
    fi
fi

# Copy Pi firmware files
echo "Copying Raspberry Pi firmware..."
if [ -d "$ROOT_MNT/usr/lib/raspi-firmware" ]; then
    cp "$ROOT_MNT/usr/lib/raspi-firmware"/*.bin "$BOOT_MNT/" 2>/dev/null || true
    cp "$ROOT_MNT/usr/lib/raspi-firmware"/*.elf "$BOOT_MNT/" 2>/dev/null || true
    cp "$ROOT_MNT/usr/lib/raspi-firmware"/*.dat "$BOOT_MNT/" 2>/dev/null || true
    echo "Copied firmware files from raspi-firmware"
elif [ -d "$ROOT_MNT/boot/firmware" ]; then
    cp "$ROOT_MNT/boot/firmware"/*.bin "$BOOT_MNT/" 2>/dev/null || true
    cp "$ROOT_MNT/boot/firmware"/*.elf "$BOOT_MNT/" 2>/dev/null || true
    cp "$ROOT_MNT/boot/firmware"/*.dat "$BOOT_MNT/" 2>/dev/null || true
    echo "Copied firmware files from /boot/firmware"
fi

# Create config.txt if not present
if [ ! -f "$BOOT_MNT/config.txt" ]; then
    echo "Creating config.txt..."
    cat > "$BOOT_MNT/config.txt" << 'EOF'
# EcoOS Raspberry Pi Configuration
# Supports Pi 3, 4, and 5

# Enable 64-bit mode
arm_64bit=1

# Kernel and initrd
kernel=vmlinuz
initramfs initrd.img followkernel

# Enable serial console for debugging
enable_uart=1

# GPU/display settings
dtoverlay=vc4-kms-v3d
gpu_mem=256

# USB and power settings (Pi 4/5)
max_usb_current=1

# Audio
dtparam=audio=on

# Camera/display interfaces
camera_auto_detect=1
display_auto_detect=1

# Pi 5 specific (ignored on older models)
[pi5]
dtoverlay=dwc2,dr_mode=host
EOF
fi

# Create cmdline.txt if not present
if [ ! -f "$BOOT_MNT/cmdline.txt" ]; then
    echo "Creating cmdline.txt..."
    # Get the UUID of the root partition
    ROOT_UUID=$(blkid -s UUID -o value "$ROOT_PART")
    if [ -n "$ROOT_UUID" ]; then
        echo "console=serial0,115200 console=tty1 root=UUID=$ROOT_UUID rootfstype=ext4 fsck.repair=yes rootwait quiet splash" > "$BOOT_MNT/cmdline.txt"
    else
        echo "console=serial0,115200 console=tty1 root=LABEL=EcoOS rootfstype=ext4 fsck.repair=yes rootwait quiet splash" > "$BOOT_MNT/cmdline.txt"
    fi
fi

# Update fstab in the root filesystem
echo "Updating /etc/fstab..."
BOOT_UUID=$(blkid -s UUID -o value "$BOOT_PART")
ROOT_UUID=$(blkid -s UUID -o value "$ROOT_PART")

cat > "$ROOT_MNT/etc/fstab" << EOF
# EcoOS fstab - Raspberry Pi
# <file system>  <mount point>  <type>  <options>  <dump>  <pass>

# Root filesystem
UUID=$ROOT_UUID  /  ext4  defaults,noatime  0  1

# Boot partition
UUID=$BOOT_UUID  /boot  vfat  defaults  0  2

# Swap (if needed)
# /swapfile  none  swap  sw  0  0
EOF

# Create symlink for boot files in rootfs
mkdir -p "$ROOT_MNT/boot"
echo "Boot partition will be mounted at /boot"

# Set hostname
echo "ecoos-rpi" > "$ROOT_MNT/etc/hostname"

# Cleanup
echo "Cleaning up..."
sync
umount "$BOOT_MNT"
umount "$ROOT_MNT"
losetup -d "$LOOP_DEV"
rm -rf "$BOOT_MNT" "$ROOT_MNT"

# Final size
FINAL_SIZE=$(ls -lh "$OUTPUT_IMG" | awk '{print $5}')
echo ""
echo "=== Raspberry Pi Image Created ==="
echo "Output: $OUTPUT_IMG"
echo "Size: $FINAL_SIZE"
echo ""
echo "To flash to SD card:"
echo "  sudo dd if=$OUTPUT_IMG of=/dev/sdX bs=4M status=progress"
echo ""
echo "Or use Raspberry Pi Imager for a safer flash."