#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

# This script helps the update of the target release in L4T APT source
# list file and optionally starts the upgrade process if specified.

set -e

APT_REPO_URL="https://repo.download.nvidia.com/jetson"
REL_META_URL="${APT_REPO_URL}/releases"
REL_META_FILE=$(mktemp -q)
IGX_REL_FILE="/etc/igx-release"
OS_NAME="l4t"
PLAT_JETSON="jetson"
PLAT_IGX="igx"
releases=()

# Defaults
do_list="n"
do_upgrade="n"
release="latest"

cleanup()
{
	rm -f "${REL_META_FILE}"
}

check_pre_req()
{
	local this_user

	this_user="$(whoami)"
	if [ "${this_user}" != "root" ]; then
		echo "ERROR: please run as sudo or root user"
		exit 1
	fi
}

is_igx()
{
	[ -f "${IGX_REL_FILE}" ]
}

download_rel_meta_file()
{
	echo "Downloading ${REL_META_URL}"
	if ! wget -q "${REL_META_URL}" -O "${REL_META_FILE}"; then
		echo "Not able to donwload release meta file!"
		exit 1
	fi
}

parse_rel_meta_file()
{
	local platform

	if is_igx; then
		platform="${PLAT_IGX}"
	else
		platform="${PLAT_JETSON}"
	fi
	mapfile -t releases < <( \
		grep -v -e "^#" "${REL_META_FILE}" | \
		awk -v platform="${platform}" '$0 ~ platform {print $1}'
	)
	if [ ${#releases[@]} -eq 0 ]; then
		echo "ERROR: no available releases found"
		exit 1
	fi
}

list_available_releases()
{
	for i in "${releases[@]}"; do
		echo "${i}"
	done
}

update_source_list()
{
	local pat_repo="deb ${APT_REPO_URL}/.*"
	local pat_rel="r[.0-9]*"
	local pat_comp="main"
	local bsp_core_pkg="nvidia-l4t-core"
	local dgpu_apt_src_pkg="nvidia-l4t-dgpu-apt-source"
	local source_list="/etc/apt/sources.list.d/nvidia-l4t-apt-source.list"
	local target="${1}"
	local current

	if [ "${target}" = "latest" ]; then
		target="${releases[-1]}"
	elif [[ ! " ${releases[*]} " == *" ${target} "* ]]; then
		echo "ERROR: release ${target} is not available!"
		exit 1
	fi

	IFS="." read -ra digits <<< "${target#r}"
	if [ ${#digits[@]} -eq 2 ]; then
		target+=".0"
	fi

	current=$(dpkg -s "${bsp_core_pkg}" | awk '/^Version:/{print $2}')
	if dpkg --compare-versions "${target#r}" lt "${current%-*}"; then
		echo "ERROR: downgrading is not supported"
		exit 1
	fi

	if dpkg-query -W -f='${Status}' "${dgpu_apt_src_pkg}" 2>/dev/null | grep -q "install ok installed"; then
		source_list="/etc/apt/sources.list.d/nvidia-l4t-dgpu-apt-source.list"
	fi

	echo "Updating source list to release ${target}"
	sed -ri "s|(^$pat_repo )($pat_rel)( $pat_comp)|\1${target}\3|" \
	    "${source_list}"
}

do_upgrade()
{
	echo "Executing upgrade..."
	apt-get update
	apt-get dist-upgrade -y
	echo "Upgrade is done successfully."
}

usage()
{
	cat << EOF

Usage: $SCRIPT_NAME [-l|--list] [-u|--do-upgrade] [-h|--help] [release]

This script is a helper script for ${OS_NAME} Debian OTA update. It will update
the target release in ${OS_NAME} APT source list accroding to the release passed
in.

Options:
 -l, --list		List available releases.
 -u, --do-upgrade	Do upgrade after updating the source list.
 -h, --help		Print the message.
 release		Target release version. Downgrading is not supported.
			If not specified, The last release version will be used.
EOF
}

usage_error()
{
	usage >&2
	exit 2
}

trap cleanup EXIT

SCRIPT_NAME="$(basename "${0}")"
OPTIONS=$(getopt -o "luh" --long list,do-upgrade,help -n "$0" -- "$@") || \
	usage_error

eval set -- "$OPTIONS"

while true; do
	case "$1" in
	-l|--list)
		do_list="y"
		shift
		;;
	-u|--do-upgrade)
		do_upgrade="y"
		shift
		;;
	-h|--help)
		usage
		exit 0
		;;
	--)
		shift
		break
		;;
	esac
done

if [ "$#" -gt 1 ]; then
	echo "ERROR: invalid number of parameters"
	usage_error
fi
if [ -n "$*" ]; then
	release="$*"
fi
check_pre_req
download_rel_meta_file
parse_rel_meta_file
if [ "${do_list}" = "y" ]; then
	list_available_releases
	exit 0
fi

update_source_list "${release}"

if [ "${do_upgrade}" = "y" ]; then
	do_upgrade
fi
