#!/bin/bash -e

# shellcheck disable=SC2119
run_sub_stage() {
	log "Begin ${SUB_STAGE_DIR}"
	pushd "${SUB_STAGE_DIR}" >/dev/null
	for i in {00..99}; do
		for DEBCONF_FILE in "${i}-debconf-${ARCH}" "${i}-debconf-arm-only" "${i}-debconf"; do
			if [ -f "${SUB_STAGE_DIR}/${DEBCONF_FILE}" ]; then
				# Skip arm-only if not arm
				if [[ "${DEBCONF_FILE}" == *"-debconf-arm-only" ]] && [[ "${ARCH}" != "armhf" && "${ARCH}" != "arm64" ]]; then
					continue
				fi

				log "Begin ${SUB_STAGE_DIR}/${DEBCONF_FILE}"

				on_chroot <<EOF
debconf-set-selections <<SELEOF
$(cat "${SUB_STAGE_DIR}/${DEBCONF_FILE}")
SELEOF
EOF

				log "End ${SUB_STAGE_DIR}/${DEBCONF_FILE}"
			fi
		done
		for PACKAGE_LIST in "${i}-packages-nr-${ARCH}" "${i}-packages-nr-arm-only" "${i}-packages-nr"; do
			if [ -f "${SUB_STAGE_DIR}/${PACKAGE_LIST}" ]; then
				# Skip arm-only packages if not arm architecture
				if [[ "${PACKAGE_LIST}" == *"-nr-arm-only" ]] && [[ "${ARCH}" != "armhf" && "${ARCH}" != "arm64" ]]; then
					continue
				fi

				log "Begin ${SUB_STAGE_DIR}/${PACKAGE_LIST}"

				PACKAGES="$(sed -f "${SCRIPT_DIR}/remove-comments.sed" <"${SUB_STAGE_DIR}/${PACKAGE_LIST}")"

				if [ -n "$PACKAGES" ]; then
					on_chroot <<EOF
apt-get -o Acquire::Retries=3 install --no-install-recommends -y $PACKAGES
EOF
				fi

				log "End ${SUB_STAGE_DIR}/${PACKAGE_LIST}"
			fi
		done
		for PACKAGE_LIST in "${i}-packages-${ARCH}" "${i}-packages-arm-only" "${i}-packages"; do
			if [ -f "${SUB_STAGE_DIR}/${PACKAGE_LIST}" ]; then
				# Skip arm-only packages if not arm architecture
				if [[ "${PACKAGE_LIST}" == *"-packages-arm-only" ]] && [[ "${ARCH}" != "armhf" && "${ARCH}" != "arm64" ]]; then
					continue
				fi

				log "Begin ${SUB_STAGE_DIR}/${PACKAGE_LIST}"

				PACKAGES="$(sed -f "${SCRIPT_DIR}/remove-comments.sed" <"${SUB_STAGE_DIR}/${PACKAGE_LIST}")"

				if [ -n "$PACKAGES" ]; then
					on_chroot <<EOF
apt-get -o Acquire::Retries=3 install -y $PACKAGES
EOF
				fi

				log "End ${SUB_STAGE_DIR}/${PACKAGE_LIST}"
			fi
		done
		for PATCH_DIR in "${i}-patches-${ARCH}" "${i}-patches-arm-only" "${i}-patches"; do
			if [ -d "${SUB_STAGE_DIR}/${PATCH_DIR}" ]; then
				# Check arm-only patches should only apply for armhf/arm64
				if [[ "${PATCH_DIR}" == *"-patches-arm-only" ]] && [[ "${ARCH}" != "armhf" && "${ARCH}" != "arm64" ]]; then
					continue
				fi

				log "Begin ${SUB_STAGE_DIR}/${PATCH_DIR}"
				pushd "${STAGE_WORK_DIR}" >/dev/null

				if [ "${CLEAN}" = "1" ]; then
					rm -rf .pc
					rm -rf ./*-pc
				fi

				QUILT_PATCHES="${SUB_STAGE_DIR}/${PATCH_DIR}"
				SUB_STAGE_QUILT_PATCH_DIR="$(basename "$SUB_STAGE_DIR")-pc"
				mkdir -p "$SUB_STAGE_QUILT_PATCH_DIR"
				ln -snf "$SUB_STAGE_QUILT_PATCH_DIR" .pc

				quilt upgrade

				if [ -e "${SUB_STAGE_DIR}/${PATCH_DIR}/EDIT" ]; then
					echo "Dropping into bash to edit patches..."
					bash
				fi

				RC=0
				quilt push -a || RC=$?
				case "$RC" in
				0 | 2) ;; # 0 = success, 2 = "already applied"
				*) false ;;
				esac

				popd >/dev/null
				log "End ${SUB_STAGE_DIR}/${PATCH_DIR}"
			fi
		done
		for RUN_SCRIPT in "${i}-run-${ARCH}.sh" "${i}-run-arm-only.sh" "${i}-run.sh"; do
			if [ -x "${SUB_STAGE_DIR}/${RUN_SCRIPT}" ]; then
				# Skip arm-only if not arm
				if [[ "${RUN_SCRIPT}" == *"-arm-only.sh" ]] && [[ "${ARCH}" != "armhf" && "${ARCH}" != "arm64" ]]; then
					continue
				fi

				log "Begin ${SUB_STAGE_DIR}/${RUN_SCRIPT}"
				"./${RUN_SCRIPT}"
				log "End ${SUB_STAGE_DIR}/${RUN_SCRIPT}"
			fi
		done
		for CHROOT_SCRIPT in "${i}-run-chroot-${ARCH}.sh" "${i}-run-chroot-arm-only.sh" "${i}-run-chroot.sh"; do
			if [ -f "${SUB_STAGE_DIR}/${CHROOT_SCRIPT}" ]; then
				# Skip arm-only if not arm
				if [[ "${CHROOT_SCRIPT}" == *"-arm-only.sh" ]] && [[ "${ARCH}" != "armhf" && "${ARCH}" != "arm64" ]]; then
					continue
				fi

				log "Begin ${SUB_STAGE_DIR}/${CHROOT_SCRIPT}"
				on_chroot <"${SUB_STAGE_DIR}/${CHROOT_SCRIPT}"
				log "End ${SUB_STAGE_DIR}/${CHROOT_SCRIPT}"
			fi
		done
	done
	popd >/dev/null
	log "End ${SUB_STAGE_DIR}"
}

run_stage() {
	log "Begin ${STAGE_DIR}"
	STAGE="$(basename "${STAGE_DIR}")"

	pushd "${STAGE_DIR}" >/dev/null

	STAGE_WORK_DIR="${WORK_DIR}/${STAGE}"
	ROOTFS_DIR="${STAGE_WORK_DIR}"/rootfs

	unmount "${WORK_DIR}/${STAGE}"

	if [ ! -f SKIP_IMAGES ]; then
		if [ -f "${STAGE_DIR}/EXPORT_IMAGE" ]; then
			EXPORT_DIRS="${EXPORT_DIRS} ${STAGE_DIR}"
		fi
	fi
	if [ ! -f SKIP ]; then
		if [ "${CLEAN}" = "1" ]; then
			if [ -d "${ROOTFS_DIR}" ]; then
				rm -rf "${ROOTFS_DIR}"
			fi
		fi
		if [ -x prerun.sh ]; then
			log "Begin ${STAGE_DIR}/prerun.sh"
			./prerun.sh
			log "End ${STAGE_DIR}/prerun.sh"
		fi
		for SUB_STAGE_DIR in "${STAGE_DIR}"/*; do
			if [ -d "${SUB_STAGE_DIR}" ]; then
				if [ ! -f "${SUB_STAGE_DIR}/SKIP" ] && [ ! -f "${SUB_STAGE_DIR}/SKIP_${ARCH}" ]; then
					run_sub_stage
				fi
			fi
		done
	fi

	unmount "${WORK_DIR}/${STAGE}"

	PREV_STAGE="${STAGE}"
	PREV_STAGE_DIR="${STAGE_DIR}"
	PREV_ROOTFS_DIR="${ROOTFS_DIR}"
	popd >/dev/null
	log "End ${STAGE_DIR}"
}

term() {
	if [ "$?" -ne 0 ]; then
		log "Build failed"
	else
		log "Build finished"
	fi
	unmount "${STAGE_WORK_DIR}"
	if [ "$STAGE" = "export-image" ]; then
		for img in "${STAGE_WORK_DIR}/"*.img; do
			unmount_image "$img"
		done
	fi
}

if [ "$(id -u)" != "0" ]; then
	echo "Please run as root" 1>&2
	exit 1
fi

BASE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

if [[ $BASE_DIR = *" "* ]]; then
	echo "There is a space in the base path of pi-gen"
	echo "This is not a valid setup supported by debootstrap."
	echo "Please remove the spaces, or move pi-gen directory to a base path without spaces" 1>&2
	exit 1
fi

export BASE_DIR

if [ -f config ]; then
	# shellcheck disable=SC1091
	source config
fi

while getopts "c:" flag; do
	case "$flag" in
	c)
		EXTRA_CONFIG="$OPTARG"
		# shellcheck disable=SC1090
		source "$EXTRA_CONFIG"
		;;
	*) ;;
	esac
done

export PI_GEN=${PI_GEN:-pi-gen}
export PI_GEN_REPO=${PI_GEN_REPO:-https://github.com/RPi-Distro/pi-gen}
export PI_GEN_RELEASE=${PI_GEN_RELEASE:-Raspberry Pi reference}

export ARCH="${ARCH:-arm64}"
export RELEASE=${RELEASE:-bookworm}
export IMG_NAME="${IMG_NAME:-chillcraftos-$RELEASE-$ARCH}"

export USE_QEMU="${USE_QEMU:-0}"
export IMG_DATE="${IMG_DATE:-"$(date +%Y-%m-%d)"}"
export IMG_FILENAME="${IMG_FILENAME:-"${IMG_DATE}-${IMG_NAME}"}"
export ARCHIVE_FILENAME="${ARCHIVE_FILENAME:-"image_${IMG_DATE}-${IMG_NAME}"}"

export SCRIPT_DIR="${BASE_DIR}/scripts"
export WORK_DIR="${WORK_DIR:-"${BASE_DIR}/work/${IMG_NAME}"}"
export DEPLOY_DIR=${DEPLOY_DIR:-"${BASE_DIR}/deploy"}

export DEPLOY_COMPRESSION=${DEPLOY_COMPRESSION:-xz}
export COMPRESSION_LEVEL=${COMPRESSION_LEVEL:-6}
export LOG_FILE="${WORK_DIR}/build.log"

export TARGET_HOSTNAME=${TARGET_HOSTNAME:-chillcraftos}
export FIRST_USER_NAME=${FIRST_USER_NAME:-system}
export FIRST_USER_ISSYSTEM=${FIRST_USER_ISSYSTEM:-true}
export WPA_COUNTRY
export ENABLE_SSH="${ENABLE_SSH:-0}"
export PUBKEY_ONLY_SSH="${PUBKEY_ONLY_SSH:-0}"

export LOCALE_DEFAULT="${LOCALE_DEFAULT:-en_US.UTF-8}"

export KEYBOARD_KEYMAP="${KEYBOARD_KEYMAP:-us}"
export KEYBOARD_LAYOUT="${KEYBOARD_LAYOUT:-English (US)}"

export TIMEZONE_DEFAULT="${TIMEZONE_DEFAULT:-Americas/Los_Angeles}"

export GIT_HASH=${GIT_HASH:-"$(git rev-parse HEAD)"}

export PUBKEY_SSH_FIRST_USER

export CLEAN
export APT_PROXY
export TEMP_REPO

export STAGE
export STAGE_DIR
export STAGE_WORK_DIR
export PREV_STAGE
export PREV_STAGE_DIR
export ROOTFS_DIR
export PREV_ROOTFS_DIR
export IMG_SUFFIX
export NOOBS_NAME
export NOOBS_DESCRIPTION
export EXPORT_DIR
export EXPORT_ROOTFS_DIR

export QUILT_PATCHES
export QUILT_NO_DIFF_INDEX=1
export QUILT_NO_DIFF_TIMESTAMPS=1
export QUILT_REFRESH_ARGS="-p ab"

# shellcheck source=scripts/common
source "${SCRIPT_DIR}/common"
# shellcheck source=scripts/dependencies_check
source "${SCRIPT_DIR}/dependencies_check"

if [ "$SETFCAP" != "1" ]; then
	export CAPSH_ARG="--drop=cap_setfcap"
fi

mkdir -p "${WORK_DIR}"
trap term EXIT INT TERM

dependencies_check "${BASE_DIR}/depends"

PAGESIZE=$(getconf PAGESIZE)
if [ "$ARCH" == "armhf" ] && [ "$PAGESIZE" != "4096" ]; then
	echo
	echo "ERROR: Building an $ARCH image requires a kernel with a 4k page size (current: $PAGESIZE)"
	echo "On Raspberry Pi OS (64-bit), you can switch to a suitable kernel by adding the following to /boot/firmware/config.txt and rebooting:"
	echo
	echo "kernel=kernel8.img"
	echo "initramfs initramfs8 followkernel"
	echo
	exit 1
fi

echo "Checking native $ARCH executable support..."
if ! arch-test -n "$ARCH"; then
	echo "WARNING: Only a native build environment is supported. Checking emulated support..."
	if ! arch-test "$ARCH"; then
		echo "No fallback mechanism found. Ensure your OS has binfmt_misc support enabled and configured."
		exit 1
	fi
fi

#check username is valid
if [[ ! "$FIRST_USER_NAME" =~ ^[a-z][-a-z0-9_]*$ ]]; then
	echo "Invalid FIRST_USER_NAME: $FIRST_USER_NAME"
	exit 1
fi

if [[ -n "${APT_PROXY}" ]] && ! curl --silent "${APT_PROXY}" >/dev/null; then
	echo "Could not reach APT_PROXY server: ${APT_PROXY}"
	exit 1
fi

if [[ -n "${WPA_PASSWORD}" && ${#WPA_PASSWORD} -lt 8 || ${#WPA_PASSWORD} -gt 63 ]]; then
	echo "WPA_PASSWORD" must be between 8 and 63 characters
	exit 1
fi

if [[ "${PUBKEY_ONLY_SSH}" = "1" && -z "${PUBKEY_SSH_FIRST_USER}" ]]; then
	echo "Must set 'PUBKEY_SSH_FIRST_USER' to a valid SSH public key if using PUBKEY_ONLY_SSH"
	exit 1
fi

log "Begin ${BASE_DIR}"

STAGE_LIST=${STAGE_LIST:-${BASE_DIR}/stage*}
export STAGE_LIST

EXPORT_CONFIG_DIR=$(realpath "${EXPORT_CONFIG_DIR:-"${BASE_DIR}/export-image"}")
if [ ! -d "${EXPORT_CONFIG_DIR}" ]; then
	echo "EXPORT_CONFIG_DIR invalid: ${EXPORT_CONFIG_DIR} does not exist"
	exit 1
fi
export EXPORT_CONFIG_DIR

for STAGE_DIR in $STAGE_LIST; do
	STAGE_DIR=$(realpath "${STAGE_DIR}")
	run_stage
done

CLEAN=1
for EXPORT_DIR in ${EXPORT_DIRS}; do
	STAGE_DIR=${EXPORT_CONFIG_DIR}
	# shellcheck source=/dev/null
	source "${EXPORT_DIR}/EXPORT_IMAGE"
	EXPORT_ROOTFS_DIR=${WORK_DIR}/$(basename "${EXPORT_DIR}")/rootfs
	if [ -e "${EXPORT_DIR}/EXPORT_ISO" ]; then
		# shellcheck source=/dev/null
		source "${EXPORT_DIR}/EXPORT_ISO"
		STAGE_DIR="${BASE_DIR}/export-iso"
		run_stage
	elif [ -e "${EXPORT_DIR}/EXPORT_IMAGE" ]; then
		# shellcheck source=/dev/null
		source "${EXPORT_DIR}/EXPORT_IMAGE"
		STAGE_DIR="${BASE_DIR}/export-image"
		run_stage
	fi
	if [ "${USE_QEMU}" != "1" ]; then
		if [ -e "${EXPORT_DIR}/EXPORT_NOOBS" ]; then
			# shellcheck source=/dev/null
			source "${EXPORT_DIR}/EXPORT_NOOBS"
			STAGE_DIR="${BASE_DIR}/export-noobs"
			run_stage
		fi
	fi
done

if [ -x "${BASE_DIR}/postrun.sh" ]; then
	log "Begin postrun.sh"
	cd "${BASE_DIR}"
	./postrun.sh
	log "End postrun.sh"
fi

log "End ${BASE_DIR}"