#!/bin/bash
#
# Copyright 2023 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

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

Options:
  -h|--help             Display this help message
  -c|--clean            Revert repositories and settings to default
  -b|--baseurl URL      Set amdgpu/rocm base package repository URL,
                        e.g. "https://repo.radeon.com"
  -r|--release ID       Set amdgpu/rocm release repository, e.g. "X.Y.Z"
  -a|--add-rocm         Add a rocm release repository instead of changing
                        configured repositories.

Note: When setting the baseurl or release, the values will be saved for next
      use, excluding when the '--add-rocm' option is used. Running 'clean' will
      revert all repositories and delete saved configuration.
END_USAGE
}

die() {
	(( $# == 0 )) || echo "ERROR: $1" >&2
	exit 1
}

die_if_missing() {
	[[ -f "$1" ]] || die "$1 is missing."
}

os_release() {
	local os_id
	os_id=$(sh -c '. /etc/os-release && echo $ID' 2>/dev/null)
	case "$os_id" in
	ubuntu|linuxmint|debian)
		PKGMAN=apt
		REPODIR=/etc/apt/sources.list.d/
		REPOEXT=list
		;;
	fedora|rhel|centos|almalinux|rocky|amzn)
		PKGMAN=yum
		REPODIR=/etc/yum.repos.d/
		REPOEXT=repo
		;;
	sles|sled|opensuse*)
		PKGMAN=zypper
		REPODIR=/etc/zypp/repos.d/
		REPOEXT=repo
		;;
	*)
		die "Unsupported OS"
		;;
	esac
}

apt_refresh() {
	$PKGMAN update \
		-o Dir::Etc::sourcelist="sources.list.d/$1.list" \
		-o Dir::Etc::sourceparts="-" -o APT::Get::List-Cleanup="0"
}

yum_refresh() {
	$PKGMAN clean metadata --disablerepo="*" --enablerepo="$1"
}

zypper_refresh() {
	$PKGMAN refresh -r "$1"
}

add_repo() {
	local repo="$1"
	local release="$2"
	local baseurl="$3"

	die_if_missing "$DATA_DIR/${repo%-*}.$REPOEXT.template"
	sed -e "s|@BASEURL@|$baseurl|" -e "s|@VER@|${release%.0}|" \
		-e "s|@REPO@|$repo|" \
		"$DATA_DIR/${repo%-*}.$REPOEXT.template" \
		> "$REPODIR/$repo.$REPOEXT"
	${PKGMAN}_refresh "$repo" || \
		die "repo refresh failed, use '$PROG -c' to revert to default"
}

revert_repos() {
	die_if_missing "$DATA_DIR/orig/amdgpu.$REPOEXT"
	die_if_missing "$DATA_DIR/orig/rocm.$REPOEXT"
	rm -f "$REPODIR/rocm-"*".$REPOEXT"
	rm -r "$CONFIG_DIR/amdgpu-setup.conf"
	cp -af "$DATA_DIR/orig/"{amdgpu,rocm}".$REPOEXT" "$REPODIR"
	${PKGMAN}_refresh "amdgpu"
	${PKGMAN}_refresh "rocm"
}

save_configuration() {
	local name="$1"
	local value="$2"

	if [[ ! -d "$CONFIG_DIR" ]]; then
		mkdir -p "$CONFIG_DIR"
	elif [[ -f "$CONFIG_DIR/$CONFIG_FILENAME" ]]; then
		sed -i "/^$name=/d" "$CONFIG_DIR/$CONFIG_FILENAME"
	fi
	echo "$name=$value" >> "$CONFIG_DIR/$CONFIG_FILENAME"
}

PROG=${0##*/}

if [[ $# -eq 0 ]]; then
    usage
    exit 0
fi

OPTIONS=()
CONFIG_DIR="$HOME/.config/amdgpu-install"
CONFIG_FILENAME="amdgpu-setup.conf"
DATA_DIR="/usr/share/amdgpu-install"

if [[ -f "/etc/amdgpu-install/$CONFIG_FILENAME" ]]; then
	. /etc/amdgpu-install/$CONFIG_FILENAME
fi
if [[ -f "$CONFIG_DIR/$CONFIG_FILENAME" ]]; then
	. "$CONFIG_DIR/$CONFIG_FILENAME"
fi

LONGOPTS="help,clean,baseurl:,release:,add-rocm"
OPTS=$(getopt -o "h,c,b:,r:,a" -l "$LONGOPTS" -- "$@")
eval set -- "$OPTS"

os_release

while true; do
	case "$1" in
	-h|--help)
		usage
		exit 0
		;;
	-c|--clean)
		revert_repos
		exit 0
		;;
	-b|--baseurl)
		shift
		OPTIONS+=("baseurl")
		BASEURL="$1"
		;;
	-r|--release)
		shift
		OPTIONS+=("release")
		if [[ $1 =~ ^[0-9]+[.][0-9]+[.][0-9]+$ ]]; then
			RELEASE="$1"
		else
			die "Invalid release format '$1'"
		fi
		;;
	-a|--add-rocm)
		OPTIONS+=("add-rocm")
		;;
	--)
		shift
		break
		;;
	*)
		usage >&2
		exit 1
		;;
	esac
	shift
done

if [[ $(id -u) -ne 0 ]]; then
	die "Please run as root"
fi

if [[ -z "$BASEURL" ]]; then
	die "Unknown BASEURL value, use '-b' to set"
fi
if [[ -z "$RELEASE" ]]; then
	die "Unknown RELEASE value, use '-r' to set"
fi

if [[ ${OPTIONS[*]} =~ add-rocm ]]; then
	add_repo "rocm-$RELEASE" "$RELEASE" "$BASEURL"
else
	if [[ ${OPTIONS[*]} =~ baseurl ]]; then
		save_configuration "BASEURL" "$BASEURL"
	fi
	if [[ ${OPTIONS[*]} =~ release ]]; then
		save_configuration "RELEASE" "$RELEASE"
	fi
	add_repo "amdgpu" "$RELEASE" "$BASEURL"
	add_repo "rocm" "$RELEASE" "$BASEURL"
fi
echo "INFO: If you run into issues, run '$PROG -c' to revert back to default"
