#!/bin/bash

if [ ! -f .config ]; then
	echo "[x] Configuration file not found!"
	echo "     Please configure build with 'make config' first."
	exit 1
fi

source .config

SUPER_IMAGE="$STOCK_FIRMWARE_PATH/super.img"
FIRMWARE_OUT_DIR="./firmware"
UNPACKED_DIR="./unpacked"
MNT_BASE="$UNPACKED_DIR/mnt"
TMP_DIR="/tmp/driedlamu"
MNT_PARTITIONS=(
	"system_a" 
	"product_a"
	"system_ext_a"
)

if [[ -d $TMP_DIR || -d $UNPACKED_DIR ]]; then
	echo "[x] Build directories detected"
	echo "     Please clean build environment with 'make clean' first."
	exit 1
fi

prepare_erofs_utils() {
	local erofs_utils_repo="https://api.github.com/repos/sekaiacg/erofs-utils/releases/latest"
	local aarch=$(lscpu | sed -n 's/Architecture:[ \t]*//p')
	local erofs_utils_url=$(curl -s "$erofs_utils_repo" | jq --arg aarch_jq $aarch -c '[ .assets[] | select( .browser_download_url | test("Linux_" + $aarch_jq + ".zip$")) ]' | jq -r ".[] | .browser_download_url")
	local download_location="$TMP_DIR/erofs-utils.zip"

	wget -O "$download_location" "$erofs_utils_url"
	unzip "$download_location" -d "./utils/bin/"
}

prepare_apatch_utils() {
	local apatch_repo="https://api.github.com/repos/bmax121/KernelPatch/releases/latest"
	local apatch_urls=$(curl -s "$apatch_repo" | jq -r '[ .assets[] | select( .browser_download_url | test("kptools-linux|kpimg-android") ) ]' | jq -r ".[] | .browser_download_url")

	for l in $apatch_urls; do
		wget -P "./utils/bin/" "$l"
		chmod +x ./utils/bin/kp*
	done
}

simg() {
	local super_image=$1
	local unpacked_dir=$2
	local super_dir=$(dirname "$super_image")
	local super_raw="$super_dir/super_raw.img"

	echo "[i] Converting a sparse image to raw"
	if [ ! -f "$super_image" ]; then
		echo "[x] super.img not found"
		exit 1
	fi

	# converting from sparse
	if ! simg2img "$super_image" "$super_raw"; then
		echo "[x] Failed to convert super image"
		exit 1
	fi

	echo "[i] Unpacking super raw image"
	if ! lpunpack "$super_raw" "$unpacked_dir"; then
		echo "[x] Failed to unpack super image"
		exit 1
	fi

	echo "[i] Super image converting successful"
}

extract_erofs() {
	local DIR=$1 
	local MNT_BASE="$2"

	mkdir -p "${MNT_BASE}/${DIR}"

	./utils/bin/extract.erofs -i "${UNPACKED_DIR}/${DIR}.img" -o "${MNT_BASE}" -x
	echo "[STAGE 2] ${DIR} EROFS partition extracted"
}

make_erofs() {
	local part_name=$1
	local target_dir=$2
	local fs_config_dir="unpacked/mnt/config/"
	local fs_uuid=$(cat ./${fs_config_dir}/${part_name}_fs_options | awk 'NR==2 {print $3}')

	./utils/bin/mkfs.erofs --quiet -zlz4hc -T 0 -U "$fs_uuid" --mount-point="/${part_name}" --fs-config-file="${fs_config_dir}/${part_name}_fs_config" --file-contexts="${fs_config_dir}/${part_name}_file_contexts" "${target_dir}/${part_name}.img" "unpacked/mnt/${part_name}"
	echo "[STAGE 3] ${part_name} EROFS partition repacked"
}

unpack_images() {
	local FS=$1

	if [ ! -d "$MNT_BASE" ]; then
		mkdir -p "$MNT_BASE"
		echo "[+] Directory $MNT_BASE created"
	else
		echo "[i] Directory $MNT_BASE already exists"
	fi

	for dir in "${MNT_PARTITIONS[@]}"; do
		if [ ! -d "${MNT_BASE}/$dir" ]; then
			mkdir -p "${MNT_BASE}/$dir"
			echo "[+] Directory $dir created"
		else
			echo "[i] Directory $dir already exists"
		fi

		if [[ "${FS}" = "erofs" ]]; then
			extract_erofs "${dir}" "${MNT_BASE}"
			mkdir -p "unpacked/original"
			mv "unpacked/${dir}.img" "unpacked/original/${dir}.img"
		fi
	done
}

pack_images() {
	for i in "${MNT_PARTITIONS[@]}"; do
		if [ ! -f "unpacked/${i}.img" ]; then
			make_erofs "$i" "${UNPACKED_DIR}"
			echo "[+] Image 'unpacked/${i}.img' created"
		else
			echo "[?] Image 'unpacked/${i}.img' already exists"
		fi
	done

	echo "[i] All images built"
}

apatch_patch_boot() {
	local apatch_key=$1
	local workdir="$TMP_DIR/root"
	local previous_dir="$(pwd)"

	cd "$workdir"

	"$previous_dir"/utils/bin/kptools-linux unpack boot.img
	"$previous_dir"/utils/bin/kptools-linux -p --image kernel --skey "$apatch_key" --kpimg "$previous_dir/utils/bin/kpimg-android" --out kernel
	"$previous_dir"/utils/bin//kptools-linux repack boot.img

	cd "$previous_dir"
}

vbmeta_add_hashtree() {
	local partition="$1"
	local properties=$(avbtool info_image --image ${UNPACKED_DIR}/original/${partition}.img | grep -E "Prop: com.android.build." | sed -e "s/[[:space:]]//g;s/Prop://g;s/->/:/g;s/^/--prop /g")
	local cmd="avbtool add_hashtree_footer \
		--image '$UNPACKED_DIR/${partition}.img' \
		--partition_name $(sed -e "s/_a//g" <<< $partition) \
		--block_size 4096 \
		--hash_algorithm sha256 \
		${properties[@]} \
		--key '$CUSTOM_AVB_KEY'"

	echo -e "[i] Adding hashtree footer to $partition"

	eval $cmd
}

make_vbmeta() {
	local CUSTOM_AVB_KEY_PATH="${AVB_KEY_PATH/#\~/$HOME}"
	local CUSTOM_AVB_KEY="$CUSTOM_AVB_KEY_PATH/driedlamu_rsa2048_private.pem"
	local CUSTOM_AVB_KEY_PUB="$CUSTOM_AVB_KEY_PATH/driedlamu_rsa2048_public.pem"

	local ORIGINAL_AVB_KEY_PATH="$TMP_DIR/vbmeta/keys/"

	mkdir -p $TMP_DIR \
		$TMP_DIR/vbmeta \
		$TMP_DIR/vbmeta/keys

	if [[ ! -f "$CUSTOM_AVB_KEY" ]]; then
		echo "[i] Generating RSA2048 private key"

		openssl genpkey -algorithm RSA \
			-pkeyopt rsa_keygen_bits:2048 \
			-outform PEM \
			-out "$CUSTOM_AVB_KEY"
	fi

	if [[ ! -f "$CUSTOM_AVB_KEY_PUB" ]]; then
		echo "[i] Extracting RSA2048 public key with AVBTool"

		avbtool extract_public_key \
			--key "$CUSTOM_AVB_KEY" \
			--output "$CUSTOM_AVB_KEY_PUB"
	fi

	# While all keys are the same, It's not known whether Motorola is going to change 
	# them between partition. 
	dd if="$STOCK_FIRMWARE_PATH/vbmeta.img" of="$ORIGINAL_AVB_KEY_PATH/vbmeta_system_key.pem" bs=1 skip="681" count="520"
	dd if="$STOCK_FIRMWARE_PATH/vbmeta.img" of="$ORIGINAL_AVB_KEY_PATH/vbmeta_vendor_key.pem" bs=1 skip="1313" count="520"
	dd if="$STOCK_FIRMWARE_PATH/vbmeta.img" of="$ORIGINAL_AVB_KEY_PATH/vbmeta_unknown_key.pem" bs=1 skip="6080" count="520"

	for i in "${MNT_PARTITIONS[@]}"; do
		vbmeta_add_hashtree "$i"
	done

	echo -e "[i] Making vbmeta_system"

	avbtool make_vbmeta_image \
		--key "$CUSTOM_AVB_KEY" \
		--algorithm 'SHA256_RSA2048' \
		--flags 0 --padding_size '4096' \
		--include_descriptors_from_image $UNPACKED_DIR/system_a.img \
		--output "$TMP_DIR/vbmeta/vbmeta_system.img"

	echo -e "[i] Making vbmeta"

	avbtool make_vbmeta_image --key "$CUSTOM_AVB_KEY" --algorithm 'SHA256_RSA2048' --flags 0 \
		--chain_partition vbmeta_system:2:"$CUSTOM_AVB_KEY_PUB" \
		--chain_partition vbmeta_vendor:4:"$ORIGINAL_AVB_KEY_PATH/vbmeta_vendor_key.pem" \
		--include_descriptors_from_image "$STOCK_FIRMWARE_PATH/boot.img" \
		--include_descriptors_from_image "$STOCK_FIRMWARE_PATH/dtbo.img" \
		--include_descriptors_from_image "$STOCK_FIRMWARE_PATH/init_boot.img" \
		--include_descriptors_from_image "$STOCK_FIRMWARE_PATH/vendor_boot.img" \
		--include_descriptors_from_image "$UNPACKED_DIR/odm_dlkm_a.img" \
		--include_descriptors_from_image "$UNPACKED_DIR/product_a.img" \
		--include_descriptors_from_image "$UNPACKED_DIR/system_dlkm_a.img" \
		--include_descriptors_from_image "$UNPACKED_DIR/system_ext_a.img" \
		--include_descriptors_from_image "$UNPACKED_DIR/vendor_dlkm_a.img" \
		--output "$TMP_DIR/vbmeta/vbmeta.img"

	echo -e "[i] VBMeta created"
}

make_super() {
	local output_path="$1"
	local size=0	
	local params=""

	mkdir -p "$UNPACKED_DIR/sparse"

	for file in `ls -1 $UNPACKED_DIR/*.img`; do
		size=$(($size + $(stat -c %s $file)))
		img2simg "$file" "$UNPACKED_DIR/sparse/$(basename $file)"
		params+="--partition $(basename $file .img):readonly:$(stat -c %s $file):main_$(echo $(basename $file .img) | tail -c 2) " 
		if [[ $(stat -c %s $file) != 0 ]]; then
			params+="--image $(basename $file .img)=${UNPACKED_DIR}/sparse/$(basename $file) "
		fi
	done

	size=$(($size * 2))

	local cmd="lpmake --metadata-size 65536 --metadata-slots 3 --block-size 4096 --virtual-ab --device super:$size --group main_a:$(($size - 2097152)) --group main_b:$(($size - 2097152)) "

	cmd+="$params"
	cmd+="--sparse --output $output_path/super.img"

	eval $cmd
}

echo -e "[i] Prepare"

# mkdir -p $FIRMWARE_OUT_DIR \
# 	$TMP_DIR

mkdir -p $TMP_DIR

if [[ ! -f "./utils/bin/mkfs.erofs" && ! -f "./utils/bin/extract.erofs" ]]; then
	echo -e "[i] Downloading erofs-utils"
	prepare_erofs_utils
fi

if [[ $INSTALL_SU == 1 ]]; then
	if [[ $SU_METHOD == "apatch" ]]; then
		if [[ ! -f "./utils/bin/kpimg-android" && ! -f "./utils/bin/kptools-linux" ]]; then
			echo -e "[i] Downloading APatch tools"
			prepare_apatch_utils
		fi
	fi
fi

echo -e "[i] Stage 1"

simg "${SUPER_IMAGE}" "${UNPACKED_DIR}"

find ./res/* -type d -exec mkdir -p /tmp/driedlamu/{} \;
find res/* -type f -exec ln -sf "$(pwd)/{}" /tmp/driedlamu/{} \;

echo -e "[i] Stage 2\n"
unpack_images "erofs"

echo -e "[i] Running scripts"
echo " "
echo "[-] - Removed"
echo "[+] - Added"
echo "[=] - Modified"
echo " "

./scripts/bloatware/main.sh "$MNT_BASE"
./scripts/props.sh "$MNT_BASE"
./scripts/configs/main.sh "$MNT_BASE" "$TMP_DIR"
./scripts/download_apks.sh "$TMP_DIR"
./scripts/apk_obtain_permissions.sh "$TMP_DIR"

if [[ $PRELOAD_MICROG == 1 ]]; then
	./scripts/microg.sh "$MNT_BASE" "$TMP_DIR"
fi

./scripts/install_apks.sh "$MNT_BASE" "$TMP_DIR/res"
./scripts/permissions/main.sh "$MNT_BASE" "$TMP_DIR"

if [[ $ADD_TONES == 1 ]]; then
	./scripts/add_media.sh "$MNT_BASE" "$TMP_DIR"
fi

echo -e "[i] Stage 3\n"
pack_images

if [[ $INSTALL_SU == 1 ]]; then
	echo -e "[i] Installing superuser"
	mkdir -p "$TMP_DIR/root"

	cp -r "$STOCK_FIRMWARE_PATH/boot.img" "$TMP_DIR/root"

	if [[ $SU_METHOD == "apatch" ]]; then
		echo -e "[i] Patching with APatch"
		apatch_patch_boot "$SU_PASSWORD"
		cp -r "$TMP_DIR/root/new-boot.img" "$UNPACKED_DIR/boot_su.img"
	fi
fi

if [[ $ENABLE_VBMETA == 1 ]]; then
	echo -e "[i] Stage 4 (vbmeta)"

	make_vbmeta
	make_super "."
fi

echo -e "DriedLamu firmware was built!"