#!/bin/bash

BACKUP_IFS=$IFS
IFS=$(echo -en "\n\b")

GETOPTS="$(which getopt)"
if [[ "$OSTYPE" == "darwin"* ]] ; then READLINK=greadlink; GETOPTS="$(brew list gnu-getopt | grep bin/getopt)"; else READLINK=readlink;fi;
BASE_DIR="$(dirname $($READLINK -f $0))"

USB_VID=8087
USB_PID=b6b6
TIMEOUT_SEC=60

DO_RECOVERY=0
DO_DFU_FLASH=1
EDISON_XML_FILE="${BASE_DIR}/pft-config-edison.xml"
MCG_XML_FILE="${BASE_DIR}/pft-config-mcg_sku.xml"
# Default to Edison
PFT_XML_FILE=${EDISON_XML_FILE}

# Handle Ifwi file for DFU update
IFWI_DFU_EDISON_FILE=${BASE_DIR}/edison_ifwi-dbg-dfu.bin
IFWI_DFU_MCG_FILE=${BASE_DIR}/ifwi_saltbay_pr2-dfu.bin
# Default to Edison SKU
IFWI_DFU_FILE=${IFWI_DFU_EDISON_FILE}

DO_LIST_VARIANTS=0
VAR_DIR="${BASE_DIR}/u-boot-envs"
VARIANT_NAME_DEFAULT="edison-default"
VARIANT_NAME_BLANK="edison-blank"
VARIANT_NAME=$VARIANT_NAME_BLANK

function print-usage {
	cat << EOF
Usage: ${0##*/} [-hl][--help][--recovery] [--keep-data] [-m][-ipdb] [-t uboot-env]
Update all software and restore board to its initial state.
 -h,--help     display this help and exit.
 --recovery    recover the board to DFU mode using a dedicated tool,
               available only on linux and window hosts.
 --keep-data   preserve user data when flashing.


 deprecated commands:
 -m            select MCG SKU soc when flashing IFWI.
 -i            flash IFWI.
 -d            flash U-boot, U-boot Environment, Linux Kernel, Rootfs.
 -b            blanck the device (eq -i -d -t ${VARIANT_NAME_BLANK}).
 -l            list availables U-boot target environments and exit.
 -t uboot-env  specify U-boot target environments to flash (default is ${VARIANT_NAME}).
EOF
	exit -5
}

function flash-command {
	dfu-util -v -d ${USB_VID}:${USB_PID} $@
	if [ $? -ne 0 ] ;
	then
		echo "Flash failed on $@"
		exit -3
	fi
}

function flash-debug {
	echo "DEBUG: lsusb"
	lsusb
	echo "DEBUG: dfu-util -l"
	dfu-util -l
}

function flash-ifwi {
	phoneflashtool --cli -f "$PFT_XML_FILE"

	if [ $? -ne 0 ];
	then
		echo "Phoneflashtool error"
		flash-debug
		exit -1
	fi
}

function dfu-wait {
	echo "Now waiting for dfu device ${USB_VID}:${USB_PID}"
	echo "Please plug and reboot the board"
	while [ `dfu-util -l | grep Found | grep -c ${USB_VID}` -eq 0 ] \
		&& [ $TIMEOUT_SEC -gt 0 ] && [ $(( TIMEOUT_SEC-- )) ];
	do
		sleep 1
	done

	if [ $TIMEOUT_SEC -eq 0 ];
	then
		echo "Timed out while waiting for dfu device ${USB_VID}:${USB_PID}"
		flash-debug
		echo "Did you plug and reboot your board?"
		echo "If yes, please try a recovery by calling this script with the --recovery option"
		exit -2
	fi
}

# Execute old getopt to have long options support
ARGS=$($GETOPTS -o hlt:meidb -l "keep-data,recovery,help" -n "${0##*/}" -- "$@");
#Bad arguments
if [ $? -ne 0 ]; then print-usage ; fi;
eval set -- "$ARGS";

while true; do
	case "$1" in
		-l) shift; DO_LIST_VARIANTS=1;;
		-t) shift; if [ -n "$1" ]; then VARIANT_NAME=$1; shift; fi;;
		-h|--help) shift; print-usage;;
		-m) shift; PFT_XML_FILE="${MCG_XML_FILE}";IFWI_DFU_FILE="${IFWI_DFU_MCG_FILE}";;
		-i) shift; DO_RECOVERY=1;;
		-d) shift; DO_DFU_FLASH=1;;
		-b) shift; DO_RECOVERY=1; DO_DFU_FLASH=1; VARIANT_NAME=$VARIANT_NAME_BLANK;;
		--recovery) shift; DO_RECOVERY=1;DO_DFU_FLASH=0;;
		--keep-data) shift; DO_RECOVERY=0; VARIANT_NAME=$VARIANT_NAME_DEFAULT;;
		--) shift; break;;
	esac
done

if [ ${DO_LIST_VARIANTS} -eq 1 ];
then
	echo "Availables U-boot targets:"
	for variant in $(ls ${VAR_DIR}) ; do
		#echo just filname without extension
		echo "${variant%.*}"
	done
	exit -5
fi



if [ ${DO_RECOVERY} -eq 1 ];
then
	if [[ "$OSTYPE" == "darwin"* ]] ; then
		echo "Recovery mode is only available on windows and linux";
		exit -3
	fi

	echo "Starting Recovery mode"
	echo "Please plug and reboot the board"
	if [ ! -f "${PFT_XML_FILE}" ];
	then
		echo "${PFT_XML_FILE} does not exist"
		exit -3
	fi
	echo "Flashing IFWI"
	flash-ifwi
	echo "Recovery Success..."
	echo "You can now try a regular flash"
fi

if [ ${DO_DFU_FLASH} -eq 1 ];
then
	echo "Using U-Boot target: ${VARIANT_NAME}"
	VARIANT_FILE="${VAR_DIR}/${VARIANT_NAME}.bin"
	if [ ! -f "${VARIANT_FILE}" ]; then
		echo "U-boot target ${VARIANT_NAME}: ${VARIANT_FILE} not found aborting"
		exit -5
	fi

	dfu-wait

	echo "Flashing IFWI"
	flash-command --alt ifwi -D "${IFWI_DFU_FILE}"
	flash-command --alt ifwib -D "${IFWI_DFU_FILE}"

	echo "Flashing U-Boot"
	flash-command --alt u-boot0 -D "${BASE_DIR}/u-boot-edison.bin"

	echo "Flashing U-Boot Environment"
	flash-command --alt u-boot-env0 -D "${VARIANT_FILE}"

	echo "Flashing U-Boot Environment Backup"
	flash-command --alt u-boot-env1 -D "${VARIANT_FILE}"

	echo "Flashing boot partition (kernel)"
	flash-command --alt boot -D "${BASE_DIR}/edison-image-edison.hddimg"

	echo "Flashing rootfs and rebooting"
	flash-command --alt rootfs -D "${BASE_DIR}/edison-image-edison.ext4" -R

	echo "U-boot & Kernel System Flash Success..."
	if [ $VARIANT_NAME == $VARIANT_NAME_BLANK ] ; then
		echo "Your board needs to reboot twice to complete the flashing procedure, please do not unplug it for 2 minutes."
	fi
fi

if [[ ${DO_DFU_FLASH} -eq 0 &&  ${DO_RECOVERY} -eq 0 ]];
then
	echo "Nothing to do ..."
	print-usage
fi

IFS=${BACKUP_IFS}
