#! /bin/bash
#
# Source: https://github.com/anyc/uboot-manager
# License: GPLv2
#
# Written by Mario Kicherer (http://kicherer.org)
#

if [ -z ${UBM_CFG_LOADED} ]; then
	if [ -x ubm_common.sh ]; then
		source ubm_common.sh
	elif [ -x /lib32/exec/ubm/ubm_common.sh ]; then
		source /lib32/exec/ubm/ubm_common.sh
	elif [ -x /usr/lib32/exec/ubm/ubm_common.sh ]; then
		source /usr/lib32/exec/ubm/ubm_common.sh
	elif [ -x /lib32/exec/ubm/ubm_common.sh ]; then
		source /lib32/exec/ubm/ubm_common.sh
	elif [ -x /lib64/exec/ubm/ubm_common.sh ]; then
		source /lib64/exec/ubm/ubm_common.sh
	else
		LIBEXEC_PATH="${LIBEXEC_PATH-/lib32/exec}"
		source "${LIBEXEC_PATH}/ubm/ubm_common.sh"
	fi
	
	UBM_CFG_LOADED=1
fi

if [ "${UBM_USE_MODEL_DIRS}" == "1" ]; then
	# first we go through all models and check if there are kernels in the model-specific
	# boot subdirectory and create a model-specific boot_${board_model}.scr. Afterwards, we
	# create the regular boot.scr that loads the model-specific boot script, if possible.
	# If the model-specific boot scripts does not boot a kernel, we fall back to the
	# non-model-specific boot script.
	
	for model in "${UBM_MODELS[@]}"; do
		echo "checking model ${model}"
		UBM_USE_MODEL_DIRS="0" UBM_KERNEL_DIR="${UBM_KERNEL_DIR}/${model}" UBM_SCRIPT_SUFFIX="_${model}" source "${BASH_SOURCE[0]}"
	done
	
	UBM_ENV_FILE=$(mktemp)
	
	echo -n "" > "${UBM_ENV_FILE}"
	echo "if ${UBM_FS_LOAD} ${UBM_FW_PART_ID} \${${UBM_SCRIPT_LADDR_ID}} \"boot_\${board_name}.scr\"; then source \${${UBM_SCRIPT_LADDR_ID}}; else echo \"model script not found\"; fi" >> "${UBM_ENV_FILE}"
	
	# we do not exit here as we look for kernel in the old standard location too
else
	UBM_ENV_FILE=$(mktemp)

	echo -n "" > "${UBM_ENV_FILE}"
fi

# override boot procedure by GPIO state?
if [ "${UBM_FAILSAFE_GPIO_PIN}" != "" ]; then
	# if GPIO pin is high, execute command
	echo "if gpio input ${UBM_FAILSAFE_GPIO_PIN}; then ; else setenv stdout serial,lcd ; echo \"failsafe pin set\"; ${UBM_FAILSAFE_COMMAND}; fi" >> "${UBM_ENV_FILE}"
fi

# override U-Boot output
if [ "${UBM_STDOUT}" != "" ]; then
	echo "setenv stdout ${UBM_STDOUT}" >> "${UBM_ENV_FILE}"
fi

# override device tree?
if [ "${UBM_PRELOADED_DEVTREE}" == "0" ]; then
	if [ "${UBM_DEVTREE_FILE}" != "" ]; then
		echo "${UBM_FS_LOAD} ${UBM_PART_ID} \${fdt_addr} ${UBM_DEVTREE_FILE}" >> "${UBM_ENV_FILE}"
	fi
else
	# check if device tree is available in memory. If yes, get bootargs set by
	# firmware and store them in $fw_bootargs. If not, load the device tree file
	# into memory for safety and backwards compatibility
	echo -n "if fdt addr \${fdt_addr}; then fdt get value fw_bootargs /chosen bootargs; printenv fw_bootargs; " >> "${UBM_ENV_FILE}"
	if [ "${UBM_DEVTREE_FILE}" != "" ]; then
		echo -n "else ${UBM_FS_LOAD} ${UBM_PART_ID} \${fdt_addr} ${UBM_DEVTREE_FILE}; "  >> "${UBM_ENV_FILE}"
	else
		echo -n "echo 'preloaded device tree not found'" >> "${UBM_ENV_FILE}"
	fi
	echo    "fi;" >> "${UBM_ENV_FILE}"
fi

# set default kernel args
if [ "${UBM_PREPEND_FW_KERNEL_ARGS}" == "1" ]; then
	for cb in "${UBM_FILTER_FW_KERNEL_ARGS[@]}"; do
		[ "$cb" == "" ] && continue
		# we use sub in a loop instead of gsub, as the latter causes uboot-2018.07 to freeze
		echo "while setexpr fw_bootargs sub ${cb}; do ; done;" >> "${UBM_ENV_FILE}"
	done
	
	# do not append fw_bootargs as it contains rootfstype=ext4
	echo "setenv bootargs \"\${fw_bootargs} root=${UBM_ROOT_DEV} rw rootwait ${UBM_KERNEL_ARGS}\"" >> "${UBM_ENV_FILE}"
else
	echo "setenv bootargs \"root=${UBM_ROOT_DEV} rw rootwait ${UBM_KERNEL_ARGS}\"" >> "${UBM_ENV_FILE}"
fi

# menu header
echo "echo ' '
echo 'Available boot options:'
echo '-----------------------'" >> "${UBM_ENV_FILE}"


# callback for each kernel
function vmlinuz_cb() {
	dir=$1
	kernel=$2
	version=$3
	
	echo -n "Found kernel \"${kernel}\""
	echo "" >> "${UBM_ENV_FILE}"
	
	#
	# check if there is a matching initramfs
	#
	
	UBM_INITRD="-"
	
	pushd ${BASE_DIR} > /dev/null
	FILES=$(eval find \"${dir}\" -maxdepth 1 ${UBM_INITRD_FILTER} | sort -r)
	[ "${FILES}" == "" ] && echo ""
	
	for fpath in $FILES; do
		# prefer initrd with a matching version suffix
		[[ "${fpath}" != *"-${version}" ]] && [ "${UBM_INITRD}" != "-" ] && continue
		
		UBM_INITRD_FPATH="$fpath"
		UBM_INITRD_NAME="$(basename "${UBM_INITRD_FPATH}")"
		
		if [ "${UBM_CONVERT_INITRD}" == "1" ] && [ "${UBM_INITRD_NAME}" == *".gz."* ]; then
			# TODO untested
			
			ORIG_FPATH="${UBM_INITRD_FPATH}"
			UBM_INITRD_FPATH="$(dirname "${UBM_INITRD_FPATH}")/${UBM_INITRD_NAME/.gz./.ugz.}"
			
			echo "Creating U-Boot image for \"${UBM_INITRD_NAME}\""
			mkimage -n "Initramfs" -A arm -O linux -T ramdisk -C gzip -d \"${ORIG_FPATH}\" \"${UBM_INITRD_FPATH}\" || { echo "error creating ${UBM_INITRD_FPATH}"; exit 1; }
		fi
		
		echo " with initrd \"${UBM_INITRD_NAME}\""
		UBM_INITRD=\${${UBM_INITRD_LADDR_ID}}
		PRECMDS="${PRECMDS}${UBM_FS_LOAD} ${UBM_PART_ID} \${${UBM_INITRD_LADDR_ID}} \"${UBM_INITRD_FPATH}\"; "
	done
	
	if [ -f "${dir}/cmdline-${version}" ]; then
		echo "setenv bootargs${CUR_IDX} 'setenv bootargs \"$(cat "${dir}/cmdline-${version}" | tr '\n' ' ' | sed "s/[[:space:]]*$//")\"'" >> "${UBM_ENV_FILE}"
		PRECMDS="${PRECMDS}run bootargs${CUR_IDX}; "
	fi
	popd >/dev/null
	
	
	echo "setenv boot${CUR_IDX} '${PRECMDS}${UBM_FS_LOAD} ${UBM_PART_ID} \${${UBM_KERNEL_LADDR_ID}} \"${dir}/${kernel}\"; ${UBM_BOOT_CMD} \${${UBM_KERNEL_LADDR_ID}} ${UBM_INITRD} \${fdt_addr}'" >> "${UBM_ENV_FILE}"
	
	echo "echo '   \"run boot${CUR_IDX}\" will boot ${kernel}'" >> "${UBM_ENV_FILE}"
}

# register callback
[[ "${UBM_KERNEL_FCTS}" == *vmlinuz_cb* ]] || UBM_KERNEL_FCTS="${UBM_KERNEL_FCTS} vmlinuz_cb"

# convert filename patterns into commandline arguments for the "find" utility
function join_array { local d="$1"; shift; echo -n "\"$1"; shift; printf "%s" "${@/#/$d}"; }
UBM_KERNEL_FILTER="-name $(join_array "\" -o -name \"" "${UBM_KERNEL_PATTERNS[@]}" )\""
UBM_INITRD_FILTER="-name $(join_array "\" -o -name \"" "${UBM_INITRD_PATTERNS[@]}" )\""

# look for kernel images
function browse_kernel_dir() {
	[ -d "$1/$2/$3" ] || return
	
	BASE_DIR=$1
	SUBVOL=$2
	KERNEL_DIR=$3
	
	echo "browsing ${BASE_DIR}/${SUBVOL}/${KERNEL_DIR}"
	
	pushd "${BASE_DIR}/${SUBVOL}/${KERNEL_DIR}" >/dev/null
		# look for kernels and execute callbacks
		FILES=$(eval find . -maxdepth 1 ${UBM_KERNEL_FILTER} | sort -r)
# 		FILES=$(find . -maxdepth 1 -name "vmlinuz*" -o -name "zImage*" -o -name "Image*" -o -name "xipImage*" -o -name "uImage*" | sort -r)
		for k in $FILES; do
			[ -e "$k" ] || continue
			
			PRECMDS=""
			[ "${SUBVOL}" != "" ] && PRECMDS="${PRECMDS}setenv bootargs \"\${bootargs} rootflags=subvol=${SUBVOL}\"; "
			kernel=$(basename "${k}")
			version=$(basename "${k}" | cut -d "-" -f2-)
			
			# execute all registered callbacks for a kernel
			for cb in ${UBM_KERNEL_FCTS}; do
				${cb} "${SUBVOL}/${KERNEL_DIR}" ${kernel} ${version}
			done
			
			CUR_IDX=$(( ${CUR_IDX} + 1 ))
		done
		echo "echo ' '" >> "${UBM_ENV_FILE}"
	popd >/dev/null
}

CUR_IDX=0

if [ "${UBM_SUBVOL_SEARCH}" != "1" ]; then
	# just browser the given kernel directory
	browse_kernel_dir "${UBM_SYSROOT}/${UBM_KERNEL_DIR}" "" ""
else
	#
	# mount the root volume of btrfs and browse the given kernel directory in each subvolume 
	#
	
	FSTYPE=$(blkid -o value -s TYPE "${UBM_LOCAL_ROOT_DEV}")
	
	[ "${FSTYPE}" == "btrfs" ] || { echo "invalid filesystem ${FSTYPE} for subvol search" >&2; exit 1; }
	
	TMP_DIR="$(mktemp -d)"
	mount "${UBM_LOCAL_ROOT_DEV}" "${TMP_DIR}/" -o subvol=/ || { echo "mount failed" >&2; exit 1; }
	
	pushd "${TMP_DIR}" >/dev/null
		for subvol in ${UBM_BTRFS_VOLUME_PREFIX}*; do
			browse_kernel_dir "${TMP_DIR}" "${subvol}" "${UBM_KERNEL_DIR}"
		done
	popd >/dev/null
	
	umount "${TMP_DIR}"
	rmdir "${TMP_DIR}"
fi

# do we alter boot procedure if failsafe file does not exist?
if [ "${UBM_FAILSAFE_FILE}" != "" ]; then
	# if file exists, execute command, else create file
	echo "if test -e ${UBM_PART_ID} ${UBM_FAILSAFE_FILE}; then setenv stdout serial,lcd ; echo \"failsafe file ${UBM_FAILSAFE_FILE} present\"; ${UBM_FAILSAFE_COMMAND}; else ${UBM_FS_WRITE} ${UBM_PART_ID} 0x02050000 ${UBM_FAILSAFE_FILE} 1; fi" >> "${UBM_ENV_FILE}"
fi

if [ "${UBM_PREBOOT_SCRIPT}" != "" ]; then
        echo "" >> "${UBM_ENV_FILE}"
        echo -n "${UBM_FS_LOAD} ${UBM_PART_ID} \${${UBM_SCRIPT_LADDR_ID}} ${UBM_PREBOOT_SCRIPT} &&" >> "${UBM_ENV_FILE}"
        echo " source \${${UBM_SCRIPT_LADDR_ID}}" >> "${UBM_ENV_FILE}"
fi

echo "
setenv autoboot run distro_bootcmd

echo 'Will execute \"run ${UBM_DEFAULT}\" in ${UBM_MENU_TIMEOUT} seconds, abort with CTRL+c...'
sleep ${UBM_MENU_TIMEOUT} && run ${UBM_DEFAULT}
" >> "${UBM_ENV_FILE}"

# create an image from the script in a format uboot understands
echo "Building uboot-manager script..."
mkdir -p "${UBM_SYSROOT}/${UBM_FW_DIR}"
mkimage -A arm -O linux -T script -C none -a 0x00000000 -e 0x00000000 -n "uboot-manager" -d "${UBM_ENV_FILE}" "${UBM_SYSROOT}/${UBM_FW_DIR}/boot${UBM_SCRIPT_SUFFIX}.scr" || { echo "error creating uboot-manager script image"; exit 1; }

EXIT_CODE=0

if [ "${UBM_FLASH_UBOOT}" == "1" ]; then
	if [ -f "${UBM_SYSROOT}/${UBM_FLASH_SOURCE_DIR}/${UBM_FLASH_BINARY}" ]; then
		echo "Flashing ${UBM_FLASH_SOURCE_DIR}/${UBM_FLASH_BINARY}..."
		dd if="${UBM_SYSROOT}/${UBM_FLASH_SOURCE_DIR}/${UBM_FLASH_BINARY}" of="${UBM_FLASH_DEVICE}" ${UBM_FLASH_DD_PARAMS} conv=fsync
	else
		echo "ERROR: u-boot binary \"${UBM_SYSROOT}/${UBM_FLASH_SOURCE_DIR}/${UBM_FLASH_BINARY}\" not found"
		EXIT_CODE=2
	fi
fi

[ "${UBM_SHOW_ENV}" != "" ] && { echo -e "\n\nResulting script:\n"; cat "${UBM_ENV_FILE}"; }
rm "${UBM_ENV_FILE}"

[ "${EXIT_CODE}" == "0" ] || exit ${EXIT_CODE}
