#!/bin/bash

# This helps to see what is going on
set -e -x

# Used to maintain intermediary state. (Not currently used)
DESTDIR="$(mktemp -d /tmp/imagebuilder-zynqmp.XXXXXX)"
UBOOT_OUT="/tmp/output"
OUTDIR="/tmp/img"
mkdir ${OUTDIR}

IMG="${OUTDIR}/zynqmp.img"

. config

UBOOT_OUT_ABS="$(readlink -f $UBOOT_OUT)"

offset=$((2*1024*1024))
_part1_size=`stat --printf="%s" $UBOOT_OUT/$XEN`
_part1_size=$(( $_part1_size + `stat --printf="%s" $UBOOT_OUT/$DOM0_KERNEL` ))
_part1_size=$(( $_part1_size + `stat --printf="%s" $UBOOT_OUT/$DEVICE_TREE` ))
_part1_size=$(( $_part1_size + `stat --printf="%s" $UBOOT_OUT/$UBOOT_SOURCE` ))
_part1_size=$(( $_part1_size + `stat --printf="%s" $UBOOT_OUT/$UBOOT_SCRIPT` ))
i=0
while test $i -lt $NUM_DOMUS
do
    _part1_size=$(( $_part1_size + `stat --printf="%s" $UBOOT_OUT/${DOMU_KERNEL[$i]}` ))
    if test "${DOMU_RAMDISK[$i]}"
    then
        _part1_size=$(( $_part1_size + `stat --printf="%s" $UBOOT_OUT/${DOMU_RAMDISK[$i]}` ))
    fi
    if test "${DOMU_PASSTHROUGH_DTB[$i]}"
    then
        _part1_size=$(( $_part1_size + `stat --printf="%s" $UBOOT_OUT/${DOMU_PASSTHROUGH_DTB[$i]}` ))
    fi
    i=$(( $i + 1 ))
done

# add 16 MB slack
_part1_size=$(( $_part1_size + 16777216 ))
_part1_size=$(( $_part1_size + $offset - 1))
_part1_size=$(( $_part1_size & ~($offset - 1) ))
_part1_size=$(( $_part1_size / 1024 / 1024 ))
echo PART1 size: "$_part1_size"MB

if test "$DOM0_RAMDISK"
then
    _part2_size=`stat --printf="%s" $UBOOT_OUT/$DOM0_RAMDISK`
fi
_part2_size=$(( $_part2_size + $offset - 1))
_part2_size=$(( $_part2_size & ~($offset - 1) ))
# account for gzip compression
_part2_size=$(( $_part2_size * 2 ))
# add some slack
_part2_size=$(( 128*1024*1024 + $_part2_size ))
_part2_size=$(( $_part2_size / 1024 / 1024 ))
echo PART2 size: "$_part2_size"MB


_sector_size=512
# _sector_first_start needs to be aligned to 2048
_sector_first_start=2048
_sector_first_end=$(( $_part1_size * 1024 * 1024 / $_sector_size + $_sector_first_start - 1))
_sector_second_start=$(( $_sector_first_end + 1 ))
_sector_second_end=$(( $_part2_size * 1024 * 1024 / $_sector_size + $_sector_second_start - 1))
_tot_size=$(( $_part1_size + $_part2_size + 16 ))

# NOTE: Increase vos_a to 256 to accomodate rootfs
# 528 MiB (256 + 256 + 16)
truncate $IMG -s "$_tot_size"M

# create GPT partition table
sgdisk -og $IMG

sgdisk -n 1:$_sector_first_start:$_sector_first_end -c 1:"Linux1" -t 1:8300 $IMG
sgdisk -n 2:$_sector_second_start:$_sector_second_end -c 2:"Linux2" -t 2:8300 $IMG

# find the first available loop device
_loop_dev=$(losetup -f)

# attach loopback device to $IMG
losetup $_loop_dev $IMG

_disksize=$(blockdev --getsize $_loop_dev)

dmsetup create diskimage --table "0 $_disksize linear $_loop_dev 0"

# ensure that /dev/mapper/diskimage exists
while [ ! -b /dev/mapper/diskimage ]
do
    sleep 2
done

kpartx -a /dev/mapper/diskimage

# ensure that /dev/mapper/{diskimage1,diskimage2} exists
while [ ! -b /dev/mapper/diskimage1 ] || [ ! -b /dev/mapper/diskimage2 ]
do
    sleep 2
done

# format vos_a partition
mkfs.ext4 -L vos_a -F /dev/mapper/diskimage1

# format vos_b partition
mkfs.ext4 -L vos_b -F /dev/mapper/diskimage2


# create mount point and mount disk2
mkdir -p ${DESTDIR}/part/disk2
mount -t ext4 /dev/mapper/diskimage2 $DESTDIR/part/disk2

# Extract rootfs cpio archive into `.../part/vos_a`
cd ${DESTDIR}/part/disk2
cat "${UBOOT_OUT_ABS}/${DOM0_RAMDISK}" | gunzip | cpio -id
cd -

# umount
sync
umount $DESTDIR/part/disk2

# create mount point and mount disk1
mkdir -p ${DESTDIR}/part/disk1
mount -t ext4 /dev/mapper/diskimage1 $DESTDIR/part/disk1

# only copy over files that were counted for the partition size
cd "$UBOOT_OUT"
cp --parents "$XEN" "${DESTDIR}/part/disk1/"
cp --parents "$DOM0_KERNEL" "${DESTDIR}/part/disk1/"
cp --parents "$DEVICE_TREE" "${DESTDIR}/part/disk1/"
cp --parents "$UBOOT_SOURCE" "${DESTDIR}/part/disk1/"
cp --parents "$UBOOT_SCRIPT" "${DESTDIR}/part/disk1/"
i=0
while test $i -lt $NUM_DOMUS
do
    cp --parents "${DOMU_KERNEL[$i]}" "${DESTDIR}/part/disk1/"
    if test "${DOMU_RAMDISK[$i]}"
    then
        cp --parents "${DOMU_RAMDISK[$i]}" "${DESTDIR}/part/disk1/"
    fi
    if test "${DOMU_PASSTHROUGH_DTB[$i]}"
    then
        cp --parents "${DOMU_PASSTHROUGH_DTB[$i]}" "${DESTDIR}/part/disk1/"
    fi
    i=$(( $i + 1 ))
done
cd -

# unmount
sync
# This fails for some reason. It could work now because we are not using qemu-user
# fstrim $DESTDIR/part/disk1
umount $DESTDIR/part/disk1

kpartx -d /dev/mapper/diskimage
dmsetup remove diskimage
losetup -d $_loop_dev