#!/usr/bin/env bash

source /hive/bin/colors
[[ -e /etc/brand.conf ]] && source /etc/brand.conf
[[ -z $BRAND ]] && BRAND="Hive"

HIVEREPO_LIST=/etc/apt/sources.list.d/hiverepo.list

# pci.ids update source, otherwise download from web
PCIIDS=/hive/etc/pci.ids
# usb.ids update source
USBIDS=/hive/etc/usb.ids
# amdgpu.ids AMD GPUs definition file used by some miners
AMDGPUIDS=/hive/etc/amdgpu.ids

[[ -f /hive-config/.DISKLESS_AMD && $(cat /proc/mounts | grep "/ " | awk '{print $1}') == tmpfs ]] && echo "${RED}Diskless rig. Selfupgrade restricted${NOCOLOR}" && exit 0

[[ -f $RIG_CONF ]] && source $RIG_CONF

#This will prevent asking interactive questions
export DEBIAN_FRONTEND=noninteractive
#If needed maybe there is even stronger option @see https://askubuntu.com/questions/146921/how-do-i-apt-get-y-dist-upgrade-without-a-grub-config-prompt
#sudo DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" dist-upgrade

function print_help {
	demourl="http://192.168.0.1/repo/binary"
	[[ -z $DEFAULT_REPO_URL ]] && DEFAULT_REPO_URL=`grep -m1 -oP "deb \K[^\s]+" /hive/etc/hiverepo.list`
	echo "${CYAN}Usage: selfupgrade [option] [repository URL] [version]${NOCOLOR}
Description:
  Tool for upgrading OS from $BRAND repository.
Options:
${GREEN}  -h|--help        ${NOCOLOR}to show current help
${GREEN}  -f|--force       ${NOCOLOR}to force upgrade by cleaning apt lists, caches and locks
		    it can help in situations when selfupgrade says 'System is up to date' but actually it isn't
${GREEN}  -g|--grub        ${NOCOLOR}to update only grub config
${GREEN}  -k|--kernel      ${NOCOLOR}to update only kernel and modules
${GREEN}  -c|--consolefont ${NOCOLOR}to update console font
${GREEN}  repository URL   ${NOCOLOR}to get updates from; should start with http like ${demourl} or ${DEFAULT_REPO_URL}
${GREEN}  version          ${NOCOLOR}to upgrade or downgrade; should be used the following format x.y-zz, like 0.6-65
Examples:
${GREEN}  selfupgrade 0.6-65 ${demourl} -f ${NOCOLOR}
                   force upgrade/downgrade to version 0.6-65 from repository ${demourl}
"
}

#screen -wipe > /dev/null 2>&1

existing_repository_list=
function set_repo {
	local repository_list="deb $1 /"
	[[ -e $HIVEREPO_LIST ]] && existing_repository_list="$(< $HIVEREPO_LIST)"
	echo "${CYAN}> Using repository ${WHITE}$param${NOCOLOR}"
	echo "$repository_list" > $HIVEREPO_LIST
}


function restore_repo {
	if [[ ! -z "$existing_repository_list" && "$existing_repository_list" != "$(< $HIVEREPO_LIST)" ]]; then
		#echo "${CYAN}> Restoring old repository URLs${NOCOLOR}"
		echo "$existing_repository_list" > $HIVEREPO_LIST
	fi
}
# restore repo on any error
trap "restore_repo" EXIT


function update_grub {
	source /etc/default/grub
	new_params=()
	readarray -t params < <( grep -oP '^\K[^#\"\\]+' /hive/etc/grub.custom | tr [:space:] '\n' | grep -v '^$' )
	for param in "${params[@]}"; do
		[[ ! " $GRUB_CMDLINE_LINUX_DEFAULT " =~ " $param " ]] && new_params+=("$param")
	done
	if [[ "${#new_params[@]}" -gt 0 ]]; then
		echo "${YELLOW}> Using additional kernel boot parameters for grub${NOCOLOR}"
		echo "${new_params[@]}"
		new_cmdline="$GRUB_CMDLINE_LINUX_DEFAULT ${new_params[@]}"
		sed -ure "s/^GRUB_CMDLINE_LINUX_DEFAULT.+?/GRUB_CMDLINE_LINUX_DEFAULT=\"${new_cmdline}\"/g" -i /etc/default/grub
		return 0
	fi
	return 1
}


# Check config grub and create symlink
function check_grub {
	[[ ! -f /hive/etc/grub ]] && return 1

	#if [[ $(readlink /etc/default/grub) != "/hive/etc/grub" ]]; then
	#	rm /etc/default/grub #>/dev/null 2>&1
	#	ln -s /hive/etc/grub /etc/default/grub #>/dev/null 2>&1
	#fi

	# remove symlink
	readlink /etc/default/grub >/dev/null &&
		rm -f /etc/default/grub

	[[ -f /etc/default/grub ]] && grub="$(< /etc/default/grub)" || grub=

	# update grub config to default if needed
	diff -q /hive/etc/grub /etc/default/grub 2>/dev/null ||
		cp -f /hive/etc/grub /etc/default/grub

	[[ ! -f /hive/etc/grub.custom && -f /hive/etc/grub.custom.template ]] &&
		cp /hive/etc/grub.custom.template /hive/etc/grub.custom

	if [[ -f /hive/etc/grub.custom ]]; then
		( update_grub ) # in subshell
	fi

	[[ "$grub" == "$(< /etc/default/grub)" ]] && return 0
	echo "${CYAN}> Updating grub${NOCOLOR}"
	update-grub
	exitcode=$?
	echo "${YELLOW}System restart is required to get the effect${NOCOLOR}"
	return $exitcode
}


function stop_services {
	[[ $asw -ne 0 ]] && echo "${CYAN}> Stopping autoswitch${NOCOLOR}" && autoswitch stop > /dev/null 2>&1
	[[ $miners -ne 0 ]] && echo "${CYAN}> Stopping miners${NOCOLOR}" && miner stop
	autofan stop > /dev/null
}


function restore_services {
	echo "${CYAN}> Restarting autofan and watchdog${NOCOLOR}"
	wd restart
	autofan restart > /dev/null

	# in maintenance mode start miner if it was running
	if [[ $MAINTENANCE -eq 0 || ( $MAINTENANCE -eq 1 && $miners -eq 1 ) ]]; then
		echo "${CYAN}> Starting miners${NOCOLOR}"
		miner start
		if [[ $asw -ne 0 ]]; then
			echo "${CYAN}> Starting autoswitch${NOCOLOR}"
			nohup bash -c 'sleep 15 && autoswitch start' > /dev/null 2>&1 &
		fi
	fi
}

kernel_upgrade() {
	echo "${CYAN}> Updating package lists${NOCOLOR}"
	[[ -e $ALT_REPO_LIST ]] && cat $ALT_REPO_LIST >> $HIVEREPO_LIST
	hpkg update
	local exitcode=$?
	if [[ $exitcode -eq 0 ]]; then
		booted_kernel=`uname -r`
		booted_version=`uname -v | sed -r 's/.*version //;s/SMP.+//'`
		short_version=${booted_kernel%.*}
		current_installed=`dpkg -s hiveos-kernel-$short_version 2>&1| grep '^Version: ' | sed 's/Version: //'`
		if [[ -z $current_installed ]]; then
			echo "${RED} No hiveos-kernel package installed. Nothing to upgrade${NOCOLOR}"
			restore_repo # Restore repository
			hello
			exit
			fi
		echo "${CYAN}> Booted kernel: ${WHITE}$short_version${NOCOLOR}"
		echo "${CYAN}> Installed package: ${WHITE}$current_installed${NOCOLOR}"
		kernel_package=`apt list --upgradable 2>&1 | grep "^hiveos-kernel-$short_version" | tr '/' ' ' | awk '{print $1 " " $3 " " $5 " " $6 " " $7}' | sort -V`
		if [[ -z "$kernel_package" ]]; then
			echo "${GREEN}> Kernel package is up to date${NOCOLOR}"
			restore_repo # Restore repository
			hello
			exit
		fi

		echo "${CYAN}> Updating packages${NOCOLOR}"
		echo "${WHITE}$kernel_package${NOCOLOR}"

		pkgs=`echo "$kernel_package" | cut -d" " -f1 | grep -E "^hiveos-kernel-$short_version"`
		if [[ ! -z $pkgs ]]; then
			apt install -y -f --only-upgrade --ignore-hold --allow-change-held-packages $pkgs
			[[ $? -ne 0 ]] && echo "${RED}(exitcode=$exitcode)${NOCOLOR}"
		fi

		current_installed=`dpkg -s hiveos-kernel-$short_version 2>&1| grep '^Version: ' | sed 's/Version: //'`
		echo "${CYAN}> Installed new package: ${WHITE}$current_installed${NOCOLOR}"
		echo "${GREEN}> Please reboot your rig.${NOCOLOR}"
		restore_repo
		hello
		exit $exitcode
	fi
	restore_repo
	hello
	exit $exitcode
	
}


VERSION=

for param in "$@"; do
	case "$param" in
		--help|-h)
			print_help
			exit 0
		;;
		--force|-f) # it can help in situations when selfupgrade says Hive is up to date but this is not true
			echo "${CYAN}> Cleaning apt lists${NOCOLOR}"
			# force HiveOS lists download
			rm -f /var/lib/apt/lists/*hiveos*
			# remove apt/dpkg lock files
			echo "${CYAN}> Removing apt/dpkg lock files and updates${NOCOLOR}"
			rm -f /var/lib/dpkg/lock*
			rm -f /var/lib/apt/lists/lock*
			rm -f /var/cache/apt/archives/lock*
			# remove partial updates
			rm -rf /var/lib/dpkg/updates/*
			# fixing dpkg install status
			dpkg-query -l hive --show >/tmp/dpkg_state.tmp 2>&1
			if [[ $? -eq 2 && $(cat /tmp/dpkg_state.tmp | grep -c "error: parsing file '/var/lib/dpkg/status'") -gt 0 ]]; then
				echo "${CYAN}> Try fixing dpkg install status db${NOCOLOR}"
				mv /var/lib/dpkg/status /var/lib/dpkg/status.BAD
				cp -f /var/lib/dpkg/status-old /var/lib/dpkg/status
			fi
			rm /tmp/dpkg_state.tmp
			apt update
		;;

		--grub|-g)
			check_grub
			exit
		;;

		--kernel|-k)
			kernel_upgrade
			exit
		;;

		--consolefont|-c)
			if [[ -f /hive/etc/console-setup ]]; then
				setupcon # update current console
				[[ ! -f /etc/default/console-setup || "$(< /hive/etc/console-setup)" != "$(< /etc/default/console-setup)" ]] &&
					cp -f /hive/etc/console-setup /etc/default/console-setup &&
					update-initramfs -u # update boot image
			fi
			exit
		;;

		http*) # Set user repository
			set_repo "$param"
		;;

		*) # version
			VERSION="$param"
			if [[ ! -z "$VERSION" ]]; then
				if [[ ! "$VERSION" =~ ^[0-9]+\.[0-9]+-[0-9]+$ && ! "$VERSION" =~ ^[0-9]+\.[0-9]+-[0-9]+\.[0-9]+$ ]]; then
					echo "${YELLOW}Invalid version format \"$VERSION\"${NOCOLOR}"
					echo "Please use the following format x.y-zz, like 0.6-200"
					exit 1
				fi
				#echo "${CYAN}> Upgrading to version ${YELLOW}$VERSION${NOCOLOR}"
			fi
		;;
	esac
done


# get current version
current_version=`dpkg -s hive | grep '^Version: ' | sed 's/Version: //'`
[[ -e /hive-config/RepoVer ]] && repover=$(< /hive-config/RepoVer) || repover=0
echo "${CYAN}> Current version: ${WHITE}$current_version@$repover${NOCOLOR}"

# remove stupid rep file for smaller updates
[[ -f /etc/apt/sources.list.d/amdgpu-pro.list ]] && rm /etc/apt/sources.list.d/amdgpu-pro.list

# remove avrdude conf file to exclude overwrite request
[[ -f /etc/avrdude.conf ]] && rm /etc/avrdude.conf

# save miner and autoswitch status
asw=$(screen -ls | grep -c autoswitch)
miner status > /dev/null && miners=1 || miners=0

# check available memory and stop miners in advance
if [[ $miners -eq 1 ]]; then
	mem=`free -b | grep "Mem" | awk '{print int($7/1024/1024)}'` || mem=0
	if [[ $mem -lt 250 ]]; then
		echo "${YELLOW}> Low memory available - $mem MB. Stopping miners${NOCOLOR}"
		stop_services
	fi
fi

# Sometimes Ubuntu fails to update repos
hpkg fix
dexitcode=$? # return on exit to show a problem
[[ $dexitcode -ne 0 ]] && echo "${RED}(exitcode=$dexitcode)${NOCOLOR}"

echo "${CYAN}> Updating package lists${NOCOLOR}"
[[ -e $ALT_REPO_LIST ]] && cat $ALT_REPO_LIST >> $HIVEREPO_LIST
hpkg update
exitcode=$?

if [[ $exitcode -eq 0 ]]; then
	hive_packages=""

	if [[ -z "$VERSION" ]]; then # no target version
		#hive_packages=`apt list --upgradable 2>&1 | grep '^hive' | tr '/' ' ' | sort -V`
		hive_packages=`apt list --upgradable 2>&1 | grep '^hive' | tr '/' ' ' | awk '{print $1 " " $3 " " $5 " " $6 " " $7}' | sort -V`
		if [[ -z "$hive_packages" ]]; then
			# force repover, hello and grub update even with no update
			repover-touch
			check_grub
			hello
			[[ -e /hive-config/RepoVer ]] && newrepover=$(< /hive-config/RepoVer) || newrepover=0
			if [[ "$repover" != "$newrepover" ]]; then
				echo "${CYAN}> Upgrade to ${WHITE}${current_version}@${newrepover}${CYAN} complete${NOCOLOR}"
				echo "${GREEN}Have a happy mining${NOCOLOR}"
			else
				echo "${GREEN}> System is up to date already${NOCOLOR}"
			fi
			restore_repo # Restore repository
			exit $dexitcode
		fi
	fi

	# Stop miners, autofan, autoswitch to prevent reboots during upgrade
	stop_services

	# allow dependent packages downgrade
	readarray -t repos < <(grep -oP "^deb.*:\/\/\K[^\/]+" $HIVEREPO_LIST)
	for repo in "${repos[@]}"; do
		list="/etc/apt/preferences.d/hive-${repo//./-}"
		[[ ! -f $list ]] && echo -e "Package: *\nPin: origin \"$repo\"\nPin-Priority: 1001\n" > $list
	done

	if [[ -z "$VERSION" ]]; then # normal upgrade to the latest available
		echo "${CYAN}> Updating packages${NOCOLOR}"
		echo "${WHITE}$hive_packages${NOCOLOR}"
		pkgs=`echo "$hive_packages" | cut -d" " -f1 | grep -E "^(hive$|hive-miners|hive-lib-dotnet|hive-opt-algomap)"` # only miners and few more pkgs
		if [[ ! -z $pkgs ]]; then
			apt-get install -y -f --only-upgrade hive $pkgs # hive must be always present
			exitcode=$?
			[[ $exitcode -ne 0 ]] && echo "${RED}(exitcode=$exitcode)${NOCOLOR}"
		fi
		# try to update extra pkgs. skip errors
		extra_pkgs=$(apt list --upgradable 2>&1 | grep '^hive' | tr '/' ' ' | cut -d" " -f1)
		if [[ ! -z $extra_pkgs ]]; then
			apt-get install -y -f --only-upgrade hive $extra_pkgs || echo "${YELLOW}Some packages were not updated${NOCOLOR}"
		fi
	else
		echo "${CYAN}> Installing version: ${BYELLOW}$VERSION${NOCOLOR}"
		# Bring hive's whole dependency set to the target release too, not just the
		# hive package. hive's Depends are '>=' minimums, so a plain 'hive=$VERSION'
		# leaves already-satisfying newer deps (hive-lib, hive-cuda-*, hive-nvflash,
		# ...) in place -> a mismatched downgrade. Name the installed hive* closure
		# so the repo pin (Pin-Priority 1001 set above) forces each down to the
		# target version. --installed --recurse keeps this bounded to what is here.
		dep_pkgs=$(apt-cache depends --installed --recurse --no-recommends --no-suggests \
			--no-conflicts --no-breaks --no-replaces --no-enhances hive 2>/dev/null |
			grep -oP '^hive[a-z0-9-]*$' | sort -u | grep -vxF hive | tr '\n' ' ')
		apt-get install -y -f --allow-downgrades --reinstall hive=$VERSION $dep_pkgs
		exitcode=$?
		[[ $exitcode -ne 0 ]] && echo "${RED}(exitcode=$exitcode)${NOCOLOR}"
	fi
fi

restore_repo # Restore repository

if [[ $exitcode -ne 0 ]]; then
	restore_services
	echo "${BRED}> Upgrade failed${NOCOLOR}"
	exit $exitcode
fi

# Save RepoVer
repover-touch

echo "${CYAN}> Cleanup packages${NOCOLOR}"
apt-get -y autoremove
apt-get clean

# Reread env variables
source /etc/environment
export $(cat /etc/environment | grep -vE '^$|^#' | cut -d= -f1) #export all variables from file

# update Grub
check_grub

# update pci.ids from local source
if [[ -e $PCIIDS && $PCIIDS -nt /usr/share/misc/pci.ids ]]; then
	echo "${CYAN}> Updating pci.ids to version ${WHITE}$(grep -m 1 -oP "Version: \K.*" $PCIIDS)${NOCOLOR}"
	[[ -f /usr/share/misc/pci.ids && -d /usr/share/misc ]] && mv -f /usr/share/misc/pci.ids /usr/share/misc/pci.ids.old
	[[ -d /usr/share/misc ]] && ln -f -s $PCIIDS /usr/share/misc/pci.ids
fi

# update usb.ids from local source: /usr/share/misc on modern images,
# /var/lib/usbutils kept for legacy ones (see hive-pkg/DEBIAN/postinst)
for ids_dir in /usr/share/misc /var/lib/usbutils; do
	[[ -e $USBIDS && -d $ids_dir ]] || continue
	[[ $USBIDS -nt $ids_dir/usb.ids ]] || continue
	echo "${CYAN}> Updating $ids_dir/usb.ids to version ${WHITE}$(grep -m 1 -oP "Version: \K.*" $USBIDS)${NOCOLOR}"
	[[ -f $ids_dir/usb.ids ]] && mv -f $ids_dir/usb.ids $ids_dir/usb.ids.old
	ln -f -s $USBIDS $ids_dir/usb.ids
done

# update amdgpu.ids from local source: /usr/share/libdrm feeds the system
# libdrm (= ROCm v6+ stack), /opt/amdgpu the legacy backend, /opt/rocm* just
# in case; per-directory freshness (see hive-pkg/DEBIAN/postinst)
for ids_dir in /usr/share/libdrm /opt/amdgpu/share/libdrm /opt/rocm*/share/libdrm; do
	[[ -e $AMDGPUIDS && -d $ids_dir ]] || continue
	[[ $AMDGPUIDS -nt $ids_dir/amdgpu.ids ]] || continue
	echo "${CYAN}> Updating $ids_dir/amdgpu.ids to version ${WHITE}$(grep -m 1 -oP "Version: \K.*" $AMDGPUIDS)${NOCOLOR}"
	ln -f -s $AMDGPUIDS $ids_dir/amdgpu.ids
done

screen -wipe > /dev/null 2>&1 # Wipe possible dead screen

# so reboot can be avoided in case it was updated
echo "${CYAN}> Running gpu-detect${NOCOLOR}"
gpu-detect listjson > $GPU_DETECT_JSON

# Moved from postinst again as hello need to see new version
#echo "${CYAN}> Saying hello to server again${NOCOLOR}"
hello nostart

# Restart agent
echo "${CYAN}> Restarting agent${NOCOLOR}"
agent-screen restart

# Restore other services
restore_services

# Restarting ROH Fan controller
if [[ `lsusb 2>/dev/null | grep -c 16c0:05dc` -ge 1 ]]; then
	screen -S octofan -X quit > /dev/null
	$OCTOFAN dontattach
fi

new_version=`dpkg -s hive | grep '^Version: ' | sed 's/Version: //'`
[[ -e /hive-config/RepoVer ]] && repover=$(< /hive-config/RepoVer) || repover=0

if dpkg --compare-versions "$current_version" "le" "$new_version" || [[ -z $VERSION ]]; then
	echo "${CYAN}> Upgrade to ${WHITE}$new_version@$repover${CYAN} complete${NOCOLOR}"
else
	echo "${CYAN}> Downgrade to ${WHITE}$new_version@$repover${CYAN} complete${NOCOLOR}"
	echo "${YELLOW}Reboot is highly recommended after downgrade!${NOCOLOR}"
fi
echo "${GREEN}Have a happy mining${NOCOLOR}"

exit $dexitcode
