#!/bin/bash
#
# Copyright 2016-2017 Advanced Micro Devices, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.

set -e
shopt -s nullglob

REPOSITORY="/var/opt/amdgpu-pro-local"

BASE_PACKAGE=amdgpu-core
META_PACKAGE=amdgpu
META_LIB_PACKAGE=amdgpu-lib
DKMS_PACKAGE=amdgpu-dkms
VERSION_REQUIRED=false
VERSIONLIST_PACKAGE=amdgpu-versionlist
VERSIONLIST_PRO_PACKAGE=amdgpu-pro-versionlist
OPENGL_META_PACKAGE=amdgpu-pro
OPENCL_LEGACY_META_PACKAGES=(clinfo-amdgpu-pro opencl-orca-amdgpu-pro-icd)
OPENCL_ROCR_META_PACKAGES=(amdgpu-pro-rocr-opencl)
VULKAN_META_PACKAGE=vulkan-amdgpu-pro
PX_PACKAGE=xorg-x11-drv-modesetting
WARN_KERNEL_SOURCES=false

usage() {
	cat <<END_USAGE
Usage: $PROG [options...]

Options:
  -h|--help                Display this help message
  --dryrun                 Print list of packages to install and exit
  --px                     (DEPRECATED) PX platform support
  --version=VERSION        Install the specified driver VERSION
  --pro                    Install "pro" support (legacy OpenGL and Vulkan)
  --opencl=legacy          Install legacy OpenCL support
  --opencl=rocr            Install ROCr OpenCL support
  --opencl=legacy,rocr     Install both legacy and ROCr OpenCL support
  --headless               Headless installation (only OpenCL support)
  --no-dkms                Do not install amdgpu-dkms package
  --compute                (DEPRECATED) Equal to --opencl=legacy --headless
  --uninstall		   Uninstall the amdgpu driver

  Unless the -h|--help option is given, '$DNF' options may be present.

  Unless headless installation is requested, pro support will be installed.

END_USAGE
}

function stderr() {
	cat - 1>&2
}

function check_kernel_headers() {
	case "$OS_CLASS" in
	"fedora")
		if ! rpm --quiet -q kernel-devel-$(uname -r); then
			return 1;
		fi
		;;
	esac
	return 0;
}

function check_for_absence_of_nomodeset() {
	if grep -q nomodeset "/proc/cmdline"; then
		echo "WARNING: nomodeset detected in kernel parameters, "`
			`"amdgpu requires KMS" | stderr
	fi
}

function check_dkms_succeeded_for_running_kernel() {
	if [ "$DKMS_PACKAGE" != "" ] && ! /usr/sbin/dkms status amdgpu | \
			grep `uname -r` | grep -q installed; then
		echo "WARNING: amdgpu dkms failed for running kernel" | stderr
		if ! check_kernel_headers; then
			echo "WARNING: kernel sources for running kernel "`
				`"may not be installed" | stderr
		fi
	fi
}

function check_install() {
	check_for_absence_of_nomodeset
	check_dkms_succeeded_for_running_kernel
}

function check_driver_install_status(){
        if rpm -q "$BASE_PACKAGE" &>/dev/null || \
                rpm -qa | grep -q "$DKMS_PACKAGE" &>/dev/null || \
                rpm -q "$VERSIONLIST_PACKAGE" &>/dev/null || \
                rpm -q "$VERSIONLIST_PRO_PACKAGE" &>/dev/null || \
                [ -e /usr/bin/amdgpu-uninstall ] || \
                [ -e /usr/bin/amdgpu-pro-uninstall ]; then
                return 0
        else
                return 1
        fi
}

function check_previous_install() {
	if check_driver_install_status; then
		echo -e "The amdgpu driver is already installed or was not "`
			`"uninstalled correctly.\nPlease fully uninstall "`
			`"the driver before proceeding with "`
			`"installation" | stderr
		if [ -e /usr/bin/amdgpu-uninstall ] || \
			[ -e /usr/bin/amdgpu-pro-uninstall ]; then
			echo -e "You can try running the amdgpu-uninstall or "`
				`"amdgpu-pro-uninstall script \npresent in "`
				`"/usr/bin to clean up the previous "`
				`"installation" | stderr
		else
			echo -e "You can try running this install script with "`
				`"the --uninstall option \n to clean up the "`
				`"previous installation" | stderr
		fi
		exit 1
	fi
}

function check_options() {
	if [[ "${OPTIONS[*]}" =~ "no-dkms" ]] && \
		[[ "${OPTIONS[*]}" =~ "pro" ]]; then
		echo -e "ERROR: using --no-dkms with a pro install is "`
			`"invalid.\nDid you want to run the following`
			` instead?\namdgpu-install --no-dkms" | stderr
		exit 1
	fi
}

function os_release() {
	if [[ -r /etc/os-release ]]; then
		. /etc/os-release

		case "$ID" in
		fedora)
			DNF=dnf
			OS_CLASS=fedora
			OS_PKG=.fc${VERSION_ID}
			RMPKG=erase
			;;
		rhel|centos)
			DNF=yum
			OS_CLASS=fedora
			OS_PKG=.el${VERSION_ID%%.*}
			RMPKG=erase
			;;
		sles|sled|opensuse*)
			DNF=zypper
			OS_CLASS=opensuse
			OS_PKG=""
			RMPKG=remove
			;;
		amzn)
			DNF=yum
			OS_CLASS=fedora
			OS_PKG=.el7
			RMPKG=erase
			;;
		*)
			echo "Unsupported RPM-based OS: `
				`/etc/os-release ID '$ID'" | stderr
			exit 1
			;;
		esac

	elif [[ -f /etc/redhat-release ]]; then
		ID=rhel
		DNF=yum
		OS_CLASS=fedora
		OS_PKG=.el6
		RMPKG=erase
	else
		echo "Unsupported OS" | stderr
		exit 1
	fi
}

function yum_repo() {
	local repo=${REPOSITORY##*/}

	if [[ "$OS_CLASS" = "opensuse" ]]; then
		echo "/etc/zypp/repos.d/$repo.repo"
	else
		echo "/etc/yum.repos.d/$repo.repo"
	fi
}

function fedora_remove_repo() {
	if [[ -d $REPOSITORY ]]; then
		local repo=${REPOSITORY##*/}
		$SUDO $DNF clean metadata --disablerepo=* --enablerepo=$repo
	fi
	$SUDO rm -rf $BIN/${PROG%-*}-uninstall $(yum_repo) $REPOSITORY
}

function opensuse_remove_repo() {
	$SUDO rm -rf $BIN/${PROG%-*}-uninstall $(yum_repo) $REPOSITORY
	$SUDO $DNF refresh
}

function install_version_list() {
	local build_id
	if [[ "$DNF" == "zypper" ]]; then
		build_id="-"$(zypper se -s '*-versionlist' \
			| grep -m 1 -o "[0-9.]*-$VER" || echo $VER)
	else
		build_id=${VER#*-}
		build_id=${build_id:+-*-$build_id$OS_PKG}
	fi

	if [[ "${OPTIONS[*]}" =~ "pro" ]] || ! $SUDO $DNF ${1+"$@"} \
		install -y "$VERSIONLIST_PACKAGE$build_id"; then
		$SUDO $DNF ${1+"$@"} install -y \
			"$VERSIONLIST_PRO_PACKAGE$build_id"
	fi

	local list="$(cat /usr/share/amdgpu*/*-versionlist | "`
		`"awk -F'\"' '/AMDGPU_VERSION_LIST_/ {printf $2}')"

	for i in ${!PACKAGES[@]}; do
		for pkg in $list; do
			if [[ "${pkg%%=*}" == "${PACKAGES[$i]}" ]]; then
				if [[ "$DNF" == "yum" ]]; then
                                        pkg="${pkg/=*:/=}"
                                fi
				PACKAGES[$i]="${pkg/=/-}"
				break
			fi
		done
	done
}

function local_repo_version() {
	local index1="$PROG_DIR/repodata"
	local index2="$PROG_DIR/RPMS"
	local sentinel=($PROG_DIR/RPMS/noarch/$VERSIONLIST_PACKAGE-*$OS_PKG.noarch.rpm)

	if [[ "${OPTIONS[*]}" =~ "pro" ]]; then
		sentinel=($PROG_DIR/RPMS/noarch/$VERSIONLIST_PRO_PACKAGE-*$OS_PKG.noarch.rpm)
	fi

	if [[ -r "$index1" && -r "$index2" && ${#sentinel[*]} -gt 0 ]]; then
		VER=$(rpm -qp $sentinel --qf "%{VERSION}-%{RELEASE}" | \
			sed "s/${OS_PKG}$//")
		OPTIONS+=(version)
		SCRIPT_DIR=$REPOSITORY
		$SUDO mkdir -p $REPOSITORY && \
			$SUDO cp -arf "$PROG_DIR"/* $REPOSITORY
		echo -e `\
			`"[${REPOSITORY##*/}]\n"`
			`"name=AMD amdgpu Pro local repository\n"`
			`"baseurl=file://$REPOSITORY\n"`
			`"enabled=1\n"`
			`"gpgcheck=0\n" \
		| $SUDO tee $(yum_repo)
	fi
}

function build_package_list() {
	local opencl_meta_packages=()
	local pro_meta_packages=()

	if [[ "${OPTIONS[*]}" =~ "no-dkms" ]]; then
		DKMS_PACKAGE=
		META_PACKAGE=$META_LIB_PACKAGE
	fi

	[[ "${OPTIONS[*]}" =~ "pro" ]] \
		&& pro_meta_packages=($OPENGL_META_PACKAGE $VULKAN_META_PACKAGE)

	if [[ "${OPTIONS[*]}" =~ "opencl" ]]; then
		[[ "${OPENCL_OPTIONS[*]}" =~ "legacy" ]] \
			&& opencl_meta_packages=(${OPENCL_LEGACY_META_PACKAGES[*]})

		if [[ "${OPENCL_OPTIONS[*]}" =~ "rocr" ]]; then
			opencl_meta_packages=(${opencl_meta_packages[*]} \
				${OPENCL_ROCR_META_PACKAGES[*]})
		fi

		if [[ -z "${opencl_meta_packages[*]}" ]]; then
			echo "ERROR: Unknown opencl option '${OPENCL_OPTIONS[*]}'"
			exit 1
		fi

		if [[ "${OPTIONS[*]}" =~ "headless" ]]; then
			PACKAGES=($DKMS_PACKAGE ${opencl_meta_packages[*]})
		else
			PACKAGES=($DKMS_PACKAGE $META_PACKAGE \
				${pro_meta_packages[*]} \
				${opencl_meta_packages[*]})
		fi
	else
		PACKAGES=($DKMS_PACKAGE $META_PACKAGE ${pro_meta_packages[*]})
	fi

	[[ "${OPTIONS[*]}" =~ "px" ]] \
		&& PACKAGES=(${PACKAGES[*]} $PX_PACKAGE)

	if ! check_kernel_headers; then
		PACKAGES=(${PACKAGES[*]} kernel-devel-$(uname -r))
	fi

	return 0
}

function amdgpu_pro_install() {
	check_previous_install
	check_options
	build_package_list

	if [[ "${OPTIONS[*]}" =~ "dryrun" ]]; then
		echo PACKAGES: ${PACKAGES[*]}
		amdgpu_pro_uninstall
		return 0
	fi

	if [[ "${OPTIONS[*]}" =~ "version" ]] || ! $VERSION_REQUIRED; then
		install_version_list ${1+"$@"}
	else
		echo "ERROR: argument '--version' must be set" | stderr
		return 1
	fi

	$SUDO ln -sf $SCRIPT_DIR/$PROG $BIN/${PROG%-*}-uninstall

	$SUDO $DNF ${1+"$@"} install ${PACKAGES[*]}

	check_install
}

function amdgpu_pro_uninstall() {
	local list=""
	rpm --quiet -q "$BASE_PACKAGE" && list="$list $BASE_PACKAGE"
	list="$list$(rpm -qa \
		| awk -F'-[0-9]' /$DKMS_PACKAGE/' { printf " %s",$1 }')"
	rpm --quiet -q "$VERSIONLIST_PACKAGE" && \
		list="$list $VERSIONLIST_PACKAGE"
	rpm --quiet -q "$VERSIONLIST_PRO_PACKAGE" && \
		list="$list $VERSIONLIST_PRO_PACKAGE"
	[[ "$list" = "" ]] || $SUDO $DNF ${1+"$@"} $RMPKG $list
	${OS_CLASS}_remove_repo
}

PROG=${0##*/}
PROG_DIR=$(cd ${0%/*} && pwd -P)
SUDO=$([[ $(id -u) -ne 0 ]] && echo "sudo" ||:)
BIN="/usr/bin"
SCRIPT_DIR=$BIN
VER=""
OPTIONS=()
OPENCL_OPTIONS=()

os_release

[[ "$PROG" =~ "pro" ]] && OPTIONS+=(pro)

while (($#))
do
	case "$1" in
	-h|--help)
		usage
		exit 0
		;;
	--px)
		echo "WARNING: --px option is deprecated\n"`
			`"All-Open with PRIME is recommended for for "`
			`"all mobile variants with hybrid graphics"
		OPTIONS+=(${1#--})
		shift
		;;
	--headless|--pro|--dryrun|--no-dkms)
		OPTIONS+=(${1#--})
		shift
		;;
	--version*)
		OPTIONS+=(version)
		VER=${1#--version=}
		shift
		;;
	--opencl*)
		OPTIONS+=(opencl)
		OPENCL_OPTIONS=${1#--opencl=}
		OPENCL_OPTIONS=(${OPENCL_OPTIONS/,/ })
		shift
		;;
	--compute)
		echo "WARNING: --compute option is deprecated"
		OPTIONS+=(opencl headless)
		OPENCL_OPTIONS=(legacy)
		shift
		;;
	--uninstall)
		amdgpu_pro_uninstall
                if check_driver_install_status; then
                        echo -e "Errors were encountered when attempting to "`
                                `" uninstall the driver" | stderr
                else
                        echo -e "Uninstall was completed sucessfully" | stderr
                fi
                exit
                ;;
	*)
		ARGS+="$1 "
		shift
		;;
	esac
done

local_repo_version

set -- $ARGS
amdgpu_pro_${0##*-} ${1+"$@"}
