#!/bin/bash

set -o posix -o pipefail
# bash < 4.4 erroneously considers empty arrays to be unset
[[ ${BASH_VERSINFO[0]} -gt 4 ]] || [[ ${BASH_VERSINFO[0]} -eq 4 && ${BASH_VERSINFO[1]} -ge 4 ]] && set -o nounset

progname="$(basename "$0")"
progdir="$(dirname "$0")"

verbose=0
num_ranks=1
num_hosts=0

virtual_ranks=0

cmd=()
charmrun_args=()
hosts=""

usage() {
  echo
  echo "Usage: $progname [options] <program> [program options]"
  echo
  echo "Options:"
  echo "    -h|--help                 Show this help"
  echo "    -v|--verbose              Verbose output"
  echo "    -n|-np|--np|-p <#ranks>   Number of ranks"
  echo "    --host <host1,host2,...>  Comma-separated list of hosts"
  echo "    -vr|--vr                  Use virtual ranks"
  echo
  echo "Other options will be passed unmodified to charmrun."
}


handle_hosts() {
  if [[ -z $hosts ]]; then
    num_hosts=1
    return
  fi

  echo "group main" > ampirun-nodelist

  IFS=","
  for n in $hosts; do
    num_hosts=$((num_hosts+1))
    echo "  host $n" >> ampirun-nodelist
  done
}


processArgs() {
  while [[ $# -ne 0 ]]
  do
    arg="$1"
    shift

    case "$arg" in
      -h|--help)
        usage
        exit 0
        ;;
      -v|--verbose)
        verbose=1
        ;;

      -n|-np|--np|-p)
        num_ranks="$1"
        shift
        ;;

      --host)
        hosts="$1"
        shift
        ;;

      -vr|--vr)
        virtual_ranks=1
        ;;

      --)
        cmd+=("$@")
        return
        ;;

      -*)
        echo "$progname error: unknown option \"$arg\""
        exit 1
        ;;

      *)
        cmd+=("$arg" "$@")
        return
        ;;
    esac
  done
}

processArgs "$@"

handle_hosts

if [[ $verbose == 0 ]]; then
  charmrun_args+=(++quiet)
fi

if [[ -z $hosts ]]; then
  charmrun_args+=(++local)
else
  charmrun_args+=(++nodelist ampirun-nodelist)
fi

if [[ $virtual_ranks == 0 ]]; then
  cmd+=("+p$num_ranks")
else
  cmd+=("+vp$num_ranks" "+p$num_hosts")
fi

runCmd() {
  [[ $verbose == 1 ]] && echo "$progname: Executing" "$@" 1>&2
  "$@"
}

runCmd "$progdir/charmrun" "${charmrun_args[@]}" "${cmd[@]}"
status=$?

rm -f ampirun-nodelist

exit $status
