#!/bin/bash
set -e


# Check if script is run as root
if [ "$EUID" -ne 0 ]; then
    echo "Please run as root"
    exit 1
fi


IMAGE_NAME="ArchLinuxARM-aarch64-odroid-m1s.img"
MOUNT_POINT="/mnt/arch"


# Create image file
echo "Creating image file..."
fallocate -l 4G "$IMAGE_NAME"


# Set up loop device
LOOP_DEV=$(losetup --find --show "$IMAGE_NAME")
echo "Using loop device: $LOOP_DEV"


# Create partition table and partitions
echo "Creating partitions..."
parted --script "$LOOP_DEV" mklabel msdos
parted --script "$LOOP_DEV" mkpart primary ext2 0% 512M
parted --script "$LOOP_DEV" mkpart primary ext4 512M 100%


# Update partition table
partx -u "$LOOP_DEV"


# Format partitions
echo "Formatting partitions..."
mkfs.ext2 "${LOOP_DEV}p1"
mkfs.ext4 -F "${LOOP_DEV}p2"


# Create mount points and mount partitions
echo "Mounting partitions..."
mkdir -p "$MOUNT_POINT"
mount "${LOOP_DEV}p2" "$MOUNT_POINT"
mkdir -p "$MOUNT_POINT/boot"
mount "${LOOP_DEV}p1" "$MOUNT_POINT/boot"


# Download and extract root filesystem
echo "Downloading and extracting root filesystem..."
wget -c http://os.archlinuxarm.org/os/ArchLinuxARM-aarch64-latest.tar.gz
tar -xpf ArchLinuxARM-aarch64-latest.tar.gz -C "$MOUNT_POINT"


# Mount necessary filesystems for chroot
echo "Mounting virtual filesystems..."
mount -t proc none "$MOUNT_POINT/proc"
mount -t sysfs none "$MOUNT_POINT/sys"
mount -o bind /dev "$MOUNT_POINT/dev"


# Backup and replace resolv.conf
cp "$MOUNT_POINT/etc/resolv.conf" "$MOUNT_POINT/etc/resolv.conf.bak"
cp /etc/resolv.conf "$MOUNT_POINT/etc/resolv.conf"


# Generate initial ramdisk
echo "Generating initial ramdisk..."
mkinitcpio -k "$MOUNT_POINT/boot/Image" -g "$MOUNT_POINT/boot/initramfs-linux.img"


# Chroot operations
echo "Performing chroot operations..."
cat << 'EOF' > "$MOUNT_POINT/setup.sh"
#!/bin/bash
set -e
pacman-key --init
pacman-key --populate
pacman -Suy --noconfirm
pacman -Scc --noconfirm
EOF
chmod +x "$MOUNT_POINT/setup.sh"
chroot "$MOUNT_POINT" /setup.sh


# Update boot files
echo "Updating boot files..."
cd "$MOUNT_POINT/boot"
cp --sparse=never Image Image.new
mv Image.new Image
cp --sparse=never initramfs-linux.img initramfs-linux.img.new
mv initramfs-linux.img.new initramfs-linux.img


# Create boot.cmd for SD card
echo "Creating boot.cmd..."
cat << 'EOF' > "$MOUNT_POINT/boot/boot.cmd"
setenv bootargs "root=/dev/mmcblk1p2 rw rootwait console=tty1 console=ttyS2,1500000"
# Load the DTB
load ${devtype} ${devnum}:${partition} ${fdt_addr_r} /dtbs/rockchip/rk3566-odroid-m1s.dtb
fdt addr ${fdt_addr_r}
# Load kernel and initramfs
load ${devtype} ${devnum}:${partition} ${kernel_addr_r} /Image
load ${devtype} ${devnum}:${partition} ${ramdisk_addr_r} /initramfs-linux.img
setenv initrd_size ${filesize}
# Boot with kernel and initramfs
booti ${kernel_addr_r} ${ramdisk_addr_r}:${initrd_size} ${fdt_addr_r}
EOF


# Generate boot.scr from boot.cmd
echo "Generating boot.scr..."
mkimage -C none -A arm64 -T script -d "$MOUNT_POINT/boot/boot.cmd" "$MOUNT_POINT/boot/boot.scr"


# Cleanup
echo "Cleaning up..."
rm "$MOUNT_POINT/setup.sh"
rm "$MOUNT_POINT/etc/resolv.conf"
mv "$MOUNT_POINT/etc/resolv.conf.bak" "$MOUNT_POINT/etc/resolv.conf"


# Unmount everything
echo "Unmounting filesystems..."
umount "$MOUNT_POINT/dev"
umount "$MOUNT_POINT/proc"
umount "$MOUNT_POINT/sys"
umount "$MOUNT_POINT/boot"
umount "$MOUNT_POINT"


# Detach loop device
echo "Detaching loop device..."
losetup --detach "$LOOP_DEV"


echo "Image preparation complete!"