#! /bin/bash

function get_kernel_arg() {
	grep "$1=" /proc/cmdline | sed -r "s/.*$1=([^ ]+).*/\\1/";
}

BOOTSETUP_CFG=/etc/bootsetup.cfg
[ -e "${BOOTSETUP_CFG}" ] || { echo "\"${BOOTSETUP_CFG}\" not found"; exit 1; }
source "${BOOTSETUP_CFG}"

EBOOT_DEV_IDX=$(get_kernel_arg boot_dev)
EBOOT_PART_IDX=$(get_kernel_arg boot_part)

UBM_LOCAL_ROOT_DEV=${UBM_LOCAL_ROOT_DEV-$(findmnt -nvo source  /)}

# /dev/sda1
if [[ ${UBM_LOCAL_ROOT_DEV} =~ [0-9]$ ]] && [ -b ${UBM_LOCAL_ROOT_DEV::-1} ]; then
	UBM_ROOT_DEV="${UBM_LOCAL_ROOT_DEV::-1}"
	PART_PREFIX="${UBM_ROOT_DEV}"
	PARTITION_SEP=""
# /dev/mmcblk0p1
elif [[ ${UBM_LOCAL_ROOT_DEV} =~ p[0-9]$ ]] && [ -b ${UBM_LOCAL_ROOT_DEV::-2} ]; then
	UBM_ROOT_DEV="${UBM_LOCAL_ROOT_DEV::-2}"
	PART_PREFIX="${UBM_ROOT_DEV}"
	PARTITION_SEP="p"
# /dev/mapper/sda1
elif [[ ${UBM_LOCAL_ROOT_DEV} =~ [0-9]$ ]] && [[ "${UBM_LOCAL_ROOT_DEV}" == "/dev/mapper/"* ]] && \
		[ -b /dev/$(basename "${UBM_LOCAL_ROOT_DEV::-1}") ]; then
	UBM_ROOT_DEV="/dev/$(basename "${UBM_LOCAL_ROOT_DEV::-1}")"
	PART_PREFIX="${UBM_LOCAL_ROOT_DEV::-1}"
	PARTITION_SEP=""
# /dev/mapper/mmcblk0p1
elif [[ ${UBM_LOCAL_ROOT_DEV} =~ p[0-9]$ ]] && [[ "${UBM_LOCAL_ROOT_DEV}" == "/dev/mapper/"* ]] && \
		[ -b /dev/$(basename "${UBM_LOCAL_ROOT_DEV::-2}") ]; then
	UBM_ROOT_DEV="/dev/$(basename "${UBM_LOCAL_ROOT_DEV::-2}")"
	PART_PREFIX="${UBM_LOCAL_ROOT_DEV::-2}"
	PARTITION_SEP="p"
else
	echo "could not determine \$UBM_ROOT_DEV from ${UBM_LOCAL_ROOT_DEV}"
	exit 1
fi


ROOT_IDX="${UBM_LOCAL_ROOT_DEV: -1}"

#if [ "${EBOOT_DEV_IDX}" == "" ]; then
# 	echo "TODO determine MMC boot partition"
# 	exit 1
# elif [ "${EBOOT_DEV_IDX}" == "1" ]; then
#else
	UBM_FLASH_DEVICE="${UBM_ROOT_DEV}"
# fi

if [ -e /proc/device-tree/model ]; then
	DEVICE_MODEL="$(cat /proc/device-tree/model)"
else
	DEVICE_MODEL=""
fi
if [ "${DEVICE_MODEL}" == "Dentsply Sirona ABK iMX8" ]; then
	if [ "${UBM_LOCAL_ROOT_DEV}" == "/dev/mmcblk0gp0p1" ] || [ "${UBM_LOCAL_ROOT_DEV}" == "/dev/mmcblk0p3" ]; then
		echo "booted ABK from eMMC"
		BOOTED_FROM=emmc
		
		UBOOT_PARTS=( "/dev/mmcblk0boot0" "/dev/mmcblk0boot1" )
		EBOOT_PARTS=( "/dev/mmcblk0p1" "/dev/mmcblk0p2" )
	else
		echo "booted ABK from SD"
		BOOTED_FROM=sd
		
		UBOOT_PARTS=( ${UBM_ROOT_DEV} )
		EBOOT_PARTS=( ${PART_PREFIX}${PARTITION_SEP}"1" )
	fi
elif [ "${DEVICE_MODEL}" == "VSK iMX7 Board" ]; then
	if [ "${UBM_LOCAL_ROOT_DEV}" == "/dev/mmcblk1gp0p1" ] || [ "${UBM_LOCAL_ROOT_DEV}" == "/dev/mmcblk1p3" ]; then
		echo "booted VSK from eMMC"
		BOOTED_FROM=emmc
		
		UBOOT_PARTS=( "/dev/mmcblk1boot0" "/dev/mmcblk1boot1" )
		EBOOT_PARTS=( "/dev/mmcblk1p1" "/dev/mmcblk1p2" )
	else
		echo "booted VSK from SD"
		BOOTED_FROM=sd
		
		UBOOT_PARTS=( ${UBM_ROOT_DEV} )
		EBOOT_PARTS=( ${PART_PREFIX}${PARTITION_SEP}"1" )
	fi
else
	echo "booted from SD on unknown device"
	BOOTED_FROM=sd
	
	UBOOT_PARTS=( ${UBM_ROOT_DEV} )
	EBOOT_PARTS=( ${PART_PREFIX}${PARTITION_SEP}"1" )
fi

# echo "Part prefix: ${PART_PREFIX} EBoot: ${EBOOT_DEV_IDX} ${EBOOT_PART_IDX}"


UPDATED=0
for force in 0 1; do 
	[ "${force}" == "1" ] && echo "warning, will force write update now" >&2
	
# 	for bidx in "${BOOT_PARTS[@]}"; do
# 		if [ "${bidx}" == "${EBOOT_PART_IDX}" ]; then
# 			[ "${force}" == "0" ] && continue
# 			
# 			echo "warning, will overwrite our boot partition" >&2
# 		fi
# 		
# 		[ "${bidx}" == "${ROOT_IDX}" ] && { echo "error, trying to mount rootfs as boot ${bidx}" >&2; continue; }
		
# 		bdev="${PART_PREFIX}${PARTITION_SEP}${bidx}"
	for bdev in "${EBOOT_PARTS[@]}"; do
		[ ! -b "${bdev}" ] && { echo "device $bdev not found" >&2; continue; }
		
		echo "trying to mount ${bdev}"
		
		umount /boot/firmware/
		mount "${bdev}" /boot/firmware/
		RES="$?"
		
		if [ "${RES}" != "0" ]; then
			# make sure we do not trigger any safety check while formating
			dd if=/dev/zero of="${bdev}" bs=1M count=2
			
			echo "error, creating new FS on ${bdev}" >&2
			mkfs.fat "${bdev}"
			
			mount "${bdev}" /boot/firmware/ || { echo "error, creating FS on ${bdev} failed: $?"; continue; }
		fi
		
		# TODO remove all?
		rm /boot/firmware/*.fit 2>/dev/null || :
		#cp -v /boot/*.fit /boot/firmware/
		for x in /boot/*.fit; do
			[ -L "$x" ] && continue
			[ ! -e "$x" ] && continue
			
			cp -v "$x" /boot/firmware/
		done
		
		echo "updated ${bdev}"
		
		UPDATED=1
		
		break
	done
	
	[ "${UPDATED}" == "1" ] && break
done

if [ "$UPDATED" != "1" ]; then
	echo "could not update any boot partition" >&2
	exit 1
fi

UBM_FLASH_DD_PARAMS="${UBM_FLASH_DD_PARAMS-bs=1024 seek=33}"
UBM_FLASH_SOURCE_DIR=/boot/

if [ "${UBM_FLASH_BINARY}" == "" ]; then
	UBM_FLASH_BINARY=$(ls -1 "${UBM_FLASH_SOURCE_DIR}/"flash*.bin 2>/dev/null | head -n1)
	if [ "${UBM_FLASH_BINARY}" == "" ]; then
		UBM_FLASH_BINARY=$(ls -1 "${UBM_FLASH_SOURCE_DIR}/"u-boot*.imx 2>/dev/null | head -n1)
	fi
fi

[ ! -e "${UBM_FLASH_BINARY}" ] && { echo "error, no flash binary"; exit 1; }

for uboot_device in "${UBOOT_PARTS[@]}"; do
	UBM_FLASH_DEVICE="${uboot_device}"
	
	[ ! -b "${UBM_FLASH_DEVICE}" ] && { echo "error, flash device \"${UBM_FLASH_DEVICE}\" does not exist"; continue; }
	
	if [ "${BOOTED_FROM}" == "emmc" ] && [ -e /sys/block/$(basename ${uboot_device})/force_ro ]; then
		echo "unlocking boot partition"
		echo 0 > /sys/block/$(basename ${uboot_device})/force_ro
	fi
	
	echo "Flashing \"${UBM_FLASH_BINARY}\" to \"${UBM_FLASH_DEVICE}\" (${UBM_FLASH_DD_PARAMS})..."
	dd if="${UBM_SYSROOT}/${UBM_FLASH_BINARY}" of="${UBM_FLASH_DEVICE}" ${UBM_FLASH_DD_PARAMS} conv=fsync
	
	if [ "${BOOTED_FROM}" == "emmc" ] && [ -e /sys/block/$(basename ${uboot_device})/force_ro ]; then
		echo "locking boot partition"
		echo 1 > /sys/block/$(basename ${uboot_device})/force_ro
	fi
done

sync
