#!/bin/sh


VERSION="1.1"


# prevent running from pipe
if [ "$0" = "sh" ] || [ "$0" = "bash" ]; then
	sleep 1
	echo "Tool hive-migrate does not work from pipe, exiting"
	exit 1
fi


if ! mount | grep -Eq "^(tmpfs|rootfs) on / "; then
	echo "Current OS is not supported. Try to use 'hive-replace'!"
	exit 1
fi


# download timeout
TIMEOUT=20
RAMDISK_SIZE=2000

set -o pipefail

export LANG=C

# main url to images repo
DOWNLOAD_URL=https://download.hiveos.farm
# direct urls to images
STABLE_URL=$DOWNLOAD_URL/latest/

OS_BRAND="hiveos"
BRAND="Hive"
HIVE_HOST_URL="https://api.hiveos.farm"

# image file mask
IMG_EXT="(zip|xz|gz)"
URL_PROTO="(http|ftp|https)"
OS_REGEX="$OS_BRAND-[^\"]+\.$IMG_EXT"
[ ! -z "$SKIP_CHECK" ] && [ $SKIP_CHECK -ne 0 ] && OS_REGEX="[^\"]+\.$IMG_EXT"
URL_MASK="^$URL_PROTO://.+/${OS_REGEX}$"


TMPDIR="/run/tmphive"
url=
farm_hash=""
rig_conf=""
force=0
no_confirm=0
configured=0
nosafe=0
batch_mode=0


NOCOLOR=$'\033[0m'
BLACK=$'\033[0;30m'
DGRAY=$'\033[1;30m'
RED=$'\033[0;31m'
BRED=$'\033[1;31m'
GREEN=$'\033[0;32m'
BGREEN=$'\033[1;32m'
YELLOW=$'\033[0;33m'
BYELLOW=$'\033[1;33m'
BLUE=$'\033[0;34m'
BBLUE=$'\033[1;34m'
PURPLE=$'\033[0;35m'
BPURPLE=$'\033[1;35m'
CYAN=$'\033[0;36m'
BCYAN=$'\033[1;36m'
LGRAY=$'\033[0;37m'
WHITE=$'\033[1;37m'
GRAY=$'\033[2;37m'


exit_handler() {
	echo -e "\npress any key to close session"
	if [ -t 0 ]; then
		read -n 1
	else
		read -n 1 -t 600
	fi
}


url_decode() {
	echo -e "${1//%/\\x}"
}


unpack_image() {
	local output="$1"
	echo "" > $TMPDIR/progress
	if [ -t 0 ]; then
		( sleep 2; while true; do pkill -USR1 -f "dd of" || break; sleep 1; echo -en "\r$(tail -n 1 $TMPDIR/progress 2>/dev/null) "; done ) &
	fi
	case "$ARCHNAME" in
		*\.xz)	xz -dkc "$TMPDIR/$ARCHNAME" | dd of=$output bs=1M 2>$TMPDIR/progress ;;
		*\.zip)	unzip -p "$TMPDIR/$ARCHNAME" | dd of=$output bs=1M 2>$TMPDIR/progress ;;
		*\.gz)	gzip -dkc "$TMPDIR/$ARCHNAME" | dd of=$output bs=1M 2>$TMPDIR/progress ;;
		*)	echo "${RED}Unsupported format: $ARCHNAME${NOCOLOR}" && return 1 ;;
	esac
	exitcode=$?
	wait
	echo -e "\r$(tail -n 1 ./progress 2>/dev/null) "
	return $exitcode
}


get_image_from_url() {
	local from_url="$( url_decode "$1" )"
	local name="$2"
	url="$(curl --connect-timeout $TIMEOUT --retry 2 -fskL --head -w '%{url_effective}' "$from_url" 2>/dev/null | tail -n 1)"
	if [ $? -eq 0 ]; then
		echo "$url" | grep -qE "$URL_MASK" && return 0
		echo -e "$RED> Got bad url for $name image - $url${NOCOLOR}"
		return 1
	fi
	echo -e "$RED> Unable to get url for $name image${NOCOLOR}"
	return 2
}


get_image() {
	local from_url="$( url_decode "$1" )"
	ARCHNAME="$(basename "$from_url")"

	[ -z "$ARCHNAME" ] && echo -e "${RED}Unable to get filename from ${BRED}$from_url${NOCOLOR}" && return 1

	# Copy|download hive image to tmpfs
	echo -e "\n${CYAN}> Downloading image file${NOCOLOR}"

	case "$from_url" in
		http*|ftp*)
			echo ""
			[ -t 0 ] && opt="--progress-bar" || opt="--no-progress-meter"
			curl -kLJ --retry 99 $opt -o "$TMPDIR/$ARCHNAME" "$from_url"
			[ $? -ne 0 ] && echo -e "${RED}Download $BRAND OS image failed. Check url${NOCOLOR}" && return 1
		;;
		*)
			[ ! -f "$from_url" ] && echo -e "${RED}$BRAND OS image not found. Check path${NOCOLOR}" && return 1
			cp -v "$from_url" "$TMPDIR/$ARCHNAME"
			[ $? -ne 0 ] && echo -e "${RED}Copy $BRAND OS image failed${NOCOLOR}" && return 1
		;;
	esac

	# check archive integrity
	if [ ! -z "$SKIP_CHECK" ] && [ $SKIP_CHECK -eq 2 ]; then
		echo -e "\n${YELLOW}> Checking archive integrity - SKIPPED!!!${NOCOLOR}"
	else
		echo -e "\n${CYAN}> Checking archive integrity${NOCOLOR}"
		unpack_image /dev/null
		[ $? -ne 0 ] && echo -e "${RED}Archive is damaged${NOCOLOR}" && return 1
	fi

	echo -e "${GREEN}$BRAND OS image ready!${NOCOLOR}"
	return 0
}


configure_image() {
	[ ! -f $TMPDIR/rig.conf ] && return 1

	echo -e "\n${CYAN}> Configuring $BRAND OS${NOCOLOR}"

	partprobe $DISK 2>/dev/null

	PART=$(blkid | grep "hiveroot" | grep -oE "$DISK[^:]+")
	if [ -z "$PART" ]; then
		# using alternative method
		losetup -P -f $DISK
		PART=$(blkid | grep "hiveroot" | grep -oE "/dev/loop[^:]+")
	fi
	[ -z "$PART" ] && echo -e "${RED}Unable to find config partition${NOCOLOR}" && return 1

	echo "Using partition $PART"

	CFGDIR=$TMPDIR/mnt-config
	mkdir -p $CFGDIR

	mount $PART $CFGDIR
	[ $? -ne 0 ] && echo -e "${RED}Unable to mount partition${NOCOLOR}" && return 1

	if [ -d $CFGDIR/hive ]; then
		# cat $TMPDIR/rig.conf > $CFGDIR/rig.conf && configured=1
		cp $TMPDIR/rig.conf $CFGDIR && echo -e "${GREEN}$BRAND OS configuration complete!${NOCOLOR}" &&
			configured=1 || echo -e "${RED}Unable to update configuration${NOCOLOR}"
		sync		
	else
		echo -e "${RED}Invalid partition${NOCOLOR}"
   	fi

	umount -f $PART || echo -e "${RED}Unable to unmount partition${NOCOLOR}"
	loop="$(expr "$PART" : '\(/dev/loop[0-9][0-9]*\)')" && losetup -d "$loop"
}


start_migration() {
	memtotal=$(free -m | grep 'Mem' | awk '{print $2}')
	memavail=$(free -m | grep 'Mem' | awk '{print $7}')

	echo
	echo -e "${CYAN}***************************************************${NOCOLOR}"
	echo -e "${YELLOW}         $BRAND OS Image Migration Tool v$VERSION${NOCOLOR}"
	echo -e "${CYAN}***************************************************${NOCOLOR}"
	echo
	echo -e "Image $WHITE$url${NOCOLOR}"
	[ ! -z "$farm_hash" ] && echo -e "Farm hash ${WHITE}$farm_hash$rig_conf${NOCOLOR}"
	[ ! -z "$rig_conf" ] && echo -e "Config ${WHITE}$farm_hash$rig_conf${NOCOLOR}"
	echo
	echo -e "RAM Total ${WHITE}$memtotal ${NOCOLOR}MB, Available ${WHITE}$memavail ${NOCOLOR}MB"
	echo

	[ $memtotal -lt 2700 ] && echo -e "\n${RED}The minimum RAM for $BRAND OS is 3 GB, exiting${NOCOLOR}" && exit 2
	[ $memavail -lt 1400 ] && echo -e "\n${RED}Not enough free memory, exiting${NOCOLOR}" && exit 3

	DISK=`mount | grep -oE "^/dev/(sd[a-z]+|nvme[0-9]+)"` 
	[ -z "$DISK" ] && echo -e "\n${RED}Unable to determine disk for installation${NOCOLOR}" && exit 4

	DISK_NAME="${DISK/\/dev\/}"
	DISK_SIZE="$(grep -w "${DISK_NAME}$" /proc/partitions | awk '{printf("%.f",$3/1024)}')"
	DISK_INFO="$(cat /sys/block/sda/device/vendor /sys/block/sda/device/model 2>/dev/null | xargs)"

	echo -e "Booted from ${BPURPLE}$DISK_NAME${NOCOLOR}: ${WHITE}$DISK_INFO $DISK_SIZE MB${NOCOLOR}"

	[ $DISK_SIZE -lt 7400 ] && echo -e "\n${RED}The minimum disk size for $BRAND OS is 8 GB${NOCOLOR}, exiting" && exit 5

	if [ $no_confirm -ne 1 ]; then
		echo -e "\n${BYELLOW}Warning: After $BRAND OS installation, all your data on ${BPURPLE}$DISK_NAME${NOCOLOR} ${BYELLOW}will be lost!${NOCOLOR}"
		echo -en "\nIf you want to install $BRAND OS on ${BPURPLE}$DISK_NAME${NOCOLOR} type ${GREEN}\"yes\"${NOCOLOR}: "
		read -t 90 answer
		[ "$answer" != "yes" ] && echo -e "${BCYAN}\nBye, bye!${NOCOLOR}" && exit 0
	fi

	command -v stop >/dev/null && echo && stop

	# create tmpfs
	grep $TMPDIR /proc/mounts >/dev/null 2>&1
	if [[ $? -ne 0 ]]; then
		mkdir -p $TMPDIR
		mount none $TMPDIR -t tmpfs -o size=${RAMDISK_SIZE}m
	fi
	rm -f $TMPDIR/*

	# self copy
	cp $0 $TMPDIR/hive-migration || exit 9

	# update config
	if [ -f "$rig_conf" ]; then
		cp $rig_conf $TMPDIR/rig.conf
	elif [ ! -z "$farm_hash" ]; then
		echo -e "HIVE_HOST_URL=\"$HIVE_HOST_URL\"" >> $TMPDIR/rig.conf
		echo -e "FARM_HASH=$farm_hash" >> $TMPDIR/rig.conf
	fi

	# turn off swap here
	swapoff -a

	# get hive zip
	get_image "$url" || exit 6

	# disable kernel message to tty
	echo 0 > /proc/sys/kernel/printk
	echo 1 > /proc/sys/kernel/sysrq
	echo 0 > /proc/sysrq-trigger
	sync

	cd $TMPDIR
	exec $TMPDIR/hive-migration --continue "$DISK" "$ARCHNAME"
}


continue_migration() {
	[ -z "$DISK" ] && echo -e "${BRED}No disk is set, exiting${NOCOLOR}" && exit 6
	[ -z "$ARCHNAME" ] || [ ! -f "$TMPDIR/$ARCHNAME" ] && echo -e "${RED}No image file, exiting${NOCOLOR}" && exit 6

	DISK_NAME="${DISK/\/dev\/}"

	# umount/remount RO
	need_ro=0
	sync
	for MOUNT in `grep $DISK /proc/mounts | awk '{print $2}'`; do
		echo -e "\n${CYAN}> Unmounting ${WHITE}$MOUNT${NOCOLOR}"
		umount -f $MOUNT > /dev/null 2>&1 && continue
		mount -o remount,ro $MOUNT > /dev/null 2>&1
		exitcode=$?
		if [[ $? -eq 0 ]]; then
			echo "${NOCOLOR}> mounted read-only${NOCOLOR}"
		else
			echo "${NOCOLOR}> mount failed ($exitcode)${NOCOLOR}"
			need_ro=1
		fi
	done
	if [[ $need_ro -eq 1 ]]; then
		echo "${NOCOLOR}> mounting all read-only"
		sync
		echo u > /proc/sysrq-trigger
		sleep 1
		[[ ! -w $TMPDIR ]] && mount -n -o remount,rw $TMPDIR
	fi

	# Writing														   
	echo -e "\n${CYAN}> Writing $BRAND OS filesystem to ${BPURPLE}$DISK_NAME${NOCOLOR}"
	echo -e "Please wait, this can take long"
	echo -e "To prevent damage to your disk device\ndo not turn off your computer until the end of the process.\n"

	unpack_image $DISK
	[ $? -ne 0 ] && echo -e "${BRED}Write image failed, exiting${NOCOLOR}" && exit 8

	sync

	echo -e
	echo -e "${BGREEN}Image writing to ${BPURPLE}$DISK_NAME${BGREEN} is successfull!${NOCOLOR}"

	# we do not need it any more
	rm "$TMPDIR/$ARCHNAME" > /dev/null 2>&1

	configure_image

	echo -e "\n${WHITE}Your rig is ready to work under $BRAND OS!${NOCOLOR}"

	if [ $configured -eq 1 ]; then
		RIG_ID=
		RIG_PASSWD=
		FARM_HASH=
		HIVE_HOST_URL=
		if [ -f $TMPDIR/rig.conf ]; then
			source $TMPDIR/rig.conf
		else
			echo -e "${BYELLOW}Empty config (default)${NOCOLOR}"
		fi
		[ ! -z "$RIG_ID" ] && [ ! -z "$RIG_PASSWD" ] && [ ! -z "$HIVE_HOST_URL" ] &&
			echo -e "RIG ID: ${BPURPLE}$RIG_ID${NOCOLOR}, API: ${BPURPLE}$HIVE_HOST_URL${NOCOLOR}"
		[ ! -z "$FARM_HASH" ] && [ ! -z "$HIVE_HOST_URL" ] && [ -z "$RIG_ID" ] &&
			echo "FARM HASH: ${BPURPLE}$FARM_HASH${NOCOLOR}, API: ${BPURPLE}$HIVE_HOST_URL${NOCOLOR}"
	else
		echo -e "${BRED}Empty config (not written)${NOCOLOR}"
	fi

	echo -e "\n${GREEN}Have a happy mining!${NOCOLOR}"
	echo -en "\n${YELLOW}Rebooting in 20 seconds${NOCOLOR}"

	sleep 1

	# Sync
	echo 1 > /proc/sys/kernel/sysrq
	echo s > /proc/sysrq-trigger

	sleep 20

	# Reboot
	echo 1 > /proc/sys/kernel/sysrq
	echo b > /proc/sysrq-trigger
}


usage() {
	local code="$1"
	echo -e "${NOCOLOR}Usage:${CYAN} hive-migrate [option] <Path or URL to ZIP/XZ file with $BRAND OS image>${NOCOLOR}"
	echo -e "By default is running in session that can be resumed by ${CYAN}screen -r hive${NOCOLOR} or ${CYAN}hive-migrate${NOCOLOR}"
	echo -e "Options:"
	echo -e "${GREEN}  -s|--stable	   ${NOCOLOR}Download and install latest ${BCYAN}Stable${NOCOLOR} image from $BRAND server"
	echo -e "${GREEN}  -y|--yes		  ${NOCOLOR}Do not ask for confirmation, answer yes to all questions"
	echo -e "${GREEN}  --hash=FARM_HASH  ${NOCOLOR}FARM_HASH from Web-interface. Optional for already configured rig"
	echo -e "${GREEN}  --rigconf[=FILE]  ${NOCOLOR}Use provided rig.conf file for new image (./rig.conf by default)"
	[ ! -z "$code" ] && exit $code
	echo -e "Examples:"
	echo -e "${GREEN}  hive-migrate -y --stable${NOCOLOR}"
	echo -e "${GREEN}  hive-migrate -y --hash=46e9602837d0bda99f0 http://mylocalserver/$OS_BRAND-0.6.zip${NOCOLOR}"
	exit 0
}


[ "$1" == "-h" ] || [ "$1" == "--help" ] && usage


if [ "$1" == "--continue" ]; then
	DISK="$2"
	ARCHNAME="$3"
	trap exit_handler EXIT
	continue_migration
	exit
fi

# Get root
if [ $(id -u) -ne 0 ]; then
	echo "Root privileges required"
	sudo chmod +x $0
	exec sudo $0 "$@"
fi


echo " $* " | grep " --no-safe " && nosafe=1
if command -v screen >/dev/null; then
	# check for running session and restore to it
	if screen -ls hive >/dev/null; then
		screen -r hive 2>/dev/null >/dev/null && exit
	elif [ $nosafe -eq 0 ]; then
		[ -z "$*" ] && usage
		echo "${BCYAN}Starting hive-migrate session${NOCOLOR}"
		sleep 1
		if [ -t 0 ]; then
			exec screen -mS hive $0 "$@"
		else
			screen -dmS hive $0 "$@"
			echo "Use interactive shell to connect to it with 'hive-migrate' or 'screen -r hive' cmd"
		fi
		exit
	fi
	# set exit handler to pause at the end
	trap exit_handler EXIT
fi


for param in "$@"; do
	case "$param" in
#		-b|--batch)
#			batch_mode=1
#		;;
		--safe|--no-safe)
			# ignore it here
		;;
		--force)
			force=1
		;;
		-y|--yes)
			no_confirm=1
		;;
		--hash=*)
			farm_hash="$(expr "$param" : '--hash=\(.*\)')"
		;;
		--rigconf)
			rig_conf="./rig.conf"
		;;
		--rigconf=*)
			rig_conf="$(expr "$param" : '--rigconf=\(.*\)')"
			[ -z "$rig_conf" ] && rig_conf=" "
		;;
		-s|--stable)
			get_image_from_url $STABLE_URL "Stable" || exit
		;;
		*)
			if echo "$param" | grep -qE "^$URL_PROTO://.+"; then
				if [ ! -z "$SKIP_CHECK" ] && [ $SKIP_CHECK -ne 0 ] || echo "$param" | grep -q "$OS_BRAND"; then
					get_image_from_url "$param" "this" || exit
				else
					echo -e "${RED}Bad url for $OS_BRAND image: ${WHITE}$param${NOCOLOR}${RED}, exiting${NOCOLOR}"
					exit 1
				fi
			elif echo "$param" | grep -qE "$OS_REGEX$"; then
				url="$( url_decode "$param" )"
			else
				echo -e "${RED}Unknown param: ${WHITE}$param${NOCOLOR}${RED}, exiting\n${NOCOLOR}" && usage 1
			fi
		;;
	esac
done


[ -z "$rig_conf" ] && [ -z "$farm_hash" ] && [ -f ./rig.conf ] && rig_conf="./rig.conf"
[ ! -z "$rig_conf" ] && [ ! -f "$rig_conf" ] && echo -e "${RED}Rig.conf not found, exiting${NOCOLOR}" && exit 1
[ ! -z "$farm_hash" ] && [ ! -z "$rig_conf" ] && echo -e "${YELLOW}--hash and --rigconf options are incompatible, exiting\n" && usage 1
[ -z "$url" ] && echo -e "${YELLOW}Filesystem path or URL to $BRAND OS image is required.${NOCOLOR}\n" && exit 1

start_migration