#!/bin/bash
# SPDX-License-Identifier: GPL-2.0
# Copyright (c) 2019 Fuzhou Rockchip Electronics Co., Ltd.

set -e

usage() {
	cat >&2 << USAGE
usage: $0 [-h] --dtb DTB

optional arguments:
  -h, --help            show this help message and exit
  --dtb DTB             the dtb file name
USAGE
}

# Parse command-line arguments
while [ $# -gt 0 ]; do
	case $1 in
		--dtb)
			DTB=$2
			shift 2
			;;
		-h)
			usage
			exit 0
			;;
		--help)
			usage
			exit 0
			;;
		*)
			shift
			;;
        esac
done

srctree=${srctree-"."}
objtree=${objtree-"."}
if [ "${ARCH}" == "" ]; then
	if [ "$($srctree/scripts/config --state CONFIG_ARM)" == "y" ]; then
		ARCH=arm
	else
		ARCH=arm64
	fi
fi

LOGO_PATH=${srctree}/logo.bmp
[ -f ${LOGO_PATH} ] && LOGO=logo.bmp

LOGO_KERNEL_PATH=${srctree}/logo_kernel.bmp
[ -f ${LOGO_KERNEL_PATH} ] && LOGO_KERNEL=logo_kernel.bmp

KERNEL_IMAGE_PATH=${objtree}/arch/${ARCH}/boot/Image
KERNEL_IMAGE_ARG="--kernel ${KERNEL_IMAGE_PATH}"
if [ "${ARCH}" == "arm" ]; then
	DTB_PATH=${objtree}/arch/arm/boot/dts/${DTB}
	ZIMAGE=zImage
else
	DTB_PATH=${objtree}/arch/arm64/boot/dts/rockchip/${DTB}
	ZIMAGE=Image.lz4
fi
KERNEL_ZIMAGE_PATH=${objtree}/arch/${ARCH}/boot/${ZIMAGE}
KERNEL_ZIMAGE_ARG="--kernel ${KERNEL_ZIMAGE_PATH}"
if [ ! -f ${DTB_PATH} ]; then
	echo "No dtb" >&2
	usage
	exit 1
fi

OUT=out
ITB=${BOOT_IMG}
ITS=${OUT}/boot.its
MKIMAGE=${MKIMAGE-"mkimage"}
MKIMAGE_ARG="-E -p 0x800"

make_boot_img()
{
	RAMDISK_IMG_PATH=${objtree}/ramdisk.img
	[ -f ${RAMDISK_IMG_PATH} ] && RAMDISK_IMG=ramdisk.img && RAMDISK_ARG="--ramdisk ${RAMDISK_IMG_PATH}"

	${srctree}/scripts/mkbootimg \
		${KERNEL_IMAGE_ARG} \
		${RAMDISK_ARG} \
		--second resource.img \
		-o boot.img && \
	echo "  Image:  boot.img (with Image ${RAMDISK_IMG} resource.img) is ready";
	${srctree}/scripts/mkbootimg \
		${KERNEL_ZIMAGE_ARG} \
		${RAMDISK_ARG} \
		--second resource.img \
		-o zboot.img && \
	echo "  Image:  zboot.img (with ${ZIMAGE} ${RAMDISK_IMG} resource.img) is ready"
}

repack_boot_img()
{
	${srctree}/scripts/repack-bootimg \
		--boot_img ${BOOT_IMG} --out ${OUT} \
		${KERNEL_IMAGE_ARG} \
		--second resource.img \
		--dtb ${DTB_PATH} \
		-o boot.img &&
	echo "  Image:  boot.img (${BOOT_IMG} + Image) is ready";
	${srctree}/scripts/repack-bootimg \
		--boot_img ${BOOT_IMG} --out ${OUT} \
		${KERNEL_ZIMAGE_ARG} \
		--second resource.img \
		--dtb ${DTB_PATH} \
		-o zboot.img && \
	echo "  Image:  zboot.img (${BOOT_IMG} + ${ZIMAGE}) is ready"
}

check_mkimage()
{
	MKIMAGE=$(type -p ${MKIMAGE} || true)
	if [ -z "${MKIMAGE}" ]; then
		# Doesn't exist
		echo '"mkimage" command not found - U-Boot images will not be built' >&2
		exit 1;
	fi
}

unpack_itb()
{
	rm -rf ${OUT}
	mkdir -p ${OUT}

	for NAME in $(fdtget -l ${ITB} /images)
	do
		# generate image
		NODE="/images/${NAME}"
		OFFS=$(fdtget -ti ${ITB} ${NODE} data-position)
		SIZE=$(fdtget -ti ${ITB} ${NODE} data-size)
		if [ -z ${OFFS} ]; then
			continue;
		fi

		if [ ${SIZE} -ne 0 ]; then
			dd if=${ITB} of=${OUT}/${NAME} bs=${SIZE} count=1 skip=${OFFS} iflag=skip_bytes >/dev/null 2>&1
		else
			touch ${OUT}/${NAME}
		fi
	done

	[ ! -f ${OUT}/kernel ] && echo "FIT ${ITB} no kernel" >&2 && exit 1 || true
}

gen_its()
{
	TMP_ITB=${OUT}/boot.tmp

	# add placeholder
	cp ${ITB} ${TMP_ITB}
	for NAME in $(fdtget -l ${ITB} /images); do
		fdtput -t s ${TMP_ITB} /images/${NAME} data "/INCBIN/(${NAME})"
	done
	dtc -I dtb -O dts ${TMP_ITB} -o ${ITS} >/dev/null 2>&1
	rm -f ${TMP_ITB}

	# fixup placeholder: data = "/INCBIN/(...)"; -> data = /incbin/("...");
	sed -i "s/\"\/INCBIN\/(\(.*\))\"/\/incbin\/(\"\1\")/" ${ITS}

	# remove
	sed -i "/memreserve/d"		${ITS}
	sed -i "/timestamp/d"		${ITS}
	sed -i "/data-size/d"		${ITS}
	sed -i "/data-position/d"	${ITS}
	sed -i "/value/d"		${ITS}
	sed -i "/hashed-strings/d"	${ITS}
	sed -i "/hashed-nodes/d"	${ITS}
	sed -i "/signer-version/d"	${ITS}
	sed -i "/signer-name/d"		${ITS}
}

gen_itb()
{
	[ -f ${OUT}/fdt ] && cp -a ${DTB_PATH} ${OUT}/fdt && FDT=" + ${DTB}"
	[ -f ${OUT}/resource ] && cp -a resource.img ${OUT}/resource && RESOURCE=" + resource.img"
	COMP=$(fdtget ${ITB} /images/kernel compression)
	case "${COMP}" in
		gzip)	EXT=".gz";;
		lz4)	EXT=".lz4";;
		bzip2)	EXT=".bz2";;
		lzma)	EXT=".lzma";;
		lzo)	EXT=".lzo";;
	esac
	cp -a ${KERNEL_IMAGE_PATH}${EXT} ${OUT}/kernel && \
	${MKIMAGE} ${MKIMAGE_ARG} -f ${ITS} boot.img >/dev/null && \
	echo "  Image:  boot.img (FIT ${BOOT_IMG} + Image${EXT}${FDT}${RESOURCE}) is ready";
	if [ "${EXT}" == "" ] && [ -f ${KERNEL_ZIMAGE_PATH} ]; then
		cp -a ${KERNEL_ZIMAGE_PATH} ${OUT}/kernel && \
		${MKIMAGE} ${MKIMAGE_ARG} -f ${ITS} zboot.img >/dev/null && \
		echo "  Image:  zboot.img (FIT ${BOOT_IMG} + zImage${FDT}${RESOURCE}) is ready";
	fi
}

repack_itb()
{
	check_mkimage
	unpack_itb
	gen_its
	gen_itb
}

# Create U-Boot FIT Image use ${BOOT_ITS}
make_fit_boot_img()
{
	ITS=${OUT}/boot.its

	check_mkimage
	mkdir -p ${OUT}
	rm -f ${OUT}/fdt ${OUT}/kernel ${OUT}/resource ${ITS}

	cp -a ${BOOT_ITS} ${ITS}
	cp -a ${DTB_PATH} ${OUT}/fdt
	cp -a ${KERNEL_ZIMAGE_PATH} ${OUT}/kernel
	cp -a resource.img ${OUT}/resource

	if [ "${ARCH}" == "arm64" ]; then
		sed -i -e 's/arch = ""/arch = "arm64"/g' -e 's/compression = ""/compression = "lz4"/' ${ITS}
	else
		sed -i -e 's/arch = ""/arch = "arm"/g' -e 's/compression = ""/compression = "none"/' ${ITS}
	fi
	FIT_DESC=$(${MKIMAGE} ${MKIMAGE_ARG} -f ${ITS} boot.img | grep "FIT description" | sed 's/FIT description: //')
	echo "  Image:  boot.img (${FIT_DESC}) is ready";
}

if [ -x ${srctree}/scripts/bmpconvert ]; then
	if [ -f ${LOGO_PATH} ]; then
		${srctree}/scripts/bmpconvert ${LOGO_PATH};
	fi
	if [ -f ${LOGO_KERNEL_PATH} ]; then
		${srctree}/scripts/bmpconvert ${LOGO_KERNEL_PATH};
	fi
fi

if [ "${srctree}" != "${objtree}" ]; then
	if [ -f ${LOGO_PATH} ]; then
		cp -a ${LOGO_PATH} ${objtree}/;
	fi
	if [ -f ${LOGO_KERNEL_PATH} ]; then
		cp -a ${LOGO_KERNEL_PATH} ${objtree}/;
	fi
fi
scripts/resource_tool ${DTB_PATH} ${LOGO} ${LOGO_KERNEL} >/dev/null
echo "  Image:  resource.img (with ${DTB} ${LOGO} ${LOGO_KERNEL}) is ready"

if [ -f "${BOOT_IMG}" ]; then
	if file -L -p -b ${BOOT_IMG} | grep -q 'Device Tree Blob' ; then
		repack_itb;
	elif [ -x ${srctree}/scripts/repack-bootimg ]; then
		repack_boot_img;
	fi
elif [ -f "${BOOT_ITS}" ]; then
	make_fit_boot_img;
elif [ -x ${srctree}/scripts/mkbootimg ]; then
	make_boot_img;
fi