#!/bin/bash

HOSTS=
CPUS=2
CORES=4
THREADS=1
HOST_PREFIX="node-"
YAML=""
CREATE_PROPERTIES=1
USE_THREADS=""

usage() {
  cat <<EOF
Usage:
    $0 -H <# of hosts> [other options]

    Generate commands to add new resources to OAR database

Options:
    -T, --use-threads           use the thread resource property
    -H, --hosts <#>             # of hosts
    -C, --cpus <#>              # of cpu per host
    -c, --cores <#>             # of core per cpu
    -t, --threads <#>           # of threads per core
    --host0 <#>                 first host id to use
    --cpu0 <#>                  first cpu id to use
    --core0 <#>                 first core id to use
    --thread0 <#>               first thread id to use
    --cpuset <#>                # of cpusets/host (default: cpus*cores*threads)
    --host-prefix <str>         hostname prefix (default: "node-")
    --host-suffix <str>         hostname suffix (e.g. ".domain")
    -a, --auto-offset           guess first host/cpu/core/thread ids
    -p, --no-create-properties  do not generate oarproperty commands
    -A, --append <str>          append a text string (extra properties)
    -o, --write-to <file>       write commands to file
    -Y, --yaml                  generate YAML output
    -h, --help                  display this message

EOF
}

die() {
  cat <<EOF 2>&1
Error: $1

EOF
  usage 2>&1
  exit 1
}

LONG_OPTS="hosts:,cpus:,cores:,threads:,host-prefix:,host-suffix:,host0:,cpu0:,core0:,thread0:,cpuset:,append:,use-threads,write-to:,auto-offset,yaml,no-create-properties,help"
SHORT_OPTS="H:C:c:t:P:S:A:o:TaYph"
args=$(getopt -l $LONG_OPTS -o $SHORT_OPTS -q -- "$@")
[ $? -gt 0 ] && die "Syntax error, $(getopt -l $LONG_OPTS -o $SHORT_OPTS -Q -- "$@" 2>&1)"

eval set -- "$args"

while [ $# -ge 1 ]; do
  case "$1" in
  --)
    # No more options left.
    shift
    break
    ;;
  -H|--hosts)
    HOSTS=$2
    shift
    ;;
  -C|--cpus)
    CPUS=$2
    shift
    ;;
  -c|--cores)
    CORES=$2
    shift
    ;;
  -t|--threads)
    THREADS=$2
    shift
    ;;
  -P|--host-prefix)
    HOST_PREFIX=$2
    shift
    ;;
  -S|--host-suffix)
    HOST_SUFFIX=$2
    shift
    ;;
  --host0)
    HOST0=$2
    shift
    ;;
  --cpu0)
    CPU0=$2
    shift
    ;;
  --core0)
    CORE0=$2
    shift
    ;;
  --thread0)
    THREAD0=$2
    shift
    ;;
  --cpuset)
    CPUSET=$2
    shift
    ;;
  -A|--append)
    APPEND=$2
    shift
    ;;
  -o|--write-to)
    WRITE_TO=$2
    shift
    ;;
  -T|--use-threads)
    USE_THREADS=1
    ;;
  -a|--auto-offset)
    AUTO_OFFSET=1
    ;;
  -p|--no-create-properties)
    CREATE_PROPERTIES=
    ;;
  -Y|--yaml)
    YAML=1
    ;;
  -h|--help)
    usage
    exit 0
    ;;
  *)
    usage
    exit 1
    ;;
  esac
  shift
done

[ -n "$HOSTS" ] && [ $HOSTS -gt 0 ] || die "Syntax error, need a # of host"

if [ -n "$WRITE_TO" ]; then
  if [ -e "$WRITE_TO" ]; then
    echo -n > $WRITE_TO
  fi
  exec 1> >(tee -a $WRITE_TO)
fi

CPUSET=${CPUSET:-$((CPUS*CORES*THREADS))}
if [ -n "$AUTO_OFFSET" ]; then
  [ -z "$HOST0" ] && echo "# Warning: guessing a new hostname is not really reliable because of a sort issue for non-numeric properties. Please double-check."
  HOST0=${HOST0:-$(($(oarnodesetting --last-property-value host | perl -pe 's/^[^\d]*(\d+).*/$1/') + 1))}
  CPU0=${CPU0:-$(($(oarnodesetting --last-property-value cpu) + 1))}
  CORE0=${CORE0:-$(($(oarnodesetting --last-property-value core) + 1))}
  if [ -n "$USE_THREADS" ]; then
    THREAD0=${THREAD0:-$(($(oarnodesetting --last-property-value thread) + 1))}
  fi
else
  HOST0=${HOST0:-1}
  CPU0=${CPU0:-0}
  CORE0=${CORE0:-0}
  THREAD0=${THREAD0:-0}
fi

host=1
thread=0
core=0
cpu=0

if [ -z "$USE_THREADS" ]; then
  # Force #threads=1 for the loop below to work. 
  THREADS=1
fi

if [ -n "$CREATE_PROPERTIES" -a -z "$YAML" ]; then
  cat <<EOF
oarproperty -c -a host || true
oarproperty -a cpu || true
oarproperty -a core || true
EOF
  if [ -n "$USE_THREADS" ]; then
    cat <<EOF
oarproperty -a thread || true
EOF
  fi
fi

if [ -n "$YAML" ]; then
  echo "---"
fi

while [ $host -le $HOSTS ]; do
  hostname="'$HOST_PREFIX$((host+HOST0-1))$HOST_SUFFIX'"
  cpuset=0
  while [ $cpu -lt $((CPUS * host)) ]; do
    while [ $core -lt $((CORES * (cpu+1))) ]; do
      while [ $thread -lt $((THREADS * (core+1))) ]; do
        if [ -n "$YAML" ]; then
          if [ -z "$USE_THREADS" ]; then
            cat <<EOF
- network_address: $hostname
  host: $hostname
  cpu: $((cpu+CPU0))
  core: $((core+CORE0))
  cpuset: $cpuset
EOF
            echo -ne "$APPEND"
          else
            cat <<EOF
- network_address: $hostname
  host: $hostname
  cpu: $((cpu+CPU0))
  core: $((core+CORE0))
  thread: $((thread+THREAD0))
  cpuset: $cpuset
EOF
            echo -ne "$APPEND"
          fi
        else
          if [ -z "$USE_THREADS" ]; then
            echo "oarnodesetting -a -h $hostname -p host=$hostname -p cpu=$((cpu+CPU0)) -p core=$((core+CORE0)) -p cpuset=$cpuset $APPEND"
          else
            echo "oarnodesetting -a -h $hostname -p host=$hostname -p cpu=$((cpu+CPU0)) -p core=$((core+CORE0)) -p thread=$((thread+THREAD0)) -p cpuset=$cpuset $APPEND"
          fi
        fi
        ((thread++))
        cpuset=$(((cpuset+1) % CPUSET))
      done
      ((core++))
    done
    ((cpu++))
  done
  ((host++))
done
