#
# Copyright (C) 2025 Amlogic, Inc. All rights reserved.
#

#!/bin/bash

KERNEL=
DTB=
RAMDISK=
OUT_KERNEL="blob-kernel.bin"
OUT_RAMDISK="blob-ramdisk.bin"
FIP_BUILD_FOLDER="fip/_tmp/"

function align_4k() {
	local file=$1
	local raw_size=
	local file_size=
	local rem_val=

	raw_size=`stat -c %s ${file}`
	file_size=`expr ${raw_size} + 4095`
	rem_val=`expr ${file_size} % 4096`
	file_size=`expr ${file_size} - ${rem_val}`
	#echo ==== align file ${file} from ${raw_size} to ${file_size} ====
	dd if=/dev/zero of=${file}_tmp.bin bs=1 count=${file_size} &> /dev/null
	dd if=${file} of=${file}_tmp.bin bs=1 count=${file_size} conv=notrunc &> /dev/null
	mv ${file}_tmp.bin ${file}
}

function amfc_compress() {
	local file=$1

	amfc_zstd_hdr=amfc_zstd_hdr.bin
	./fip/tools/aml_zstd ${file} -19 -o ${file}.tmp
	bin_org_size=`stat -c %s ${file}`
	bin_zstd_size=`stat -c %s ${file}.tmp`
	printf "%s" "@ZST" >  $amfc_zstd_hdr

	printf "%02x%02x%02x%02x" $[(bin_org_size) & 0xff] \
	$[((bin_org_size) >> 8) & 0xff] $[((bin_org_size) >> 16) & 0xff] \
	$[((bin_org_size) >> 24) & 0xff] | xxd -r -ps >>  $amfc_zstd_hdr

	printf "%02x%02x%02x%02x" $[(bin_zstd_size) & 0xff] \
	$[((bin_zstd_size) >> 8) & 0xff] $[((bin_zstd_size) >> 16) & 0xff] \
	$[((bin_zstd_size) >> 24) & 0xff] | xxd -r -ps >>  $amfc_zstd_hdr

	cat $amfc_zstd_hdr ${file}.tmp > ${file}.zstd
	rm $amfc_zstd_hdr ${file}.tmp -f
}

function usage() {
	cat << EOF
	Usage:
		gen-kernel.sh  --help

		create kernel image for c4, input payload will be compressed before create bootblob
		gen-kernel.sh --kernel Image --dtb c4_pxp.dtb --ramdisk rootfs.cpio.uboot --out-kernel blob-kernel.bin --out-ramdisk blob-ramdisk.bin

EOF
	exit 1
}

function parser() {
	local i=0
	local argv=()
	for arg in "$@" ; do
		argv[$i]="$arg"
		i=$((i + 1))
	done
	i=0
	while [ $i -lt $# ]; do
		arg="${argv[$i]}"
		i=$((i + 1)) # must place here
		case "$arg" in
			--h|--help)
				usage
				exit ;;
			--kernel)
				KERNEL="${argv[$i]}"
				i=$((i + 1))
				continue ;;
			--dtb)
				DTB="${argv[$i]}"
				i=$((i + 1))
				continue ;;
			--ramdisk)
				RAMDISK="${argv[$i]}"
				i=$((i + 1))
				continue ;;
			--out-kernel)
				OUT_KERNEL="${argv[$i]}"
				i=$((i + 1))
				continue ;;
			--out-ramdisk)
				OUT_RAMDISK="${argv[$i]}"
				i=$((i + 1))
				continue ;;
			*)
				echo "unsupported args: ${arg}"
				usage
				exit ;;
		esac
	done
}

function create_blob() {
	local dtb_file_size

	# process kernel/dtb
	cp ${KERNEL} ${FIP_BUILD_FOLDER}/Image
	cp ${DTB} ${FIP_BUILD_FOLDER}/dtb
	amfc_compress ${FIP_BUILD_FOLDER}/Image
	amfc_compress ${FIP_BUILD_FOLDER}/dtb
	if [ "${RAMDISK}" != "" ]; then
		cp ${RAMDISK} ${FIP_BUILD_FOLDER}/ramdisk
		amfc_compress ${FIP_BUILD_FOLDER}/ramdisk
	fi

	# align
	dtb_file_size=`stat -c %s ${FIP_BUILD_FOLDER}/dtb.zstd`
	dtb_file_size=`expr ${dtb_file_size} + 16`
	rem_size=`expr ${dtb_file_size} % 16`
	dtb_file_size=`expr ${dtb_file_size} - ${rem_size}`

	dd if=/dev/zero of=${FIP_BUILD_FOLDER}/dtb.tmp bs=${dtb_file_size} count=1 &> /dev/null
	dd if=${FIP_BUILD_FOLDER}/dtb.zstd of=${FIP_BUILD_FOLDER}/dtb.tmp bs=${dtb_file_size} count=1 conv=notrunc &> /dev/null

	rm -rf ${FIP_BUILD_FOLDER}/Image.tmp
	cat ${FIP_BUILD_FOLDER}/dtb.tmp ${FIP_BUILD_FOLDER}/Image.zstd >> ${FIP_BUILD_FOLDER}/Image.tmp
	align_4k ${FIP_BUILD_FOLDER}/Image.tmp

	image_size=`stat -c %s ${FIP_BUILD_FOLDER}/Image.tmp`
	echo final kernel image size:${image_size}
	fip/c4/binary-tool/acpu-imagetool create-boot-blobs --scs-family=c4 --infile-bl33-payload=${FIP_BUILD_FOLDER}/Image.tmp \
		--size-payload-bl33=${image_size} \
		--outfile-blob-bl3-ree=${OUT_KERNEL}
	if [ "${RAMDISK}" != "" ]; then
		cp ${FIP_BUILD_FOLDER}/ramdisk.zstd ${OUT_RAMDISK}
	fi
	# clean up
}

parser $@

if [ -z ${KERNEL} ]; then
	echo "ERROR: NO kernel"
	usage
	exit -1
fi

if [ -z ${DTB} ]; then
	echo "ERROR:NO dtb"
	usage
	exit -1
fi

if [ -z ${RAMDISK} ]; then
	echo "WARN:NO ramdisk"
fi

create_blob ${KERNEL} ${DTB} ${RAMDISK} ${OUT_FILE}