#!/bin/bash

[ -z $BUILDROOT ] && BUILDROOT='./buildroot'

check_error() {
	if [ $? -ne 0 ]; then
		echo "failed"
		exit -1
	fi
	echo "OK"
}

do_flash() {
	if [ ! -f $2 ]; then
		echo "fail: cannot find $2"
		exit -1
	fi

	$BUILDROOT/output/host/usr/bin/fastboot flash $1 $2 &> /dev/null
	check_error
}

usage() {
	echo -e "Usage: $0 [OPTIONS]"
	echo -e "\t-a\t\t\tuse images in the buildroot output directory (automatic)"
	echo -e "\t-m <mbr image>\t\tflash the given mbr image"
	echo -e "\t-b <bootloader image>\tflash the given bootloader image"
	echo -e "\t-r <rootfs image>\tflash the given root filesystem image"
	echo -e "\t-k\t\t\tdo not wipe the u-boot environment (keep)"
	echo -e "\t-h\t\t\tprint this help"
}

if [ $(id -u) -ne 0 ]; then
	echo -e "Run this script as root"
	exit -1
fi

if [ ! -f $BUILDROOT/output/host/usr/bin/fastboot ] ||
	[ ! -f $BUILDROOT/output/host/usr/bin/imx_usb ]; then
	echo "Cannot find needed tools. Please verify the buildroot source tree is at \"$BUILDROOT\""
	exit -1
fi


while getopts ":m::b::r::k:a" o; do
	case "${o}" in
		a)
			mbr_img=./mbr
			uboot_img=$BUILDROOT/output/images/u-boot.imx
			rootfs_img=$BUILDROOT/output/images/rootfs.ext4
			break
			;;
		m)
			mbr_img=${OPTARG}
			;;
		b)
			uboot_img=${OPTARG}
			;;
		r)
			rootfs_img=${OPTARG}
			;;
		k)
			KEEP_ENV=1
			;;
	esac
done

if [ ! -z $mbr_img ] && [ -z $uboot_img ]; then
	echo "Add a bootlader image to flash"
	exit -1
fi

if [ ! -z $rootfs_img ] && [ -z $mbr_img ]; then
	echo "WARNING: flashing a root filesystem without flashing the MBR"
fi

if [ -z $mbr_img ] && [ -z $uboot_img ] && [ -z $rootfs_img ]; then
	usage
	exit 0
fi

echo -n "Launching U-Boot... "
$BUILDROOT/output/host/usr/bin/imx_usb -c $BUILDROOT/output/host/usr/etc/imx-loader.d/ \
		uboot_DO_NOT_TOUCH 2> /dev/null | grep "succeeded" &> /dev/null
check_error

if [ ! -z $mbr_img ]; then
	echo -n "Flashing the MBR table... "
	do_flash 0x0 $mbr_img
fi

if [ -z $KEEP_ENV ]; then
	echo -n "Resetting the environment... "
	do_flash 0x600 zeroenv
fi

if [ ! -z $uboot_img ]; then
	echo -n "Flashing the bootloader... "
	do_flash 0x400 $uboot_img
fi

if [ ! -z $rootfs_img ]; then
	echo -n "Flashing the root filesystem... "
	do_flash 0x7d8200 $rootfs_img
fi

echo -e "\nDone."
exit 0