#!/bin/bash
#  vim:sw=8:noet
# SPDX-License-Identifier:  GPL-2.0+
#
# rootfs installer for ALT Linux
# Current version
VERSION=0.4.0

# TODO:
#  * Check $TMP and mktemp if could not found

. shell-error
. shell-terminfo

# Variable
EXT4=0FC63DAF-8483-4772-8E79-3D69D8477DE4
TMPROOT=
LOG="$TMP/alt-rootfs-installer.log"
PROGNAME="${0##*/}"
unset ROOTPART_UID

# one Mebibyte = 2^20
MB=1048576

# Print the message to the terminal. This function is intended to be used with
# print_done() or print_fail().
print_message() {
	echo -n "${1-}"
	PRINT_TEXT_END=1
}

print_done() {
	terminfo_col "$PRINT_COLUMN"
	color_message "[DONE]" green
	PRINT_TEXT_END=
}

print_fail() {
	terminfo_col "$PRINT_COLUMN"
	color_message "[FAIL]" red
	PRINT_TEXT_END=
}

fatal_error() {
	[ -n "${PRINT_TEXT_END-}" ] && print_fail
	color_text "Error: " red
	fatal "${1-}"
}

# usage message
usage() {
	echo "
Usage: $PROGNAME --target=TARGET <--rootfs=ROOTFS | --image-in=IMAGE> <--media=DEVICE | --image-out=FILE> [options]

Options:
	--target=TARGET   - target board

	--rootfs=ROOTFS   - archive rootfs file name
	--image-in=IMAGE  - image file name (img|img.xz)

	--media=DEVICE    - media device file (/dev/[sdX|mmcblkX])
	--image-out=FILE  - out image file

	--image-size=SIZE - size out image file (MiB)
	--resize          - Resize ROOT partition
	--supported       - List of supported hardware
	--version         - Display version and exit
	--log             - Log file (default $LOG)
	-y                - Assumes yes, will not wait for confirmation

Example: $PROGNAME --rootfs=regular-lxqt-20190213-aarch64.tar.xz --target=rpi3 --media=/dev/mmcblk0

Update U-Boot for another Target
Update to a new u-boot from a local host install.

Example: $PROGNAME --target=orangepi_prime --media=/dev/mmcblk0

"
}

pre_exit() {
	local rc=$?
	trap - EXIT
	sync
	if [ -n "$TMPROOT" ] && [ -d "$TMPROOT" ]; then
		umount "$TMPROOT" >> "$LOG" 2>&1
		rmdir "$TMPROOT" >> "$LOG" 2>&1
	fi

	if [ -n "$TMPFIRM" ] && [ -d "$TMPFIRM" ]; then
		umount "$TMPFIRM" >> "$LOG" 2>&1
		rmdir "$TMPFIRM" >> "$LOG" 2>&1
	fi

	LC_ALL=C type soc_exit >&/dev/null && soc_exit

	if [ -n "$IMAGE_OUT" -a -n "$MEDIA" ]; then
		kpartx -d -s "$MEDIA" || {
			sleep 10
			kpartx -d -s -v "$MEDIA"
		}
		losetup --detach "$MEDIA"
	fi
	exit $rc
}

final_exit() {
	echo "Installation Complete!"
	if [ -n "$IMAGE_OUT" ]; then
		echo "Image created successfully."
	else
		echo "Insert into the "$TARGET" and boot."
	fi
	exit 0
}

# Initialize shell-terminfo from libshell
terminfo_init
PRINT_COLUMN="$(( COLUMNS - 6 ))"

# Set some global variables for the command directory, target board directory,
# and valid targets.
DIR="$(dirname $0)"
if [ -d "${DIR}/boards.d" ]; then
	BOARDDIR="${DIR}/boards.d"
	SOCSDIR="${DIR}/socs.d"
	DOC_DIR="${DIR}/"
else
	if [ -d "/usr/share/alt-rootfs-installer/boards.d" ]; then
		BOARDDIR="/usr/share/alt-rootfs-installer/boards.d"
		SOCSDIR="/usr/share/alt-rootfs-installer/socs.d"
		DOC_DIR="/usr/share/doc/alt-rootfs-installer/"
	else
		fatal_error "directory boards.d not exist"
	fi
fi

. ${SOCSDIR}/socs-utils

# missing or wrongs arguments
[ $# -ge 1 ] &&
TMPARGS="$(getopt -n "$0" -o hy -l debug,help,target:,rootfs:,media:,\
image-in:,image-out:,image-size:,resize,supported,version,log: -- "$@")"
if [ $? -ne 0 ]; then
	usage
	exit 1
fi

eval set -- "$TMPARGS"
# check the args
while :; do
	case $1 in
		--)
			shift; break
			;;
		--debug)
			set -x
			;;
		-h|--help)
			usage; exit 0
			;;
		--target)
			shift; TARGET="$1"
			;;
		--rootfs)
			shift; ROOTFS="$1"
			;;
		--image-in)
			shift; IMAGE="$1"
			;;
		--media)
			shift; MEDIA="$1"
			;;
		--image-out)
			shift; IMAGE_OUT="$1"
			;;
		--image-size)
			shift; IMAGE_SIZE="$1"
			;;
		--resize)
			RESIZE="1"
			;;
		--supported)
			cat "$DOC_DIR/SUPPORTED-BOARDS"
			exit 0
			;;
		--version)
			echo "$PROGNAME-$VERSION"
			exit 0
			;;
		--log)
			shift
			LOG="$(readlink -f "$1")"
			;;
		-y)
			NOASK=1
			;;
		*)
			echo "Invalid argument $1"
			usage
			exit 1
			;;
	esac
	shift
done

# getting full PATH
[ -n "$ROOTFS" ] && ROOTFS_PATH="$(readlink -f "$ROOTFS")"
[ -n "$IMAGE" ] && IMAGE_PATH="$(readlink -f "$IMAGE")"

# checking file existing
if [ -n "$ROOTFS" ]; then
	[ -f "$ROOTFS_PATH" ] || fatal_error "ROOTFS does not exist"
	ROOTFS="$ROOTFS_PATH"
fi
if [ -n "$IMAGE" ]; then
	[ -f "$IMAGE_PATH" ] || fatal_error "IMAGE does not exist"
	IMAGE="$IMAGE_PATH"
fi

# replace old targets
[ "$TARGET" = "rpi2" ] && TARGET="rpi_3_32b"
[ "$TARGET" = "rpi3" ] && TARGET="rpi_3"

# checking the entered value --image-size
if [ -n "${IMAGE_SIZE-}" ]; then
	[[ "$IMAGE_SIZE" =~ [^0-9]+ ]] && fatal_error "bad image size"
fi

# ensure sudo user
if [ "$(whoami)" != "root" ]; then
	fatal_error "this script requires 'sudo' privileges in order to write to disk & mount media"
fi

# cleanup log file
if [ ! -d "${LOG%/*}" ]; then
	fatal_error "can't locate directory ${LOG%/*} for log file"
fi
> "$LOG"

if [ -z "$IMAGE_OUT" -a -z "$MEDIA" ]; then
	fatal_error "one of the options must be defined: --media or --image-out"
fi

if [ -n "$IMAGE_OUT" -a -n "$MEDIA" ]; then
	fatal_error "At the same time, mutually exclusive options are indicated: --media and --image-out"
fi

if [ -n "$IMAGE" -a -n "$ROOTFS" ]; then
	fatal_error "At the same time, mutually exclusive options are indicated: --image and --rootfs"
fi

if [ -z "$TARGET" -a -z "$ROOTFS" -a -z "$IMAGE" ]; then
	usage
	exit 1
fi

if [ -n "$MEDIA" ] && [[ "$(readlink -f /dev/root)" ==  "$MEDIA"* ]]; then
	fatal_error "$MEDIA is a device contains a host system rootfs"
fi

# rootfs exists
if [ -n "$ROOTFS" -a ! -f "$ROOTFS" ]; then
	fatal_error "$ROOTFS not found! Please choose an existing rootfs"
fi

# image exists
if [ -n "$IMAGE" -a ! -f "$IMAGE" ]; then
	fatal_error "$IMAGE not found! Please choose an existing image"
fi

# device exists
if [ -z "$IMAGE_OUT" -a ! -e "$MEDIA" ]; then
	fatal_error "$MEDIA not found! Please choose an existing device"
fi

# target board exists
if [ -n "$TARGET" -a ! -e "${BOARDDIR}/${TARGET}" ]; then
	echo "Please choose a supported board"
	usage
	exit 1
fi

# check image format
if [ -n "$IMAGE" ]; then
	case "$IMAGE" in
		*.img)
			CAT='cat'
			;;
		*.img.xz)
			CAT='xzcat'
			;;
		*) fatal_error "unknown image format"
			;;
	esac
	# NOTE: We assume there is no arch names containing '-' or '.'
	ARCH="${IMAGE##*-}"
	ARCH="${ARCH%%.*}"
fi

# definition of the rootfs archive format and ARCH
if [ -n "$ROOTFS" ]; then
	case "$ROOTFS" in
		*.tar)
			TAR='tar -xf'
			let "TAR_SIZE = $(stat -Lc %s "$ROOTFS") / $MB"
			;;
		*.tar.gz)
			TAR='tar -zxf'
			let "TAR_SIZE = $(gzip --list "$ROOTFS" | awk 'END{print $2}') / $MB"
			;;
		*.tar.xz)
			TAR='tar -Jxf'
			let "TAR_SIZE = $(xz --robot --list "$ROOTFS" | awk '/^totals/{print $5}') / $MB"
			;;
		*) fatal_error "unknown rootfs archive format"
			;;
	esac
	# NOTE: We assume there is no arch names containing '-' or '.'
	ARCH="${ROOTFS##*-}"
	ARCH="${ARCH%%.*}"
fi

if [ -n "$TARGET" ]; then
	if [ -L "$BOARDDIR/$TARGET" ]; then
		SOCS="$(basename $(readlink -e "$BOARDDIR/$TARGET"))"
		SOCS_ARCH="${SOCS#*-}"
	else
		SOCS="$TARGET".sh
		unset SOCS_ARCH
	fi
else
	SOCS="unknown-soc.sh"
	unset SOCS_ARCH
fi

[ -n "$SOCS" ] || SOCS="unknown-soc.sh"

case "$MEDIA" in
	/dev/mmcblk*|/dev/loop*)
		partsuffix="p"
		;;
	*)
		unset partsuffix
		;;
esac

# Source SoCs specific shell script
. "$SOCSDIR/$SOCS"
soc_init

# Find Firmware Partition
if [ -z "$ROOTFS" -a -z "$IMAGE" ]; then
	FIRMPART=
	ROOTPART=1
	parted "${MEDIA}${partsuffix}1" print | grep "fat" >> "$LOG" 2>&1 &&
		FIRMPART=1
	[ -n "$FIRMPART" ] && ROOTPART_NUM=2
fi

[ -n "$FIRMPART_NUM" ] && FIRMPART="${MEDIA}${partsuffix}${FIRMPART_NUM}"
[ -n "$ROOTPART_NUM" ] && ROOTPART="${MEDIA}${partsuffix}${ROOTPART_NUM}"
[ -n "$BBLPART_NUM" ] && BBLPART="${MEDIA}${partsuffix}${BBLPART_NUM}"

# Image size calculation
if [ -n "$ROOTFS" -a -n "$IMAGE_OUT" ]; then
	# Add 256 MiB to the image size for bios, firmware, bootloader, etc and increase size by 10%
	# which is reserved by FS, take into account i-node ratio and reserved 32Mb. Add 1 just in case.
	IMAGE_SIZE_MIN="$(awk -vTS="$TAR_SIZE" -vM="$MB" 'BEGIN{ printf ("%d", ((TS+512+32)*(1+256/8192)*1.1+1)) }')"
	# Change image size
	if [ -z "${IMAGE_SIZE-}" ] || [ "$IMAGE_SIZE_MIN" -ge "$IMAGE_SIZE" ]; then
		IMAGE_SIZE="$IMAGE_SIZE_MIN"
	fi
	let "END_SECTOR = $IMAGE_SIZE - 1"
	END_SECTOR="$END_SECTOR"MiB
fi

# Last chance to back out
if [ -n "$IMAGE_OUT" -a -n "$IMAGE_SIZE" ]; then
	echo -n 'Image out:'; terminfo_col 20; color_message "$IMAGE_OUT" green
	echo -n "Image size:"; terminfo_col 20; color_message "$IMAGE_SIZE MiB" green
else
	echo -n 'Selected Media:'; terminfo_col 20; color_message "$MEDIA" green
fi
# target hardware ARCH
if [ -n "$TARGET" ]; then
	echo -n 'Target:'; terminfo_col 20; color_message "$TARGET" green
fi
# rootfs if included
if [ -n "$ROOTFS" ]; then
	echo -n 'Selected rootfs:'; terminfo_col 20; color_message "$ROOTFS" green
fi
# image-in if included
if [ -n "$IMAGE" ]; then
	echo -n 'Selected image:'; terminfo_col 20; color_message "$IMAGE" green
fi
echo -n "Log file: "; terminfo_col 20; color_message "$LOG" green

if [ -z "$IMAGE_OUT" ]; then
	if [ -n "$IMAGE" -o -n "$ROOTFS" ]; then
		color_message "WARNING! ALL DATA WILL BE DESTROYED" red
	else
		color_message "WARNING! BOOTLOADER WILL BE OVERWRITTEN" red
	fi
fi

if [ "$NOASK" != 1 ]; then
	# Ask user before start
	color_message "Would you like to continue? [Yes/No] " bold
	read PROCEED
	case "$PROCEED" in
		[Yy][Ee][Ss])
			;;
		*)
			echo "User exit, no rootfs written."
			exit 0
			;;
	esac
fi

trap pre_exit EXIT HUP PIPE INT QUIT TERM

# Check to make image file
if [ -n "$IMAGE_OUT" ]; then
	MEDIA="$(losetup --find)"
	[ -n "$MEDIA" ] || fatal_error "/dev/loop* not allowed"
fi

# Call SoC specific function to write system
[ -n "$IMAGE" -o -n "$ROOTFS" ] && soc_write_rootfs
[ -n "$TARGET" ] && soc_write_bootloader

final_exit
