#!/bin/sh

SUBMIT_COMMON=$(which work_queue_submit_common)
if [ -z "$SUBMIT_COMMON" ];
then
	echo "Please add 'work_queue_submit_common' to your PATH." 1>&2
	exit 1
else
	. $SUBMIT_COMMON
fi


show_help()
{
		echo "Use: ec2_submit_workers [options] <servername> <port> <ec2-key-name> <ec2-key-file> <num-ec2-instances>"
		echo "where options are:"
		echo "  -k,--key-name <string>  Name of the EC2 key name. (required)"
		echo "  -K,--key-file <path>    Path to the EC2 key file. (required)"
		echo "  -n,--instances <number> Set the number of EC2 instances to launch. Each instance will"
		echo "                          launch <num-workers> workers. (required)"
		echo "  -I <image_id>           EC2 OS image ID. Default = ami-fa01f193 (Ubuntu 10.04 x86_64)."
		echo "  -Z <instance_size>      EC2 instance size. Default = m1.large."
		echo "  -p <parameters>         Set parameters for ec2-run-instances."
		echo "  -h                      Show this help message."
		exit 1
}

image=ami-fa01f193
instance_size=m1.large
ec2_parameters=""
ec2_key_name=""
ec2_key_path=""
ec2_num_of_instances=""

parse_arguments()
{
	while [ $# -gt 0 ]
	do
		case $1 in
			-n | --instances)
			shift
			ec2_num_of_instances=$1
			;;
			-k | --key-name)
			shift
			ec2_key_name=$1
			;;
			-K | --key-file)
			shift
			ec2_key_path=$1
			;;
			-p)
			shift
			ec2_parameters="$ec2_parameters $1"
			;;
			*)
			break
			;;
		esac
		shift
	done

	if [ -z $ec2_key_name ]
	then
		echo "A key name should be provided with the -k option."
		exit 1
	fi

	if [ -z $ec2_key_path ]
	then
		echo "A path to the key file should be provided with the -K option."
		exit 1
	fi

	if [ -z $ec2_num_of_instances ]
	then
		echo "The number of instances to launch should be provided with the -n option."
		exit 1
	fi
}


submit_workers_command()
{
	ec2run=`which ec2-run-instances 2>/dev/null`
	if [ $? != 0 ]
	then
		echo "$0: please add 'ec2-run-instances' to your PATH."
		exit 1
	fi

	cat >worker.sh <<EOF
#!/bin/sh

for i in `eval echo {1..$count}`
 do
	./work_queue_worker $arguments $host $port &
done
EOF

	chmod 755 worker.sh

	#Start EC2 instances of the required count.
	ec2run_out=$(ec2-run-instances $image --instance-type $instance_size -k $keyname $parameters -n $count)
	reservation_id=$(echo $ec2run_out | awk '{print $2}')
	echo "Created $count instance(s) under reservation $reservation_id."
	sleep 60

	#Wait until all of the requested instances come on-board..
	running_instances=0
	while [ $running_instances -lt $count ]
	do
		sleep 20
		running_instances=$(ec2-describe-instances -F 'reservation-id'=$reservation_id | grep $image | grep running | wc -l)
	done

	inst_count=1
	while [ $inst_count -le $count ]
	do
		instance=$(ec2-describe-instances -F 'reservation-id'=$reservation_id | grep running | awk 'NR==var{print $4 }' var="${inst_count}")
		echo "------------------------------------------------"
		echo "Starting worker #$inst_count on $instance"
		scp -o StrictHostKeyChecking=no -o ConnectTimeout=20 -o ConnectionAttempts=1 -i $keyfile $worker worker.sh ubuntu@$instance:.
		ssh -o StrictHostKeyChecking=no -o ConnectTimeout=20 -o ConnectionAttempts=1 -i $keyfile ubuntu@$instance './worker.sh 1>worker.log 2>&1 &'
		return_status=$?
		inst_count=$((inst_count+1))
	done
}

submit_workers "$@"
