#!/bin/bash

# Copyright 2019-2023 NVIDIA CORPORATION

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
SAVED_SETTINGS="$-"
set -eu

if echo ${PM_VERBOSITY-} | grep -i "debug" > /dev/null ; then
	set -x
	PM_CURL_SILENT=""
	PM_WGET_QUIET=""
else
	PM_CURL_SILENT="-s -S"
	PM_WGET_QUIET="--quiet"
fi
export PM_PACKMAN_VERSION=7.23.1

# This is necessary for newer macOS
if [ `uname` == 'Darwin' ]; then
	export LC_ALL=en_US.UTF-8
	export LANG=en_US.UTF-8
fi

# We cannot rely on realpath, it isn't installed on macOS and some Linux distros
get_abs_filename() {
  echo "$(cd "$(dirname "$1")" && pwd)/$(basename "$1")"
}

# Specify where packman command exists
export PM_INSTALL_PATH="$(get_abs_filename "$(dirname "${BASH_SOURCE}")")"

# The packages root may already be configured by the user
if [ -z "${PM_PACKAGES_ROOT:-}" ]; then
	# Set variable temporarily in this process so that the following execution will work
	if [ `uname` == 'Darwin' ]; then
		export PM_PACKAGES_ROOT="${HOME}/Library/Application Support/packman-cache"
	else
		if [ -z "${XDG_CACHE_HOME:-}" ]; then
			export PM_PACKAGES_ROOT="${HOME}/.cache/packman"
		else
			export PM_PACKAGES_ROOT="${XDG_CACHE_HOME}/packman"
		fi
	fi
fi

# Ensure the packages root path exists:
if [ ! -d "$PM_PACKAGES_ROOT" ]; then
	echo "Creating packman packages cache at $PM_PACKAGES_ROOT"
	mkdir -p -m a+rwx "$PM_PACKAGES_ROOT"
fi

fetch_file_from_s3()
{
	SOURCE=$1
	SOURCE_URL=http://bootstrap.packman.nvidia.com/$SOURCE
	TARGET=$2
	echo "Fetching $SOURCE from bootstrap.packman.nvidia.com ..."
	if command -v wget >/dev/null 2>&1; then
		wget $PM_WGET_QUIET -O$TARGET $SOURCE_URL
	else
		curl -o $TARGET $SOURCE_URL $PM_CURL_SILENT
	fi
}

generate_temp_file_name()
{
	if [ `uname` == "Darwin" ]; then
		local tmpfile=`mktemp -t packman`
	else
		local tmpfile=`mktemp -t packman.XXXXXXXX`
	fi
	echo "$tmpfile"
}

install_python()
{
	PLATFORM=`uname`
	PROCESSOR=`uname -m`
	PYTHON_VERSION=3.10.5-1

	if [ $PLATFORM == 'Darwin' ]; then
		PYTHON_PACKAGE=$PYTHON_VERSION-macos-x86_64
	elif [ $PLATFORM == 'Linux' ] && [ $PROCESSOR == 'x86_64' ]; then
		PYTHON_PACKAGE=$PYTHON_VERSION-linux-x86_64
	elif [ $PLATFORM == 'Linux' ] && [ $PROCESSOR == 'aarch64' ]; then
		PYTHON_PACKAGE=$PYTHON_VERSION-linux-aarch64
	else
		echo "Operating system not supported"
		exit 1
	fi

	PYTHON_INSTALL_FOLDER="$PM_PACKAGES_ROOT/python/$PYTHON_PACKAGE"
	if [ ! -d "$PYTHON_INSTALL_FOLDER" ]; then
		mkdir -p "$PYTHON_INSTALL_FOLDER"
	fi

	export PM_PYTHON="$PYTHON_INSTALL_FOLDER/python"

	if [ ! -f "$PM_PYTHON" ]; then
		PYTHON_PACKAGE_TMP=$(generate_temp_file_name)
		fetch_file_from_s3 "python@$PYTHON_PACKAGE.tar.gz" "$PYTHON_PACKAGE_TMP"
		if [ "$?" -eq "0" ]; then
			echo "Unpacking python"
			tar -xf "$PYTHON_PACKAGE_TMP" -C "$PYTHON_INSTALL_FOLDER"
			rm "$PYTHON_PACKAGE_TMP"
		else
			echo "Failed downloading the Python interpreter"
			exit $?
		fi
	fi
}

# Ensure python is available:
if [ -z "${PM_PYTHON_EXT:-}" ]; then
	install_python
else
	PM_PYTHON="$PM_PYTHON_EXT"
fi

# The packman module may be externally configured
if [ -z "${PM_MODULE_DIR_EXT:-}" ]; then
	PM_MODULE_DIR="$PM_PACKAGES_ROOT/packman-common/$PM_PACKMAN_VERSION"
else
    PM_MODULE_DIR="$PM_MODULE_DIR_EXT"
fi
export PM_MODULE="$PM_MODULE_DIR/run.py"

# Ensure the packman package exists:
if [ ! -f "$PM_MODULE" ]; then
	# Remove a previously corrupt packman-common if it's there
	if [ -d "$PM_MODULE_DIR" ]; then
		rm -rf "$PM_MODULE_DIR"
	fi
	PM_MODULE_PACKAGE="packman-common@$PM_PACKMAN_VERSION.zip"
	TARGET=$(generate_temp_file_name)
	# We always fetch packman from S3:
	fetch_file_from_s3 "$PM_MODULE_PACKAGE" "$TARGET"
	if [ "$?" -eq "0" ]; then
		echo "Unpacking ..."
		"$PM_PYTHON" -S -s -u -E "$PM_INSTALL_PATH/bootstrap/install_package.py" "$TARGET" "$PM_MODULE_DIR"
		rm "$TARGET"
	else
		echo "Failure while fetching packman module from S3!"
		exit 1
	fi
fi

# Generate temporary file name for environment variables:
PM_VAR_PATH=`mktemp -u -t tmp.$$.pmvars.XXXXXX`

if [ $# -ne 0 ]
  then
    PM_VAR_PATH_ARG=--var-path="$PM_VAR_PATH"
fi

"$PM_PYTHON" -S -s -u -E "$PM_MODULE" "$@" ${PM_VAR_PATH_ARG:-}
exit_code=$?
# Export the variables if the file was used and remove the file:
if [ -f "$PM_VAR_PATH" ]; then
	while read -r line
	do
        if [ ${#line} -gt 0 ]; then
    		export "$line"
        fi
	done < "$PM_VAR_PATH"
    rm -f "$PM_VAR_PATH"
fi

# avoid leaking -e and -u into the host script if they weren't originally set
if [[ ! ( "$SAVED_SETTINGS" =~ e ) ]]; then
    set +e
fi

if [[ ! ( "$SAVED_SETTINGS" =~ u ) ]]; then
    set +u
fi

# Return the exit code from python
if [ "$exit_code" != 0 ]; then
    exit "$exit_code"
fi