#!/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:-$RK_CHIP_DIR/$RK_PARAMETER}"

	if [ ! -r "$PARAMETER" ]; then
		error "$PARAMETER not exists!" >&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_id()
{
	grep -v "^$HDR_PART " "$PART_TABLE" | \
		grep -n -E -m 1 "^$1(|_[a-z]) " | cut -d':' -f1 || true
}

rk_partition_name()
{
	[ "$#" -eq 1 ] || return 1

	# Part name
	if ! echo $1 | grep -qE "^[0-9]*$"; then
		if ! grep -q "^$1 " "$PART_TABLE"; then
			error "No such part ($1)!" >&2
			return 1
		fi

		echo $1
		return 0
	fi

	# Part idx
	IDX=$1
	if [ "$IDX" -lt 1 ]; then
		error "Index should not be less than 1!" >&2
		return 1
	fi

	NUM=$(rk_partition_num)
	if [ "$IDX" -gt "$NUM" ]; then
		error "Index should not be greater than $NUM!" >&2
		return 1
	fi

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

rk_partition_start()
{
	OFFSET=0
	while read NAME SIZE; do
		if echo "$NAME" | grep -qE "^$1(|_[a-z])$"; then
			echo $OFFSET | awk '{printf "0x%08x\n",$1}'
			return 0
		fi

		# NOTE: Assuming partitions are contiguous
		OFFSET=$(( $OFFSET + ${SIZE/-/0} ))
	done < "$PART_TABLE"
}

rk_partition_size()
{
	grep -E -m 1 "^$1(|_[a-z]) " "$PART_TABLE" | cut -d' ' -f2 | \
		tr -d '\-' || true
}

# Get partition size limit, 0 means unlimited or not exists.
rk_partition_size_kb() {
	PART_SIZE="$(rk_partition_size "$1")"
	echo $(( ${PART_SIZE:-0} / 2))
}

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

# Print partition table info
rk_partition_print()
{
	message "\n=========================================="
	message "          Partition table"
	message "=========================================="
	{
		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}')
				message "$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
	message "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:-$RK_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()
{
	TEMP_FILE="$(mktemp)"

	echo "# name size" > "$TEMP_FILE"
	while read NAME SIZE; do
		SIZE_STR=$(rk_partition_size_sector_to_readable $SIZE)
		echo "$NAME $SIZE # $SIZE_STR" >> "$TEMP_FILE"
	done < "$PART_TABLE"

	eval ${EDITOR:-vi} "$TEMP_FILE"

	sed -i -e "/^#/d" -e "s/[ ]*#.*//" "$TEMP_FILE"
	mv "$TEMP_FILE" "$PART_TABLE"

	rk_partition_save
}

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

	PART_NAME="$(rk_partition_name $1)"
	[ "$PART_NAME" ] || 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_name $2 &>/dev/null; then
		error "Part already exists ($2)!"
		return 1
	fi

	PART_NAME="$(rk_partition_name $1)"
	[ "$PART_NAME" ] || 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
	PART_NAME="$(rk_partition_name $2)"
	[ "$PART_NAME" ] || return 1

	PART_NAME="$(rk_partition_name $1)"
	[ "$PART_NAME" ] || return 1

	PART=$(sed -n "/^$PART_NAME /p" "$PART_TABLE")

	NUM=$(rk_partition_num)
	if [ "$2" -eq "$NUM" ] && grep -q "\-$" "$PART_TABLE"; then
		error "Cannot move after unlimited part!"
		return 1
	fi

	if echo "$PART" | grep -q "\-$"; then
		error "Cannot move unlimited part ($1)!"
		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

	PART_NAME="$(rk_partition_name $1)"
	[ "$PART_NAME" ] || 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
		error "Index should not be less than 1!"
		return 1
	fi

	NUM=$(rk_partition_num)
	if [ "$IDX" -gt "$(($NUM + 1))" ]; then
		error "Index should not be greater than $(($NUM + 1))!"
		return 1
	fi

	echo $2 | grep -qE "^[a-zA-Z]" || return 1
	if rk_partition_name $2 &>/dev/null; then
		error "Part already exists ($2)!"
		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
		error "Only latest part can be unlimited!"
		return 1
	fi

	if [ "$IDX" -eq "$(( $NUM + 1 ))" ] && grep -q "\-$" "$PART_TABLE"; then
		error "Cannot insert after unlimited part!"
		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
				error "Only latest part can be unlimited!"
				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()
{
	DEV=$(rk_extra_part_arg ${1:-1} 1)
	case "${DEV:-auto}" in
		/*|*=*) echo $DEV ;;
		auto) echo PARTLABEL=$(rk_extra_part_name $@) ;;
		*) echo PARTLABEL=$DEV ;;
	esac
}

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

rk_extra_part_mountpoint()
{
	MOUNTPOINT="$(rk_extra_part_arg ${1:-1} 3)"
	case "${MOUNTPOINT:-auto}" in
		auto) echo "/$(rk_extra_part_name $@)" ;;
		*) echo "$MOUNTPOINT" ;;
	esac
}

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
		if [ -z "$src" -o "$src" = empty ]; then
			return 0
		elif echo "$src" | grep -q "^/"; then
			echo "$src"
		elif [ -d "$RK_CHIP_DIR/$PART_NAME/$src" ]; then
			echo "$RK_CHIP_DIR/$PART_NAME/$src"
		elif [ -d "$RK_CHIP_DIR/$src" ]; then
			echo "$RK_CHIP_DIR/$src"
		elif [ -d "$RK_EXTRA_PARTS_DIR/$PART_NAME/$src" ]; then
			echo "$RK_EXTRA_PARTS_DIR/$PART_NAME/$src"
		else
			echo "$RK_EXTRA_PARTS_DIR/$src"
		fi
	done
}

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

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

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

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

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

rk_extra_part_img()
{
	echo "$RK_EXTRA_PART_OUTDIR/$(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)"
	DST="$(rk_extra_part_img $1)"
	MOUNT_DIR="$(rk_extra_part_mount_dir $1)"
	FAKEROOT_SCRIPT="$(rk_extra_part_fakeroot_script $1)"
	SRCS="$(rk_extra_part_src $1)"

	notice "Preparing partiton $PART_NAME"

	rm -rf "$OUTDIR" "$DST" "$FAKEROOT_SCRIPT" \
		"$RK_FIRMWARE_DIR/$(basename "$DST")"

	mkdir -p "$OUTDIR"

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

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

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