#!/bin/bash 
apppath=$(readlink -f $0)
appdir=$(dirname ${apppath})
app=$(basename $apppath)

DEV_TYPE=Agilex5
IMAGE_FILE=sd_card.img
SD_CARD_SIZE_MB=1024
FAT_SIZE_MB=256
GZIP=N
ZIP=N
FAT_FILES=""
BOOT_FILES=""
ROOT_FS_OVERLAY=""
FAT_OVERLAY=""
USE_SPARSE="no"
ROOT_BALL=""

#
# function die
# print an error message and exit with status
# args 1- are the message to print
die()
{
       shift
       echo $(hostname): FATAL: "$*" >&2
       exit 1
}
usage()
{
	echo Usage:
	echo "$app [-s|--size <size>] [-o|--outfile <outfile>]" 
	echo "     [-a|--sparse]"
	echo "     [-f|--fatfile <file>] [-r|--rfsoverlay <dir>] [-t|--fatoverlay <dir>]"
	echo "     [-v img] [-g] [-d|--device <Agilex5>] [-F|--fatsize <size>] [<root.tar.gz>]"
	echo "    -d device type [Agilex5]"
	echo "    -s size of image [default $SD_CARD_SIZE_MB MB]"
	echo "    -o name of the output file [default is $IMAGE_FILE]"
	echo "    -a create sparse image file (for bmaptool writing)"
	echo "    -f Files to add to the FAT partition (rbf images)"
	echo "    -r File structure overlay to be put on top of root fs"
	echo "    -t FAT structure overlay to be put on top of FAT partition"
	echo "    -v verity rootfs"
	echo "    -g will result in the output file being compressed with gzip"
	echo "    -z will result in the output file being compressed with zip"
	echo "    -F size of the FAT partition in MB"
	echo "    <root.tar.gz> is the yocto filesystem tarball to use"
}

guestfish=$(which guestfish)
[ -n "$guestfish" ] || die guestfish missing.. run sudo apt-get install libguestfs-tools

# Note that we use `"$@"' to let each command-line parameter expand to a 
# separate word. The quotes around `$@' are essential!
# We need TEMP as the `eval set --' would nuke the return value of getopt.
TEMP=`getopt -o o:s:d:f:b:r:t:F:ghv:az --long outfile:,size:,device:,fatfile:,bootfile:,rfsoverlay:,fatoverlay:,fatsize:,gzip,zip,help,verbose,sparse\
     -n $app -- "$@"`

[ $? != 0 ] && die 2 "Terminating... getopt failed"
eval set -- "$TEMP"
while true ; do
        case "$1" in
	-a|--sparse) USE_SPARSE="yes"; shift ;;
        -o|--outfile) IMAGE_FILE="$2"; shift 2;;
        -s|--size) SD_CARD_SIZE_MB="$2"; shift 2;;
        -g|--gzip) GZIP=Y; shift ;;
        -z|--zip) ZIP=Y; shift ;;
        -k|--keep) keepfiles=Y; shift ;;
        -v|--verity) vimg="$2"; shift 2;;
        -h|--help) usage ; exit 0; shift;;
	-d|--device) DEV_TYPE="$2"; shift 2;;
	-f|--fatfile) FAT_FILES="${FAT_FILES} $2"; shift 2;;
	-b|--bootfile) BOOT_FILES="${BOOT_FILES} $2"; shift 2;;
	-r|--rfsoverlay) ROOT_FS_OVERLAY="$2"; shift 2;;
	-t|--fatoverlay) FAT_OVERLAY="$2"; shift 2;;
	-F|--fatsize) FAT_SIZE_MB="$2"; shift 2;;
        --) shift ; break ;;
        *) die 2 "Internal error!: $*" ; shift;;
        esac
done


if [ "$#" -eq 1 ] ; then
	ROOT_BALL=$(readlink -f $1)
fi

rm -rf $IMAGE_FILE
if [ $USE_SPARSE == "yes" ]
then
	# generate image file using a sparse file
	truncate -s${SD_CARD_SIZE_MB}M $IMAGE_FILE || die unable to create SD image file
else
	# Create a blank image 
	dd if=/dev/zero of=$IMAGE_FILE bs=1M count=$SD_CARD_SIZE_MB || die unable to create SD image file
fi

#######################################
###           PARTITIONING          ###
#######################################
fat_start=8192
fat_sectors=$(($FAT_SIZE_MB * 2 * 1024))
fat_end=$(($fat_start + $fat_sectors))
root_start=$(($fat_end + 1))
root_end=$((($SD_CARD_SIZE_MB-1) * 2 * 1024))

echo "Creating partitions"
# Create partitions
# /dev/sda1 fat for fpga images
# /dev/sda2 ext3 for root fs
guestfish --rw --format=raw -a ${IMAGE_FILE} << EOF
	run
	part-init /dev/sda mbr 
	part-add /dev/sda p ${fat_start} ${fat_end}
	part-set-mbr-id /dev/sda 1 0xb
	mkfs vfat /dev/sda1 
	part-add /dev/sda p ${root_start} ${root_end}
	part-set-mbr-id /dev/sda 2 0x83
EOF

if [ -n "$vimg" ]; then
	guestfish --rw -a "$IMAGE_FILE" --format=raw -a "$vimg" -- run : copy-device-to-device /dev/sdb /dev/sda2
fi

#############################################
###        COPYING ROOTFS TARBALL         ###
#############################################
if [ "${ROOT_BALL}" != "" ]
then
echo "Copying in rootfs"
# Copy in rootfs tarball to /dev/sda3
guestfish -a ${IMAGE_FILE} << EOF
	run
	mount /dev/sda2 /
	tar-in  ${ROOT_BALL} / compress:gzip
	chown 0 0 /
	umount /
	sync
EOF
fi

#############################################
###           COPYING FAT FILES           ###
#############################################
if [ "${FAT_FILES}" != "" ]
then
echo "Copying in FAT files"
guestfish -a ${IMAGE_FILE} << EOF
	run
	mount /dev/sda1 /
	copy-in ${FAT_FILES} /
	umount /
	sync
EOF
fi

#############################################
###           COPYING BOOT FILES          ###
#############################################
if [ "${BOOT_FILES}" != "" ]
then
echo "Copying in boot files"
guestfish -a ${IMAGE_FILE} << EOF
	run
	mount /dev/sda2 /
	copy-in ${BOOT_FILES} /boot
	umount /
	sync
EOF
fi

#############################################
###           COPYING ROOT FS Overlay     ###
#############################################
if [ ! -z "$ROOT_FS_OVERLAY" ] 
then
echo "Copying in root fs overlay"
tar -czf root_fs_overlay.tar.gz -C ${ROOT_FS_OVERLAY} .
guestfish -a ${IMAGE_FILE} << EOF
	run
	mount /dev/sda2 /
	tar-in root_fs_overlay.tar.gz / compress:gzip
	chown 0 0 /
	umount /
	sync
EOF
rm root_fs_overlay.tar.gz
fi

#############################################
###           COPYING FAT Overlay         ###
#############################################
if [ ! -z "$FAT_OVERLAY" ] 
then
echo "Copying in fat overlay"
tar -czf fat_overlay.tar.gz -C ${FAT_OVERLAY} .
guestfish -a ${IMAGE_FILE} << EOF
	run
	mount /dev/sda1 /
	tar-in fat_overlay.tar.gz / compress:gzip
	chown 0 0 /
	umount /
	sync
EOF
rm fat_overlay.tar.gz
fi


#############################################
###           COMPRESSING IMAGE           ###
#############################################
if [ "Y" == "$GZIP" ] ; then
	echo compressing ${IMAGE_FILE}
	gzip -f ${IMAGE_FILE}
fi

if [ "Y" == "$ZIP" ] ; then
	echo compressing ${IMAGE_FILE}
	zip ${IMAGE_FILE}.zip ${IMAGE_FILE}
fi

exit 0