#!/usr/bin/env bash

if [ "$#" -ne 7 ]; then
	echo "Usage: $0 OUTFILE BOOT_MEGABYTES TOTAL_MEGABYTES BOOT_DIR ROOTFS_DIR UBOOT_BIN_DIR RAMFS_FILE"
	exit 1
fi

set -e

# umount_retry DEVICE
function umount_retry {
	sync
	until umount $1
	do
		sleep 1
	done
}

# get_part_info start|size DISK_IMAGE PARTITION_NUMBER
function get_part_info {
	local cmd=$1
	local img_file=$2
	local part_no=$3

	val=`sfdisk -d $img_file | grep "^$img_file$part_no" | grep -oP "(?<=$cmd=)[ \\d+]*"`
	if [ "$?" -eq 0 ]; then
		echo "$(( ${val//[[:blank:]]/} * 512 ))"
	else
		return 1
	fi
}

# losetup_partition DISK_IMAGE PARTITION_NUMBER
function losetup_partition {
	local part_start=$(get_part_info start $1 $2)
	local part_size=$(get_part_info size $1 $2)
	losetup -f --show --offset $part_start --sizelimit $part_size $1
}

function losetup_delete_retry {
	sync
	until losetup -d $1
	do
		sleep 1
	done
}

# mount_partition DISK_IMAGE PARTITION_NUMBER MOUNT_LOCATION
#function mount_partition {
#	local part_start=$(get_part_info start $1 $2)
#	local part_size=$(get_part_info size $1 $2)
#	mount -o loop,offset=$part_start -t auto $1 $3
#}

outfile=$1
boot_mb=$2
total_mb=$3
boot_dir=$4
rootfs_dir=$5
uboot_bin_dir=$6
ramfs_file=$7

# Calculate desired image size in bytes
img_size=$(( $total_mb*1024*1024 ))

# Bytes/Sector
bytes=512
# Sectors/Track
sectors=63
# Heads/Track
heads=255

# Bytes/Cylinder = head * sectors * bytes
bytes_per_cylinder=$(( $heads*$sectors*$bytes ))
sectors_per_cylinder=$(( $heads*$sectors ))

# Calculate number of cylinders (round up)
cylinders=$(( $(($img_size+$bytes_per_cylinder-1)) / $bytes_per_cylinder ))

# Recalculate the size (pad with one sector)
img_size=$(( $cylinders*$bytes_per_cylinder + $bytes ))

# Create disk image
#qemu-img create -f raw $outfile $img_size
#echo -e "\x55\xaa" | dd bs=1 count=2 seek=510 of=$outfile conv=notrunc
dd if=/dev/zero of=$outfile bs=1024 count=$(( $img_size/1024 ))
echo -e "\x55\xaa" | dd bs=1 count=2 seek=510 of=$outfile conv=notrunc
dd if=$uboot_bin_dir/bl1.bin.hardkernel of=$outfile bs=1 count=442 conv=notrunc
dd if=$uboot_bin_dir/bl1.bin.hardkernel of=$outfile bs=512 skip=1 seek=1 conv=notrunc
dd if=$uboot_bin_dir/u-boot.bin of=$outfile bs=512 seek=64 conv=notrunc
sync

# Using --unit M with sfdisk would be too easy... unfortunately it seems like it
# has issues when the first partition starts at an arbitrary offset...
first_part_start_sector=$((1*1024*1024/$bytes))
first_part_end_sector=$(( $(( $(( $(($boot_mb*1024*1024/$bytes)) + $first_part_start_sector + $sectors_per_cylinder - 1 )) / $sectors_per_cylinder )) * $sectors_per_cylinder ))
second_part_start_sector=$first_part_end_sector
second_part_end_sector=$(($cylinders*$sectors_per_cylinder))

# Partition the disk
sfdisk --unit S --in-order --no-reread $outfile  <<EOF
$first_part_start_sector,$(($first_part_end_sector-$first_part_start_sector)),0x0C,*
$second_part_start_sector,$(($second_part_end_sector-$second_part_start_sector)),L
EOF

if mountpoint -q mnt; then
	device=`df -h mnt | grep "^/" | cut -f1 -d' '`
	umount_retry mnt
	if [ "$device" != "" ]; then
		losetup_delete_retry $device
	fi
fi
if [ -d "mnt" ]; then
	rm -rf mnt/
fi


# Mount boot partition, format it and copy files
lodev=$(losetup_partition $outfile 1)
partprobe $lodev
echo "Formatting boot partition on $lodev..."
mkfs.vfat -F 32 $lodev
boot_uuid=`blkid $lodev | sed -n 's/.*UUID=\"\([^\"]*\)\".*/\1/p'`
losetup_delete_retry $lodev

# Mount rootfs partition, format it and copy files
lodev=$(losetup_partition $outfile 2)
partprobe $lodev
echo "Formatting root partition on $lodev..."
mkfs.ext4 -b 4096 -E stride=16384,stripe-width=16384 -m 1 -L root $lodev
#tune2fs -i 0 -c 0 $lodev
root_uuid=`blkid $lodev | sed -n 's/.*UUID=\"\([^\"]*\)\".*/\1/p'`
mkdir -p mnt
mount -t ext4 $lodev mnt/
rsync --quiet --archive --devices --specials --hard-links --acls --xattrs --sparse --exclude '/boot/' --exclude '/tmp/' --exclude '/media/' --exclude '/root/' $rootfs_dir/* mnt/
mkdir -p mnt/boot
mkdir -p mnt/media
mkdir -p mnt/root
mkdir -p mnt/tmp
sed -e "s/\${BOOT_UUID}/$boot_uuid/" -e "s/\${ROOT_UUID}/$root_uuid/" fstab.template > mnt/etc/fstab
if [ -x "./createimg-rootfs-post" ] ; then
	./createimg-rootfs-post mnt/ $rootfs_dir/
fi
umount_retry mnt
losetup_delete_retry $lodev

# Wait a bit, weird race condition (see above)
sleep 2

# Mount the boot partition again and copy boot.ini
lodev=$(losetup_partition $outfile 1)
partprobe $lodev
echo "Mounting boot partition on $lodev..."
mount -t vfat $lodev mnt/
sed -e "s/\${BOOT_UUID}/$boot_uuid/" -e "s/\${ROOT_UUID}/$root_uuid/" boot.ini.template > mnt/boot.ini
cp -rp $boot_dir/* mnt/
cp -p $ramfs_file mnt/
if [ -x "./createimg-bootfs-post" ] ; then
	./createimg-bootfs-post mnt/
fi
umount_retry mnt
losetup_delete_retry $lodev