#!/bin/bash

# SPDX-License-Identifier:  GPL-2.0+

# Installer rootfs for ALT Linux
# Current version
VERSION=0.2.4

. shell-error

# Variable
TMPFIRM=
TMPROOT=
LOG="$TMP/alt-rootfs-installer.log"
PROGNAME="${0##*/}"
unset ROOTPART_UID
unset FIRMPART_UID

# usage message
usage() {
	echo "
Usage: $PROGNAME <options>

	--rootfs=ROOTFS	- archiv rootfs file name
	--media=DEVICE	- media device file (/dev/[sdX|mmcblkX])
	--target=TARGET	- target board
	--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 "$TMPFIRM" ] && [ -d "$TMPFIRM" ]; then
		umount "$TMPFIRM" >> "$LOG" 2>&1
		rmdir "$TMPFIRM" >> "$LOG" 2>&1
	fi
	if [ -n "$TMPROOT" ] && [ -d "$TMPROOT" ]; then
		umount "$TMPROOT" >> "$LOG" 2>&1
		rmdir "$TMPROOT" >> "$LOG" 2>&1
	fi
	if [ "$TARGET" = "jetson-nano" ]; then
		rm -fr "$TMP/Linux_for_Tegra"
	fi
	exit $rc
}

final_exit () {
	echo ""
	echo "= Installation Complete! Insert into the "$TARGET" and boot."
	exit 0
}

wget_jetson_nano () {
	if [ "$NOASK" != 1 ]; then
		rm -fr "$JETSON_PATH/$JETSON_ARCHIVE"
		echo "*** $JETSON_ARCHIVE not downloaded"
		read -p "= Download? " DOWNLOAD
		case "$DOWNLOAD" in
			[Yy][Ee][Ss])
				;;
			*)
				fatal "Error: $JETSON_ARCHIVE not downloaded!!!"
				;;
		esac
	fi
	echo "*** Download https://developer.nvidia.com/embedded/$JETSON_ARCHIVE..."
	wget "https://developer.nvidia.com/embedded/dlc/$JETSON_ARCHIVE" \
	  -P "$JETSON_PATH"
	echo "*** Extract Jetson Nano driver archive..."
	tar -xf "$JETSON_PATH/$JETSON_ARCHIVE" -C "$TMP"
	if [ $? -ne 0 ]; then
		rm -fr "$JETSON_PATH/$JETSON_ARCHIVE"
		fatal "Error: bad archive. Removed $JETSON_ARCHIVE"
	fi
}

# 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"
	DOC_DIR="${DIR}/"
else
	if [ -d "/usr/share/alt-rootfs-installer/boards.d" ]; then
		BOARDDIR="/usr/share/alt-rootfs-installer/boards.d"
		DOC_DIR="/usr/share/doc/alt-rootfs-installer/"
	else
		fatal "ERROR: directory boards.d not exist"
	fi
fi

# missing or wrongs arguments
[ $# -ge 1 ] &&
TMPARGS="$(getopt -n "$0" -o hy -l debug,help,target:,rootfs:,media:,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"
			;;
		--media)
			shift; MEDIA="$1"
			;;
		--supported)
			cat "$DOC_DIR/SUPPORTED-BOARDS"
			exit 0
			;;
		--version)
			echo "$PROGNAME-$VERSION"
			exit 0
			;;
		--log)
			shift
			LOG="$(readlink -f "$1")"
			;;
		-y)
			NOASK=1
			;;
		*)
			message "Error - ${1}"
			usage
			exit 1
			;;
	esac
	shift
done

# 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 "Can't locate directory ${LOG%/*} for log file. Exit"
fi
> "$LOG"

# check to make sure populated
if [ "$MEDIA" = "" ]; then
	usage
	exit 1
fi

if [ "$TARGET" = "" ] && [ "$ROOTFS" = "" ]; then
	usage
	exit 1
fi

if [ "$MEDIA" = "/dev/sda" ]; then
	echo ""
	echo " ***********************************************************"
	echo " ** WARNING: You have requested the rootfs be written to sda."
	echo " ** /dev/sda is usually the root filesystem of the host. "
	echo " ***********************************************************"
	echo " ** Do you wish to continue? (type 'yes' to continue)"
	echo " ***********************************************************"
	# wait for agreement
	read -p " = Continue? " AGREE
	case "$AGREE" in
		[Yy][Ee][Ss])
			;;
		*)
			echo "User exit, no rootfs written."
			exit 0
			;;
	esac
fi

# rootfs exists
if [ ! -f "$ROOTFS" ] && [ "$ROOTFS" != "" ]; then
	fatal "Error: $ROOTFS not found! Please choose an existing rootfs."
fi


# definition of the rootfs archive format and ARCH
if [ -n "$ROOTFS" ]; then
	case "$ROOTFS" in
		*.tar) TAR='tar -xf'
			;;
		*.tar.gz) TAR='tar -zxf'
			;;
		*.tar.xz) TAR='tar -Jxf'
			;;
		*) fatal 'Error: unknown rootfs archive format'
			;;
	esac
	# NOTE: We assume there is no arch names containing '-' or '.'
	ARCH="${ROOTFS##*-}"
	ARCH="${ARCH%%.*}"
fi

# device exists
if [ ! -e "$MEDIA" ]; then
	fatal "Error: $MEDIA not found! Please choose an existing device."
fi

# change cubietruck target to uppercase
if [ "$TARGET" = "cubietruck" ]; then
	TARGET="Cubietruck"
fi

# check for boards
if [ "$TARGET" != "" -a ! -e "${BOARDDIR}/${TARGET}" ]; then
	message "Error: You must choose a supported board or none at all." 
	usage
	exit 1
fi

if [ "$TARGET" != "" -a -L "$BOARDDIR/$TARGET" ]; then
	SOCS="$(basename $(readlink -e "$BOARDDIR/$TARGET"))"
	SOCS_ARCH="${SOCS#*-}"
else
	SOCS="$TARGET"
	unset SOCS_ARCH
fi

# socs and ARCH compliance check
if [ "$SOCS_ARCH" != "" -a "$ARCH" != "" -a "$SOCS_ARCH" != "$ARCH" ]; then
	fatal "Error: You must choose a supported board or none at all."
fi

if [ "$TARGET" = "rpi3" ] && [ "$ARCH" = "armh" ]; then
	fatal "Error: This rootfs does not support Raspberry Pi 3"
fi
if [ "$TARGET" = "rpi2" ] && [ "$ARCH" = "aarch64" ]; then
	fatal 'Error: This rootfs does not support Raspberry Pi 2'
fi

# Nvidia Jetson Nano driver packages download
if [ "$TARGET" = "jetson-nano" ]; then
	JETSON_ARCHIVE="l4t-jetson-driver-package-32-1-jetson-nano"
	JETSON_PATH="/root"
	rm -fr "$TMP/Linux_for_Tegra"
	if [ -f  ]; then
		echo "Extract Jetson Nano driver archive..."
		tar -xf "$JETSON_PATH/$JETSON_ARCHIVE" -C "$TMP" ||
			wget_jetson_nano
	else
		wget_jetson_nano
	fi
fi

clear
# Last chance to back out
echo ""
echo "====================================================="
echo "= Selected Media:  $MEDIA"
# target hardware ARCH
if [ "$TARGET" != "" ]; then 
	echo "= U-Boot Target:   $TARGET"
fi
# rootfs if included
if [ "$ROOTFS" != "" ]; then
	echo "= Selected rootfs: $ROOTFS"
	echo "= Log file:        $LOG"
	echo "====================================================="
	echo " "
	echo "*****************************************************"
	echo "*****************************************************"
	echo "******** WARNING! ALL DATA WILL BE DESTROYED ********"
	echo "*****************************************************"
	echo "*****************************************************"
else
	echo "= Log file:        $LOG"
	echo "====================================================="
	echo " "
	echo "*****************************************************"
	echo "*****************************************************"
	echo "******* WARNING! BOOTLOADER WILL BE REWRITED ********"
	echo "*****************************************************"
	echo "*****************************************************"

fi

if [ "$NOASK" != 1 ]; then
	echo " "
	echo " Type 'YES' to proceed, anything else to exit now "
	echo " "
	# wait for agreement
	read -p "= Proceed? " PROCEED
	case "$PROCEED" in
		[Yy][Ee][Ss])
			;;
		*)
			echo "User exit, no rootfs written."
			exit 0
			;;
	esac
fi

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

# umount before starting
umount "${MEDIA}${partsuffix}"?* >> "$LOG" 2>&1

# Create partitions
if [ "$ROOTFS" != "" ]; then
	case "$SOCS" in
		jetson_nano-aarch64)
			unset FIRMPART
			ROOTPART="${MEDIA}${partsuffix}1"
			cd "$TMP/Linux_for_Tegra" \
			&& sed -i 's/sudo//g' create-jetson-nano-sd-card-image.sh \
			&& yes |./create-jetson-nano-sd-card-image.sh -o sd.img -s 128M -r 100  >> "$LOG" 2>&1 \
			&& dd if=sd.img of="$MEDIA" bs=1M >> "$LOG" 2>&1 \
			&& sgdisk -g "$MEDIA" >> "$LOG" 2>&1 \
			&& yes |parted -s "$MEDIA" resizepart 1 100% >> "$LOG" 2>&1 \
			&& yes |mkfs.ext4 -L ROOT "$ROOTPART" >> "$LOG" 2>&1 \
			|| fatal "Error: create partition failed!!!"
			;;
		Rockchips-*)
			unset FIRMPART
			ROOTPART="${MEDIA}${partsuffix}1"
			echo 'y' | sgdisk -g --clear \
			--new=1:32768: --change-name=1:ROOT --typecode=1:0FC63DAF-8483-4772-8E79-3D69D8477DE4 \
			"$MEDIA" >> "$LOG" 2>&1 \
			&& sleep 2 && echo 'y' | mkfs.ext4 -L ROOT "$ROOTPART" >> "$LOG" 2>&1 \
			&& parted -s "$MEDIA" set 1 "legacy_boot" on >> "$LOG" 2>&1 \
			|| fatal "Error: create partition failed!!!"
			;;
		*-riscv64)
			unset FIRMPART
			ROOTPART="${MEDIA}${partsuffix}2"
			echo 'y' | sgdisk -g --clear \
			--new=1:2048:67583 --change-name=1:bootloader --typecode=1:2E54B353-1271-4842-806F-E436D6AF6985 \
			--new=2:264192:    --change-name=2:root       --typecode=2:0FC63DAF-8483-4772-8E79-3D69D8477DE4 \
			"$MEDIA" >> "$LOG" 2>&1 \
			&& sleep 2 && echo 'y' | mkfs.ext4 -L ROOT "$ROOTPART" >> "$LOG" 2>&1 \
			|| fatal "Error: create partition table failed!!!"
			;;
		rpi2|rpi3)
			FIRMPART="${MEDIA}${partsuffix}1"
			ROOTPART="${MEDIA}${partsuffix}2"
			parted -s "$MEDIA" mktable "msdos" >> "$LOG" 2>&1 \
			&& parted -s -a optimal "$MEDIA" mkpart primary fat16 2MiB 40MiB >> "$LOG" 2>&1 \
			&& sleep 2 && echo 'y' | mkfs.fat $FIRMPART >> "$LOG" 2>&1 \
			&& parted -s -a optimal "$MEDIA" mkpart primary ext4 40MiB 100% >> "$LOG" 2>&1 \
			&& sleep 2 && echo 'y' | mkfs.ext4 -L ROOT "$ROOTPART" >> "$LOG" 2>&1 \
			&& parted -s "$MEDIA" set 2 boot on >> "$LOG" 2>&1 \
			|| fatal "Error: create partitions failed!!!"
			;;
		*)
			unset FIRMPART
			ROOTPART="${MEDIA}${partsuffix}1"
			parted -s "$MEDIA" mktable "msdos" >> "$LOG" 2>&1 && sync \
			&& parted -s -a optimal "$MEDIA" mkpart primary ext4 2MIB 100% >> "$LOG" 2>&1 \
			&& sleep 2 && echo 'y' | mkfs.ext4 -L ROOT "$ROOTPART" >> "$LOG" 2>&1 \
			&& parted -s "$MEDIA" set 1 boot on >> "$LOG" 2>&1 \
			|| fatal "Error: create partitions failed!!!"
			;;
	esac
	echo "Partitions created successfully"
	# check to see how many partitions on the rootfs
	sleep 2
	partprobe "$MEDIA"
	if [ $? -ne 0 ]; then
		fatal "Error: create partitions failed!!!"
	fi
	# UID definition
	if [ "$FIRMPART" != "" ]; then
		FIRMPART_UID="$(blkid $FIRMPART |sed -E 's;.*\sUUID="([^"]*)".*;\1;')"
	fi
	ROOTPART_UID="$(blkid $ROOTPART |sed -E 's;.*\sUUID="([^"]*)".*;\1;')"
else
	# check device Partition Table compliance target
	case "$SOCS" in
		Rockchips-*|*-riscv64)
			PARTTABLE=gpt
			;;
		*)
			PARTTABLE=msdos
			;;
	esac
	unset MEDIATABLE
	parted "${MEDIA}" print | grep gpt >> "$LOG" 2>&1 && MEDIATABLE="gpt"
	parted "${MEDIA}" print | grep msdos >> "$LOG" 2>&1 && MEDIATABLE="msdos"
	if [ "$PARTTABLE" != "$MEDIATABLE" ]; then
		fatal "Error: Partition Table not $PARTTABLE"
	fi
	unset TYPEFS
	parted "${MEDIA}${partsuffix}1" print | grep "fat" >> "$LOG" 2>&1 && TYPEFS="fat"
	if [ "$TYPEFS" = "" ]; then
		unset FIRMPART
		if [ "$SOCS" = bbl-riscv64 ]; then
			ROOTPART="${MEDIA}${partsuffix}2"
		else
			ROOTPART="${MEDIA}${partsuffix}1"
		fi
		if [ "$TARGET" = "rpi2" -o "$TARGET" = "rpi3" ]; then
			fatal "Error: Firmware partition not found!!!"
		fi
	else
		FIRMPART="${MEDIA}${partsuffix}1"
		ROOTPART="${MEDIA}${partsuffix}2"
	fi
fi

trap pre_exit EXIT HUP PIPE INT QUIT TERM

# make temp mount points
if [ "$FIRMPART" != "" ]; then
	TMPFIRM="$(mktemp -d --tmpdir 'firmpart.XXXXXXXX' 2>>"$LOG")"
	mount "$FIRMPART" "$TMPFIRM" >> "$LOG" 2>&1
	if [ $? -ne 0 ]; then
		fatal "Error: mount $FIRMPART $TMPFIRM failed!!!"
	fi
fi
if [ "$ROOTPART" != "" ]; then
	TMPROOT="$(mktemp -d --tmpdir 'rootpart.XXXXXXXX' 2>>"$LOG")"
	mount "$ROOTPART" "$TMPROOT" >> "$LOG" 2>&1
	if [ $? -ne 0 ]; then
		fatal "Error: mount $ROOTPART $TMPROOT failed!!!"
	fi
fi

sleep 1

# Write the disk rootfs to media
if [ "$ROOTFS" != "" ]; then
	echo "= Writing: "
	echo "= $ROOTFS "
	echo "= To: $MEDIA ...."
	$TAR "$ROOTFS" -C "$TMPROOT"
	if [ $? -ne 0 ]; then
		fatal "Error: writing rootfs!!!"
	fi
	sync;
	echo "= Writing rootfs complete!"
fi

PREFIX="$TMPROOT"

# Update /etc/fstab
if [ "$FIRMPART_UID" != "" ]; then
	mkdir -p "$PREFIX"/mnt/FIRMPART
	echo "UUID=$FIRMPART_UID	/mnt/FIRMPART	vfat	ro	0 0" >> "$PREFIX"/etc/fstab
fi
if [ "$ROOTPART_UID" != "" ]; then
	sed -i "s/LABEL=ROOT/UUID=$ROOTPART_UID/" "$PREFIX"/etc/fstab
fi

if [ "$TARGET" = "jetson-nano" ]; then
	echo "= Install drivers for Nvidia Jetson Nano"
	. "${BOARDDIR}/${TARGET}" >> "$LOG" 2>&1
	if [ $? -ne 0 ]; then
		fatal "Error: install drivers for Nvidia Jetson Nano failed!!!"
	fi
	final_exit
fi

if [ "$TARGET" = "rpi2" ]; then
	echo "= Install bootloader for Raspberry Pi 2"
	rm -f "$TMPFIRM"/kernel8.img
	cp -f "${PREFIX}"/usr/share/u-boot/rpi_3_32b/* "$TMPFIRM"/
	if [ $? -ne 0 ]; then
		fatal "Error: install bootloder for Raspberry Pi 2 (armh)!!!"
	fi
	final_exit
fi

if [ "$TARGET" = "rpi3" ]; then
	echo "= Install bootloader for Raspberry Pi 3"
	rm -f "$TMPFIRM"/kernel7.img
	cp -f "${PREFIX}"/usr/share/u-boot/rpi_3/* "$TMPFIRM"/
	if [ $? -ne 0 ]; then
		fatal "Error: install bootloder for Raspberry Pi 3 (aarch64)!!!"
	fi
	final_exit
fi

if [ "$SOCS" = bbl-riscv64 ]; then
	unset KERNEL
	KERNEL="$(readlink -e $PREFIX/boot/vmlinuz |sed 's;.*/;;')"
	if [ "$KERNEL" = "" ]; then
		fatal "Error: kernel not found"
	fi
	. "$BOARDDIR/$TARGET"
	if [ $? -ne 0 ]; then
		fatal "Error: install bootloader for $TARGET!!!"
	fi
	final_exit
fi

# determine uboot and write to disk
if [ "$TARGET" != "" ]; then
	if [ -d "${PREFIX}/usr/share/u-boot/${TARGET}" ]; then
		. "${BOARDDIR}/${TARGET}"
	else
		fatal "No U-Boot files found for $TARGET!!!"
	fi
fi

final_exit

