#!/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
VERSIONLIST_PACKAGE=amdgpu-pro-versionlist
OPENGL_META_PACKAGE=amdgpu-pro
OPENCL_LEGACY_META_PACKAGES=(clinfo-amdgpu-pro opencl-orca-amdgpu-pro-icd)
OPENCL_PAL_META_PACKAGES=(clinfo-amdgpu-pro opencl-amdgpu-pro-icd)
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=pal             Install PAL OpenCL support
  --opencl=legacy,pal      Install both legacy and PAL OpenCL support
  --headless               Headless installation (only OpenCL support)
  --no-dkms                Do not install amdgpu-dkms package
  --compute                (DEPRECATED) Equal to --opencl=legacy --headless

  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_and_install_kernel_headers() {
	case "$OS_CLASS" in
	"fedora")
		if ! rpm --quiet -q kernel-devel-$(uname -r); then
			if ! $SUDO $DNF ${1+"$@"} install \
					kernel-devel-$(uname -r); then
				WARN_KERNEL_SOURCES=true
			fi
		fi
		;;
	esac
}

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 $WARN_KERNEL_SOURCES; 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 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
			;;
		*)
			echo "Unsupported RedHat derivative OS" | 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 [[ "$ID" = "sle"[s,d] ]] || [[ "$ID" = "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() {
	$SUDO $DNF ${1+"$@"} install "$VERSIONLIST_PACKAGE-$VER$OS_PKG"
	. /usr/share/amdgpu-pro/amdgpu-versionlist

	for i in ${!PACKAGES[@]}; do
		for pkg in ${AMDGPU_VERSION_LIST}; do
			if [[ "${pkg%%=*}" == "${PACKAGES[$i]}" ]]; then
				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/$BASE_PACKAGE-*$OS_PKG.noarch.rpm)

	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=()

	[[ "${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[*]})

		[[ "${OPENCL_OPTIONS[*]}" =~ "pal" ]] \
			&& opencl_meta_packages=(${opencl_meta_packages[*]} \
				${OPENCL_PAL_META_PACKAGES[*]})

		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)

	return 0
}

function amdgpu_pro_install() {
	$SUDO ln -sf $SCRIPT_DIR/$PROG $BIN/${PROG%-*}-uninstall

	build_package_list

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

	if [[ "${OPTIONS[*]}" =~ "version" ]]; then
		install_version_list ${1+"$@"}
	elif rpm --quiet -q "$VERSIONLIST_PACKAGE" ; then
		$SUDO $DNF ${1+"$@"} $RMPKG $VERSIONLIST_PACKAGE
	fi

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

	check_install
}

function amdgpu_pro_uninstall() {
	rpm --quiet -q "$BASE_PACKAGE" && \
		$SUDO $DNF ${1+"$@"} $RMPKG $BASE_PACKAGE
	rpm --quiet -q "$VERSIONLIST_PACKAGE" && \
		$SUDO $DNF ${1+"$@"} $RMPKG $VERSIONLIST_PACKAGE
	${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
local_repo_version

[[ "$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)
		OPTIONS+=(${1#--})
		shift
		;;
	--no-dkms)
		DKMS_PACKAGE=
		META_PACKAGE=$META_LIB_PACKAGE
		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
		;;
	*)
		ARGS+="$1 "
		shift
		;;
	esac
done

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