#!/bin/bash

#    ltask: lightweight remote task execution
#    Copyright (C) 2020 Jonas Gunz <himself@jonasgunz.de>
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.

ARGC=$#
ARGV=($@)

DIRNAME=$(dirname $0)
[ -d $DIRNAME ] && cd $DIRNAME
unset DIRNAME

[ ! -f ./$(basename $0) ] && echo "Could not find working directory" && exit 1

BASEDIR=$(pwd)

export readonly TOOL_DIR=$BASEDIR/tools
export readonly INCLUDE_DIR=$BASEDIR/libs

TASK=
TARGET=

PARALLEL="no"

HOSTS=()
FILES=()
INCLUDES=()

#block run when sourcing task
export readonly HOSTMODE="yes"

function perror() {
	echo -e $@ 1>&2
}

selector()
{
	local regex="^-?[0-9]+\$"
	local cnt=0
	for selection in "$@"
	do
		echo "$cnt) $selection"
		((cnt=$cnt + 1))
	done

	read -p "(default=0) >" inp
	if [[ "$inp" =~ $regex ]] && [ $inp -ge 0 -a $inp -le $# ]
	then
		return $inp
	fi
	return 0
}

function print_help() {
	cat << EOF
Usage:
$0 [OPTIONS] TASK TARGET
Options:
	-H HOST		Add HOST to target hosts. can be used
			to target single hosts without defining
			a custom target
	-p		Force parallel execution
	-i		specify a SSH identity file. This is overridden by tasks!
	-h		print this help text
	-d DIR		specify a different folder containing "tasks", "targets", and "assets"
EOF
	exit $1
}

#Parse args
for (( i=0; i < $ARGC;i++ )); do
	ARGREGEX="^-.*"
	if [[ ! ${ARGV[$i]} =~ $ARGREGEX ]]; then
		[ -z $TASK ] && TASK=${ARGV[$i]} && continue
		[ -z $TARGET ] && TARGET=${ARGV[$i]} && continue

		print_help 1
	fi

	case ${ARGV[$i]} in
		-H)
			i=$((i+1))
			HOSTS+=(${ARGV[$i]});;
		-p)
			PARALLEL="yes";;
		-h)
			print_help 0;;
		-d)
			i=$((i+1))
			BASEDIR="${ARGV[$i]}";;
		-i)
			i=$((i+1))
			SSH_IDENTITY_FILE="${ARGV[$i]}";;

		*)
			echo Invalid Argument ${ARGV[$i]}
			echo $0 -h for help
			exit 1;;
	esac
done

# Set after BASEDIR might have been changed
export readonly TASK_DIR=$BASEDIR/tasks
export readonly ASSET_DIR=$BASEDIR/assets
export readonly TARGET_DIR=$BASEDIR/targets

if [ -z $TASK ]; then
	TASK_LIST=( $(ls $TASK_DIR) )
	TARGET_LIST=( $(ls $TARGET_DIR) )

	perror "\nSelect a Task"
	>&2 selector ${TASK_LIST[@]}
	TASK=${TASK_LIST[$?]}

	perror "\nSelect a Target"
	>&2 selector ${TARGET_LIST[@]}
	TARGET=${TARGET_LIST[$?]}
fi

unset TASK_ISSET

[ -f $TARGET_DIR/$TARGET ] && source $TARGET_DIR/$TARGET
[ -f $TASK_DIR/$TASK ] && source $TASK_DIR/$TASK

if [ -z $TASK_ISSET ]; then
	perror Task did not load correctly
	exit 1
fi

HOSTS_UNIQUE=($(tr ' ' '\n' <<< "${HOSTS[@]}" | sort -u))

COUNTER=0
HOSTCOUNT=${#HOSTS_UNIQUE[@]}

[ $PARALLEL = "yes" ] && AND="&"
declare -A HOST_PID

for hoststring in "${HOSTS_UNIQUE[@]}"; do
	rexecargs=""
	IFS=$':' read hostname port <<< "$hoststring"

	[ ! -z $port ] && rexecargs+="-p $port "
	[ ! -z $SSH_IDENTITY_FILE ] && rexecargs+="-i $SSH_IDENTITY_FILE "
	[ ${#FILES[@]} -gt 0 ] && for f in ${FILES[@]}; do rexecargs+="-f $f "; done
	[ ! -z $SSH_USER ] && rexecargs+="$SSH_USER@"

	rexecargs+="$hostname ${INCLUDES[@]} $TASK_DIR/$TASK"

	set -o pipefail
	eval "$TOOL_DIR/rexec.sh $rexecargs | (while read l; do echo [$hostname] \$l; done) $AND"
	RET=$?

	if [ $PARALLEL = "yes" ]; then
		HOST_PID[$hostname]=$!
	else
		[ $RET -ne 0 ] && FAILED_HOSTS+=($hostname)
		COUNTER=$(($COUNTER + 1))
		PERCENT=$(($COUNTER * 100 / $HOSTCOUNT))
		perror -n "\r[$COUNTER: $PERCENT%] "
	fi
done

if [ $PARALLEL = "yes" ]; then
	perror Waiting for asynchronous tasks to finisch...

	for hostname in ${!HOST_PID[@]}; do
		wait ${HOST_PID[$hostname]}
		[ $? -ne 0 ] && FAILED_HOSTS+=($hostname)
	done

	perror Done!
fi

if [ ${#FAILED_HOSTS[@]} -gt 0 ]; then
	perror
	perror ${#FAILED_HOSTS[@]} targets failed:
	perror ${FAILED_HOSTS[@]}
	exit 1
fi