#!/bin/bash

# log
# Write a status message to the log.
#
# Parameters:
#   $1: The message to write.
log() {
	echo "$1" >&3
}

# dry_run
# Write a debug message to the log. Use instead of a destructive operation when
# DRY_RUN is set.
#
# Parameters:
#   $1: The message to write.
dry_run() {
	if test -n "$DRY_RUN" ; then
		echo "DRY_RUN: ${1}" >&3
	fi
}

# fatal
# Exits the script after writing an error message to the log.
#
# Parameters:
#   $1: A string describing the error encountered.
fatal() {
	echo "FAILURE in ${_EXTRACTED_PATH}/imager.sh: ${1}" | tee -a /etc/motd >&3
	exit 1;
}

# is_at_least_version
# Compares the provided version string to the current FlashCast version.
#
# Parameters:
#   $1: The version string to check against.
# Returns: 0 if the given version is less than or equal to the current version,
# 1 otherwise.
is_at_least_version() {
	test "$(vercmp "$(cat /etc/flasher-version)" "$1")" -ne -1
}

# has_mod_option
# Checks for a flag in the mod_options file.
#
# Parameters:
#   $1: The flag to check for.
# Returns: 0 if the flag is present, 1 otherwise.
has_mod_option() {
	if test -f "${_EXTRACTED_PATH}/mod_options" ; then
		grep -qe "^${1}\$" "${_EXTRACTED_PATH}/mod_options"
	else
		return 1
	fi
}

# require_version
# Just like is_at_least_version, except fatal is called automatically on
# failure.
require_version() {
	if ! is_at_least_version "$1" ; then
		fatal "This version of FlashCast ($(cat /etc/flasher-version)) is too old. Please update to at least version ${1}."
	fi
}

# mtd_lookup
# Converts an MTD partition name into a device name of the form "mtdX".
#
# Parameters:
#   $1: The name of the MTD partition.
# Returns: Prints the device name.
mtd_lookup() {
	_LOOKUP_RESULT="$(tail -n '+1' '/proc/mtd' | grep -e '"'"$1"'"$' | cut -d ':' -f 1)"
	if test -z "$_LOOKUP_RESULT" ; then
		log "Failed to look up MTD partition ${1}"
		return 1
	else
		echo "$_LOOKUP_RESULT"
	fi
}

# flash_mtd_partition
# Writes an image to an MTD partition.
#
# Parameters:
#   $1: The name (NOT number) of the MTD partition to flash.
#   $2: The path to an image file to flash.
# Returns: 0 on success, 1 on failure.
flash_mtd_partition() {
	_PARTITION="$(mtd_lookup "$1")"
	if test "$?" -ne 0 ; then
		return 1;
	fi

	log "Flashing ${2} to ${1} (${_PARTITION})"
	dry_run "Skipping flash on dry run"
	if test -z "$DRY_RUN" ; then
		flash_erase -q "/dev/${_PARTITION}" 0 0 &&
		nandwrite -qp "/dev/${_PARTITION}" "$2"
	fi
}

# mount_mtd_partition
# Mounts an MTD partition to a temporary directory.
#
# Parameters:
#   $1: The name (NOT number) of the MTD partition to mount.
# Returns: Prints the path to a mounted temporary directory.
mount_mtd_partition() {
	_PARTITION="$(mtd_lookup "$1")"
	if test "$?" -ne 0 ; then
		return 1;
	fi

	_MOUNT_DIR="$(mktemp -d)"
	log "Mounting ${1} (${_PARTITION}) at ${_MOUNT_DIR}"
	if ! mount "/dev/mtdblock${_PARTITION#mtd}" "$_MOUNT_DIR" ; then
		log "Could not mount MTD partition ${1}"
		rmdir "$_MOUNT_DIR"
		return 1;
	fi
	echo "$_MOUNT_DIR"
}

# cleanup_mount
# Unmounts and cleans up a directory mounted using a mount_* helper.
#
# Parameters:
#   $1: The path of the mount to clean up.
cleanup_mount() {
	log "Unmounting ${1}"
	umount "$1"
	rmdir "$1"
}

# clear_data
# Empties all stock files from the data partition.
_CLEAR_DATA_SKIP='. .. lost+found flashcast'
clear_data() {
	_DATA="$(mount_mtd_partition userdata)"
	log "Clearing userdata partition (${_DATA})"
	ls -a "$_DATA" | while read _PATH ; do
		for _I in $_CLEAR_DATA_SKIP ; do
			if test "$_PATH" = "$_I" ; then
				continue 2
			fi
		done

		dry_run "Would remove ${_DATA}/${_PATH}"
		if test -z "$DRY_RUN" ; then
			rm -rf "${_DATA}/${_PATH}"
		fi
	done
	cleanup_mount "$_DATA"
}

# begin_squashfs_edit
# Mounts a read-only squashfs partition in such a way that it can be modified
# and the changes stored in RAM. Be sure to call end_squashfs_edit when you're
# done to flash the edited filesystem back. NOTE: Due to limitations in the
# provided kernel source, this function can currently only be used with the
# "rootfs" MTD partition.
#
# Parameters:
#   $1: The name (NOT number) of the MTD partition to mount.
# Returns: Prints the path to the editable mount.
begin_squashfs_edit() {
	if test "$1" != 'rootfs' ; then
		log "begin_squashfs_edit can only be used with the rootfs MTD partition"
		return 1;
	fi

	_PARTITION="$(mtd_lookup "$1")"
	if test "$?" -ne 0 ; then
		return 1;
	fi

	_WRAPPER="$(mktemp -l -d)"
	log "Setting up squashfs union at ${_WRAPPER}"
	mkdir "${_WRAPPER}/base" "${_WRAPPER}/overlay" "${_WRAPPER}/union"

	if ! mount -t squashfs "/dev/mtdblock${_PARTITION#mtd}" "${_WRAPPER}/base" ; then
		log "Could not mount MTD partition ${1}"
		rm -r "$_WRAPPER"
		return 1;
	fi

	unionfs -o cow "${_WRAPPER}/overlay=RW:${_WRAPPER}/base=RO" "${_WRAPPER}/union"

	echo "$1" > "${_WRAPPER}/mtdpart"
	echo "${_WRAPPER}/union"
}

# begin_squashfs_image_edit
# Mounts a local squashfs image file to be edited and then written to disk with
# end_squashfs_edit.
#
# Parameters:
#   $1: The path of the local squashfs image.
#   $2: The name (NOT number) of the MTD partition to write the image to.
# Returns: Prints the path to the editable mount.
begin_squashfs_image_edit() {
	_WRAPPER="$(mktemp -l -d)"
	log "Setting up squashfs image union at ${_WRAPPER}"
	mkdir "${_WRAPPER}/base" "${_WRAPPER}/overlay" "${_WRAPPER}/union"

	if ! mount -o loop -t squashfs "$1" "${_WRAPPER}/base" ; then
		log "Could not mount local squashfs image ${1}"
		rm -r "$_WRAPPER"
		return 1
	fi

	unionfs -o cow "${_WRAPPER}/overlay=RW:${_WRAPPER}/base=RO" "${_WRAPPER}/union"

	echo "$2" > "${_WRAPPER}/mtdpart"
	echo "${_WRAPPER}/union"
}

# end_squashfs_edit
# Reflashes a squashfs from a previously-edited partition and cleans up the
# mount points.
#
# Parameters:
#   $1: The directory returned by begin_squashfs_edit.
end_squashfs_edit() {
	_END_SQUASHFS_EDIT_RETVAL=0

	_WRAPPER="${1%/union}"
	log "Reflashing squashfs union at ${_WRAPPER}"

	_REPACKED_FS="${_WRAPPER}/repack.sqfs"
	if ! mksquashfs "${_WRAPPER}/union" "$_REPACKED_FS" -no-progress ; then
		log "Failed to repack squashfs image at ${_WRAPPER}"
		_END_SQUASHFS_EDIT_RETVAL=1
	else
		flash_mtd_partition "$(cat "${_WRAPPER}/mtdpart")" "$_REPACKED_FS"
		_END_SQUASHFS_EDIT_RETVAL="$?"
	fi

	fusermount -u "${_WRAPPER}/union"
	umount "${_WRAPPER}/base"
	rm -r "${_WRAPPER}"

	return "$_END_SQUASHFS_EDIT_RETVAL"
}

# mktemp
# Returns a temporary file or directory name. Syntax is the same as mktemp, with
# the addition of an optional -l flag which must be the first argument if used.
#
# Parameters:
#   -l: If passed and a large temporary space is available, create the temporary
#       file in in instead of /tmp.
# Returns: Prints the name of the temporary file or directory.
_REAL_MKTEMP="$(which mktemp)"
mktemp() {
	if test "$1" = '-l' ; then
		shift
		if test -d "$LARGE_TEMP_DIR" ; then
			"$_REAL_MKTEMP" -p "$LARGE_TEMP_DIR" "$@"
		else
			log "Large temp file was requested but we have nowhere to put it!"
			"$_REAL_MKTEMP" "$@"
		fi
	else
		"$_REAL_MKTEMP" "$@"
	fi
}

# Set up an fd for log messages.
exec 3>&1

# If we're being sourced, return after defining functions.
if test "${BASH_SOURCE[0]}" != "$0" ; then
	return
fi

# Put the mod in a temp directory.
_MOD_FILE="${1%/}"
if test -f "$_MOD_FILE" -a "${_MOD_FILE%.zip}" != "$_MOD_FILE" ; then
	_EXTRACTED_PATH="$(mktemp -d)"
	echo "Extracting ${_MOD_FILE} (zip) to ${_EXTRACTED_PATH}"

	unzip -d "$_EXTRACTED_PATH" "$_MOD_FILE"
elif test -f "$_MOD_FILE" -a "${_MOD_FILE%.tar*}" != "$_MOD_FILE" ; then
	_EXTRACTED_PATH="$(mktemp -d)"
	echo "Extracting ${_MOD_FILE} (tar) to ${_EXTRACTED_PATH}"

	tar -xv -C "$_EXTRACTED_PATH" -f "$_MOD_FILE"
elif test -d "$_MOD_FILE" ; then
	_EXTRACTED_PATH="$(mktemp -u)"
	echo "Copying ${_MOD_FILE} to ${_EXTRACTED_PATH}"

	cp -r "$_MOD_FILE" "$_EXTRACTED_PATH"
	chmod 755 "$_EXTRACTED_PATH"
else
	echo "${_MOD_FILE} is not of a recognized mod format"
	exit 2
fi

# Check for a mod options file and copy if needed.
_MOD_OPTIONS_PATH="${_MOD_FILE}.options"
if test -f "$_MOD_OPTIONS_PATH" ; then
	cp "$_MOD_OPTIONS_PATH" "${_EXTRACTED_PATH}/mod_options"
fi

# Make sure to clean up when we're done.
_ORIG_PWD="$(pwd)"
_clean_up() {
	echo "Removing temporary mod directory ${_EXTRACTED_PATH}"
	cd "$_ORIG_PWD"
	rm -rf "$_EXTRACTED_PATH"
}

# Check for and run the imager script.
echo "Flashing mod at ${_EXTRACTED_PATH}"
trap _clean_up EXIT
cd "$_EXTRACTED_PATH"

if ! test -f './imager.sh' ; then
	echo "ERROR: $(pwd)/imager.sh does not exist"
	exit 1
fi

shift
source './imager.sh' "$@"