#!/usr/bin/env bash

if [ "$#" -ne 6 ]; then
	echo "Usage: $0 OUTFILE BOOT_MEGABYTES TOTAL_MEGABYTES BOOT_DIR ROOTFS_DIR ROOT_DEV"
	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

	# Wait a bit, weird race condition
	sleep 2
}

outfile=$1
boot_mb=$2
root_mb=$3
boot_dir=$4
rootfs_dir=$5
root_dev=$6

# needed + 40mb for first boot and resizing to work, just for logs and such
needed_mb=$(($(du -s --block-size=1MB ${rootfs_dir} | awk {'print $1'}) + 40))

if [ ${root_mb} -lt $needed_mb ];
then
	if [ ${root_mb} -eq -1 ];
	then
		echo "Auto sizing root partition..."
	else
		echo "Requested root partition (${root_mb} MB) is too small, auto resizing..."
	fi
	root_mb=$needed_mb
fi

echo "Building image, p1 ${boot_mb} MB, p2 ${root_mb} MB"

# Create disk image
dd if=/dev/zero of=$outfile bs=1024 count=$(((${boot_mb} + ${root_mb} + 1) * 1024))
sync

# Partition the disk
fdisk $outfile <<EOF
n
p
1

+${boot_mb}M
t
b
a
n
p
2


p
w
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 -n BOOT -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 -O $(./fsckoptlist $rootfs_dir/etc/mke2fs.conf) -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_dir/' --exclude '/tmp/' --exclude '/media/' $rootfs_dir/* mnt/
mkdir -p mnt/boot
mkdir -p mnt/media
mkdir -p mnt/root
mkdir -p mnt/tmp
if [ -x "./createimg-rootfs-post" ] ; then
	./createimg-rootfs-post mnt/ $boot_dir $boot_uuid $root_uuid
fi
umount_retry mnt
losetup_delete_retry $lodev

# 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/
# Copy links as FAT32 can not support links
rsync --quiet --archive --devices --specials --hard-links --acls --xattrs --sparse --copy-links $rootfs_dir/$boot_dir/* mnt/
if [ -x "./createimg-bootfs-post" ] ; then
	./createimg-bootfs-post mnt/ $boot_uuid $root_dev
fi
umount_retry mnt
losetup_delete_retry $lodev