#!/bin/bash

# Make sure that we are sourced and called inside of RK build scripts.
if [ "$BASH_SOURCE" = "$0" -o -z "$RK_SESSION" ];then
	echo "$(realpath $0) is not supposed to be executed directly"
	exit 1
fi

# Parsed partition table file
PART_TABLE="$RK_LOG_DIR/part-table"

# Hidden header partition (part-table and idblock)
HDR_PART="<hidden>"

rk_partition_size_sector_to_readable()
{
	# Consider 0 as unlimited
	case "${1:-grow}" in
		- | 0 | grow)
			echo grow
			return 0 ;;
	esac

	SIZE=$(( $1 * 512 ))
	if [ "$SIZE" -lt 1024 ]; then
		echo $SIZE
	elif [ "$SIZE" -ge $(( 1024 * 1024 * 1024 )) ]; then
		echo "$(echo "scale=1; $SIZE / 1024 / 1024 / 1024" | bc | \
			sed 's/\.0$//')G"
	elif [ "$SIZE" -ge $(( 1024 * 1024 )) ]; then
		echo "$(echo "scale=1; $SIZE / 1024 / 1024" | bc | \
			sed 's/\.0$//')M"
	else
		echo "$(echo "scale=1; $SIZE / 1024" | bc | \
			sed 's/\.0$//')K"
	fi
}

rk_partition_size_readable_to_sector()
{
	SIZE=${1%B}

	case "${SIZE:-grow}" in
		# Consider 0 as unlimited
		- | 0 | grow)
			echo '-'
			return 0
			;;
		0x*)
			echo $SIZE
			return 0
			;;
	esac

	{
		case "$SIZE" in
			*K) echo "${SIZE%K} * 2" | bc ;;
			*M) echo "${SIZE%M} * 2 * 1024" | bc ;;
			*G) echo "${SIZE%G} * 2 * 1024 * 1024" | bc ;;
			*) echo "$SIZE / 512" | bc ;;
		esac
	} | cut -d'.' -f1 | awk '{printf "0x%08x",$1}'
}

# Parse parameter to "<name> <size>" pairs
rk_partition_parse()
{
	PARAMETER="${1:-$CHIP_DIR/$RK_PARAMETER}"

	if [ ! -r "$PARAMETER" ]; then
		echo -e "\e[36m$PARAMETER not exists!\e[0m" >&2
		exit 1
	fi

	PARTS="$(grep "^CMDLINE:" "$PARAMETER" | grep -o "0x.*")"

	# NOTE: Assuming partitions are contiguous
	echo "$HDR_PART $(echo $PARTS | awk -F '[@():]' '{print $2}')"
	echo "${PARTS//,/ }" | xargs -n 1 | \
		awk -F '[@():]' '{print $3,$1}'
}

# Parse parameter to "<name>" arrays
rk_partition_parse_names()
{
	rk_partition_parse "$1" | grep -v "^$HDR_PART " | cut -d' ' -f1 | xargs
}

# Cache parsed partition table
rk_partition_init()
{
	rk_partition_parse > "$PART_TABLE"
}

rk_partition_size()
{
	grep "^$1 " "$PART_TABLE" | cut -d' ' -f2 | tr -d '\-' || true
}

rk_partition_num()
{
	echo $(( $(cat "$PART_TABLE" | wc -l) - 1 ))
}

# Print partition table info
rk_partition_print()
{
	echo
	echo "=========================================="
	echo "          Partition table"
	echo "=========================================="
	{
		OFFSET=0
		while read NAME SIZE; do
			OFFSET=$(echo $OFFSET | awk '{printf "0x%08x",$1}')
			SIZE_STR=$(rk_partition_size_sector_to_readable $SIZE)

			if [ "$NAME" != "$HDR_PART" ]; then
				NAME=$(echo $NAME | awk '{printf "%12s",$1}')
				echo -e "$NAME at $OFFSET size=$SIZE($SIZE_STR)"
			fi

			# NOTE: Assuming partitions are contiguous
			OFFSET=$(( $OFFSET + ${SIZE/-/0} ))
		done < "$PART_TABLE"
	} | sed "=" | sed "N;s/\n/: /"

	echo
	echo "Legacy cmdline:"
	rk_partition_to_cmdline
	echo
}

# Convert partition table to Rockchip legacy cmdline format
rk_partition_to_cmdline()
{
	OFFSET=0
	while read NAME SIZE; do
		case "$SIZE" in
			-)
				# Latest grow part
				echo "$NAME $OFFSET" | \
					awk '{printf "-@0x%08x(%s:grow)",$2,$1}'
				break ;;
			*) SIZE=$(rk_partition_size_readable_to_sector $SIZE) ;;
		esac

		# Visible parts
		if [ "$NAME" != "$HDR_PART" ]; then
			echo "$NAME $OFFSET $(( $SIZE ))" | \
				awk '{printf "0x%08x@0x%08x(%s)",$3,$2,$1}'
		fi

		# NOTE: Assuming partitions are contiguous
		OFFSET=$(( $OFFSET + $SIZE ))
	done < "$PART_TABLE" | sed 's/)\([^$]\)/),\1/g'
}

# Save partition table to parameter
rk_partition_save()
{
	PARAMETER="${1:-$CHIP_DIR/$RK_PARAMETER}"
	[ -r "$PARAMETER" ] || return 1

	PARTS=$(rk_partition_to_cmdline)
	[ "$PARTS" ] || return 1

	sed -i "/^CMDLINE:/s/0x.*/$PARTS/" "$PARAMETER"

	# Reflush
	rk_partition_init
}

# Edit raw partition table
rk_partition_edit()
{
	sed -i '1i # name size' "$PART_TABLE"
	eval ${EDITOR:-vi} "$PART_TABLE"
	sed -i '/#/d' "$PART_TABLE"
	rk_partition_save
}

rk_partition_parse_name()
{
	unset PART_NAME

	[ "$#" -eq 1 ] || return 1

	# Part name
	if ! echo $1 | grep -qE "^[0-9]*$"; then
		if ! grep -q "^$1 " "$PART_TABLE"; then
			echo -e "\e[35mNo such part ($1)!\e[0m"
			return 1
		fi

		PART_NAME=$1
		return 0
	fi

	# Part idx
	IDX=$1
	if [ "$IDX" -lt 1 ]; then
		echo -e "\e[35mIndex should not be less than 1!\e[0m"
		return 1
	fi

	NUM=$(rk_partition_num)
	if [ "$IDX" -gt "$NUM" ]; then
		echo -e "\e[35mIndex should not be greater than $NUM!\e[0m"
		return 1
	fi

	PART_NAME=$(sed -n "$(($IDX + 1))s/\(^[^ ]*\) .*/\1/p" "$PART_TABLE")
}

rk_partition_del()
{
	[ "$#" -gt 0 ] || return 1

	rk_partition_parse_name $1 || return 1
	sed -i "/^$PART_NAME /d" "$PART_TABLE"

	rk_partition_save
}

rk_partition_rename()
{
	[ "$#" -gt 1 ] || return 1

	echo $2 | grep -qE "^[a-zA-Z]" || return 1
	if rk_partition_parse_name $2 >/dev/null; then
		echo -e "\e[35mPart already exists ($2)!\e[0m"
		return 1
	fi

	rk_partition_parse_name $1 || return 1
	sed -i "s/^$PART_NAME /$2 /" "$PART_TABLE"

	rk_partition_save
}

rk_partition_move()
{
	[ "$#" -gt 1 ] || return 1

	echo $2 | grep -qE "^[0-9]*$" || return 1
	rk_partition_parse_name $2 || return 1

	rk_partition_parse_name $1 || return 1
	PART=$(sed -n "/^$PART_NAME /p" "$PART_TABLE")

	NUM=$(rk_partition_num)
	if [ "$2" -eq "$NUM" ] && grep -q "\-$" "$PART_TABLE"; then
		echo -e "\e[35mCannot move after unlimited part!\e[0m"
		return 1
	fi

	if echo "$PART" | grep -q "\-$"; then
		echo -e "\e[35mCannot move unlimited part ($1)!\e[0m"
		return 1
	fi

	sed -i "/^$PART$/d" "$PART_TABLE"
	sed -i "$2 a$PART" "$PART_TABLE"

	rk_partition_save
}

rk_partition_resize()
{
	[ "$#" -gt 1 ] || return 1

	case "$2" in
		0x*) SIZE=$2 ;;
		*) SIZE="$(rk_partition_size_readable_to_sector $2)" ;;
	esac

	rk_partition_parse_name $1 || return 1
	sed -i "s/^$PART_NAME .*/$PART_NAME $SIZE/" "$PART_TABLE"

	rk_partition_save
}

rk_partition_insert()
{
	[ "$#" -gt 1 ] || return 1

	echo $1 | grep -qE "^[0-9]*$" || return 1

	IDX=$1
	if [ "$IDX" -lt 1 ]; then
		echo -e "\e[35mIndex should not be less than 1!\e[0m"
		return 1
	fi

	NUM=$(rk_partition_num)
	if [ "$IDX" -gt "$(($NUM + 1))" ]; then
		echo -e "\e[35mIndex should not be greater than $(($NUM + 1))!\e[0m"
		return 1
	fi

	echo $2 | grep -qE "^[a-zA-Z]" || return 1
	if rk_partition_parse_name $2 >/dev/null; then
		echo -e "\e[35mPart already exists ($2)!\e[0m"
		return 1
	fi

	case "${3:-grow}" in
		0x*) SIZE=$3 ;;
		*) SIZE="$(rk_partition_size_readable_to_sector $3)" ;;
	esac

	if [ "$SIZE" = "-" ] && [ "$IDX" -lt "$(( $NUM + 1 ))" ]; then
		echo -e "\e[35mOnly latest part can be unlimited!\e[0m"
		return 1
	fi

	if [ "$IDX" -eq "$(( $NUM + 1 ))" ] && grep -q "\-$" "$PART_TABLE"; then
		echo -e "\e[35mCannot insert after unlimited part!\e[0m"
		return 1
	fi

	sed -i "$IDX a$2 $SIZE" "$PART_TABLE"

	rk_partition_save
}

# Usage: <offset> <name> <size> <name> <size>...
rk_partition_create()
{
	[ "$#" -gt 1 ] || return 1

	{
		echo "$HDR_PART $(echo $(( $1 )) | awk '{printf "0x%08x",$1}')"
		shift

		while [ "$1" ]; do
			NAME=$1
			shift

			SIZE="$(rk_partition_size_readable_to_sector $1)"
			[ -z "$1" ] || shift

			if [ "$1" -a "$SIZE" = "-" ]; then
				echo -e "\e[35mOnly latest part can be unlimited!\e[0m"
				break
			fi

			echo "$NAME $SIZE"
		done
	} > "$PART_TABLE"

	rk_partition_save
}

# -------- extra partition helpers -------- #

rk_extra_part_num()
{
	echo ${RK_EXTRA_PARTITION_NUM:-0}
}

rk_extra_part_cfg()
{
	[ "$RK_EXTRA_PARTITION_STR" ] || return 0

	RK_EXTRA_PARTITION_ARRAY=( $(echo ${RK_EXTRA_PARTITION_STR//@/ } | \
		xargs -n 1 | sort) )
	PART_IDX=$(( ${1:-1} - 1 ))
	echo "${RK_EXTRA_PARTITION_ARRAY[$PART_IDX]}"
}

rk_extra_part_arg()
{
	PART="$(rk_extra_part_cfg ${1:-1})"
	ARG="$(echo "$PART" | cut -d':' -f${2:-1})"
	echo "${ARG:-$3}"
}

rk_extra_part_dev()
{
	rk_extra_part_arg ${1:-1} 1
}

rk_extra_part_name()
{
	rk_extra_part_arg ${1:-1} 2
}

rk_extra_part_mountpoint()
{
	rk_extra_part_arg ${1:-1} 3 "/$(rk_extra_part_name $1)"
}

rk_extra_part_fstype()
{
	rk_extra_part_arg ${1:-1} 4 ext4
}

rk_extra_part_options()
{
	rk_extra_part_arg ${1:-1} 5 defaults
}

rk_extra_part_src()
{
	PART_NAME="$(rk_extra_part_name $1)"
	for src in $(rk_extra_part_arg ${1:-1} 6 | tr ',' ' '); do
		# Src is either relative path to $RK_IMAGE_DIR/<name>/, or \
		# absolute path.
		if echo "$src" | grep -vq "^/"; then
			echo "$RK_IMAGE_DIR/$PART_NAME/${PART_NAME}_${src}"
		else
			echo "$src"
		fi
	done
}

rk_extra_part_size()
{
	rk_extra_part_arg ${1:-1} 7 auto
}

rk_extra_part_fixed()
{
	rk_extra_part_arg ${1:-1} 8 | grep -wq fixed
}

rk_extra_part_builtin()
{
	rk_extra_part_arg ${1:-1} 8 | grep -wq builtin
}

# Symlink to the mountpoint in rootfs dir
rk_extra_part_mount_dir()
{
	echo "$RK_OUTDIR/$(rk_extra_part_name $1)-mount"
}

rk_extra_part_outdir()
{
	echo "$RK_OUTDIR/$(rk_extra_part_name $1)"
}

rk_extra_part_fakeroot_script()
{
	echo "$(rk_extra_part_outdir $1).fs"
}

rk_extra_part_img()
{
	echo "$RK_FIRMWARE_DIR/$(rk_extra_part_name $1).img"
}

# Prepare extra part's files and fakeroot script
rk_extra_part_prepare()
{
	PART_NAME="$(rk_extra_part_name $1)"
	OUTDIR="$(rk_extra_part_outdir $1)"
	MOUNT_DIR="$(rk_extra_part_mount_dir $1)"
	FAKEROOT_SCRIPT="$(rk_extra_part_fakeroot_script $1)"
	SRCS="$(rk_extra_part_src $1)"

	echo "Preparing partiton $PART_NAME"

	rm -rf "$OUTDIR" "$FAKEROOT_SCRIPT"

	if [ ! "$SRCS" ]; then
		echo "Ignoring $PART_NAME for no sources"
		return 0
	fi

	mkdir -p "$OUTDIR"

	if rk_extra_part_fixed $idx; then
		if ! rk_extra_part_builtin $idx; then
			# Skip boot-time resizing by adding a tag file
			touch "$OUTDIR/.fixed"
		fi
	fi

	echo "#!/bin/sh -e" > "$FAKEROOT_SCRIPT"
	chmod a+x "$FAKEROOT_SCRIPT"

	for src in $MOUNT_DIR $SRCS; do
		[ -d "$src" ] || continue
		[ "$(ls "$src/")" ] || continue

		echo "Merging $src into $OUTDIR"
		rsync -a "$src/" "$OUTDIR"
		for f in $(ls "$OUTDIR" | grep "\.fs$" || true); do
			echo "Merging $src/$f into $FAKEROOT_SCRIPT"
			cat "$OUTDIR/$f" >> "$FAKEROOT_SCRIPT"
			rm -f "$OUTDIR/$f"
		done
	done
}