#!/bin/bash
#
# Copyright (c) 2016-2018 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# See file LICENSE for terms.
#

BASE_NAME=$(basename $0)
BASE_DIR=$(dirname $0)
SYSTEM_UNIT_FILES_DIR=${SYSTEM_UNIT_FILES_DIR:-"/etc/systemd/system"}

. $BASE_DIR/sharp_funcs.sh


# Get the linux disto name (it returns systemd in case systemd is supported)
distro_name=$(get_linux_distro)

start_serv()
{
	local serv=$1

	case $distro_name in
	'systemd')
		systemctl --quiet start $serv > /dev/null
	;;
	*)         service $serv start  &> /dev/null      ;;
	esac
}

#This function creates a conf file that overrides settings in
#the original service file. Multiple settings can be send to the function
#in the following format:
#update_unit Section key value key value ....
function update_unit()
{
	local serv=$1
	local section=$2
	shift
	shift
	local drop_folder=${SYSTEM_UNIT_FILES_DIR}/${serv}.service.d
	local drop_file=${drop_folder}/${section}.conf


	mkdir -p ${drop_folder}
	touch ${drop_file}
	#Create section title
	cat > ${drop_file} << EOF
[$section]
EOF
	#Loop over Key Value arguments
	while [ $# -gt 0 ]; do
		key=$1
		shift
		val=$1
		shift
		cat >> ${drop_file} << EOF
$key=$val
EOF
	done

	systemctl daemon-reload > /dev/null
}

add_serv_systemd()
{
	local serv=$1
	local location_dir=$2

	local unit_base_name=${serv}.service
	local unit_source_path=${location_dir}/systemd/system/${unit_base_name}
	local unit_path=${SYSTEM_UNIT_FILES_DIR}/${unit_base_name}

	if ! [[ -f ${unit_source_path} ]]; then
		echo "Error: $unit_source_path doesn't exist."
		echo "Cannot install ${serv} as systemd service. Exit."
		exit 3
	fi

	local conf_dir="${SYSTEM_UNIT_FILES_DIR}/"$unit_base_name".d"

	# Copy service file to systemd files directory
	echo "Copying ${unit_source_path} to ${unit_path}"
	\cp -rf ${unit_source_path} ${unit_path}
	if ! [[ -f ${unit_path} ]]; then
		echo "Error: $unit_path doesn't exist."
		echo "Cannot install ${serv} as systemd service. Exit."
		exit 3
	fi

	# Enable service
	systemctl --quiet enable ${unit_path} > /dev/null 2>&1

	if [[ ${serv} = "sharp_am" ]]; then
		update_unit $serv Service Environment "CONF=-${location_dir}/conf/${serv}.cfg"  ExecStart "" ExecStart "${location_dir}/bin/${serv} -O \$CONF"
	fi
}

add_serv()
{
	local serv=$1
	local location_dir=$2
	local init=$3

	case $distro_name in
	'systemd')
		add_serv_systemd ${serv} ${location_dir}
		;;
	'ubuntu')
		chmod 755 $init
		ln -sf $init /etc/init.d/$serv &>/dev/null
		update-rc.d $serv defaults &> /dev/null
		;;
	*)
		chmod 755 $init
		ln -sf $init /etc/init.d/$serv &>/dev/null
		chkconfig --add $serv &> /dev/null
		;;
	esac
}

rm_serv()
{
	serv=$1

	case $distro_name in
	'systemd')
		local unit_base_name=${serv}.service
                local conf_dir="${SYSTEM_UNIT_FILES_DIR}/"$unit_base_name".d"

		systemctl disable ${serv}.service &> /dev/null

		[[ -f "${SYSTEM_UNIT_FILES_DIR}/${unit_base_name}" ]] && \rm -f ${SYSTEM_UNIT_FILES_DIR}/${unit_base_name} &> /dev/null

		#if the service.d directory is empty, remove it.
		if [[ -d ${conf_dir} ]]; then
			if ! find "$conf_dir" -mindepth 1 -print -quit | grep -q .; then
				\rm -df $conf_dir &> /dev/null
				echo "removed $conf_dir"
			fi
		fi

		systemctl daemon-reload
	 ;;
	'ubuntu') (ls /etc/rc*.d | grep -q $serv) &> /dev/null && update-rc.d -f $serv remove &> /dev/null ;;
	*)        (chkconfig --list | grep $serv) &> /dev/null && chkconfig --del $serv &> /dev/null       ;;
	esac
}

unset_level()
{
	serv=$1

	case $distro_name in
	'systemd') ;;
	'ubuntu') update-rc.d $serv disable 0123456 &> /dev/null   ;;
	*)        chkconfig --level 0123456 $serv off &> /dev/null ;;
	esac
}

is_added()
{
	serv=$1
	ok=1

	case $distro_name in
	'systemd') systemctl list-unit-files --type=service | grep -q $serv && ok=0 ;;
	'ubuntu') ls /etc/rc*.d | grep -q $serv && ok=0                ;;
	*)        (chkconfig --list | grep $serv) &> /dev/null && ok=0 ;;
	esac

	if [ $ok -eq 0 ]; then
		echo "Service $serv is installed"
		return 0
	else
		echo "Error: failed to install service $serv"
		return 1
	fi
}

is_removed()
{
	serv=$1
	ok=1

	case $distro_name in
		'systemd') ! (systemctl list-unit-files --type=service | grep -q $serv) && ok=0 ;;
		'ubuntu') ! (ls /etc/rc*.d | grep -q $serv) && ok=0              ;;
	*)        ! (chkconfig --list | grep $serv) &> /dev/null && ok=0 ;;
	esac

	if [ $ok -eq 0 ]; then
		echo "Service $serv is removed"
		return 0
	else
		echo "Error: failed to remove service $serv"
		return 1
	fi
}

# $1 - SHARP location dir
# $2 - list of daemons
setup()
{
	local location_dir=$1
	local init=${location_dir}/sbin/sharp.init

	if [ ! -f "$init" ]; then
		echo "Error: $init doesn't exist. Exit."
		exit 3
	fi

	for daemon in $2; do
		add_serv $daemon ${location_dir} ${init}

		if [[ ${daemon} = "sharp_am" ]]; then
			unset_level sharp_am
		else
			start_serv ${daemon}
		fi

		is_added $daemon
	done
}

# $1 - list of daemons
unsetup()
{
	for daemon in $1; do
		[ -x /etc/init.d/$daemon ] && /etc/init.d/$daemon stop
		pkill $daemon &>/dev/null
		rm -f /tmp/d_${daemon}.log /var/run/{$daemon}.pid &> /dev/null

		rm -f /etc/init.d/$daemon &> /dev/null
		rm_serv $daemon

		is_removed $daemon
	done
}

usage()
{
	echo "Usage: `basename $0` (-s | -r) [-p SHARP location dir]"
	echo "   or: `basename $0` -h"
	echo "$BASE_NAME helps system administrator to manage SHARP daemons."
	echo "All management operations require root access."
	echo ""
	echo "Operation flags (must specify one for setup/remove):"
	echo "	-s - Setup SHARP daemon (sharp_am)"
	echo "	-r - Remove SHARP daemon (sharp_am)"
	echo "	-h - Print this help and exit"
	echo ""
	echo "Additional options:"
	echo "	-p - Path to SHARP source installation directory (default: auto-detected)"
	echo ""
	echo "Examples:"
	echo "	Show this help message:"
	echo "		$0 -h"
	echo "	Setup SHARP daemon with default path:"
	echo "		$0 -s"
	echo "	Setup SHARP daemon with custom path:"
	echo "		$0 -s -p /path/to/sharp"
	echo "	Remove SHARP daemon:"
	echo "		$0 -r"
	echo ""
	echo "Note: sharp_am.service is NOT started automatically and does NOT support by default auto-restart."
	echo "Starting sharp_am.service resets SHARP trees and cleans allocated resources."
	echo "It could affect running applications."

	exit 2
}

# Hardcode sharp_am as the only daemon
dlist="sharp_am"
location_dir=""
to_setup=""
to_remove=""
while getopts "sp:rh" flag_arg; do
	case $flag_arg in
	s) to_setup="yes"         ;;
	r) to_remove="yes"        ;;
	p) location_dir="$OPTARG" ;;
	h) usage                  ;;
	*) usage                  ;;
	esac
done

if [ $(whoami) != "root" ]; then
	echo "Error: root permission required. Exit."
	echo ""
	usage
	exit 1
fi

[ $OPTIND -eq 1 ] && usage

[[ -z ${to_setup} && -z ${to_remove} ]] && usage 
[[ -n ${to_setup} && -n ${to_remove} ]] && usage 

[ -z "$location_dir" ] && _dir=$(readlink -f $(dirname $0)/..)

[ -z "$location_dir" ] && [ -f "$_dir/bin/sharp_am" ] && location_dir=$_dir

[ -n "$to_setup" ] && setup "$location_dir" "$dlist"
[ -n "$to_remove" ] && unsetup "$dlist"

exit 0
