#!/bin/bash

###############################################
# SCRIPT CONSTANTS DECLARATION
###############################################
DEFAULT_SC_CFG="default"

# Next default values need to be consistent with runcompss
DEFAULT_LIBRARY_PATH=$(pwd)
DEFAULT_APPDIR=$(pwd)
DEFAULT_CLASSPATH=$(pwd)
DEFAULT_PYTHONPATH=$(pwd)
DEFAULT_MASTER_PORT_BASE=43000
DEFAULT_MASTER_PORT_RAND_RANGE=1000
DEFAULT_CPU_AFFINITY="automatic"
DEFAULT_FPGA_REPROGRAM=""


###############################################
# ERROR CONSTANTS DECLARATION
###############################################
ERROR_CFG_SC="SuperComputer CFG file doesn't exist"
ERROR_CFG_Q="Queue system CFG file doesn't exist"
ERROR_MASTER_NODE="Missing master node parameter"
ERROR_WORKER_NODES="Missing worker nodes parameter"
ERROR_NUM_CPUS="Invalid number of CPUS per node"
ERROR_WORKER_WD="Invalid Worker Working Dir option"
ERROR_NETWORK="Invalid network option"
ERROR_WORKER_IN_MASTER_CPUS="Parameter worker_in_master_cpus is bigger than the maximum number of cpus_per_node"
ERROR_WORKER_IN_MASTER_MEMORY="Incorrect worker_in_master_memory parameter. Only disabled or <int> allowed. I.e. 33000, 66000"
ERROR_WORKER_IN_MASTER_MEMORY_TOO_HIGH="Parameter worker_in_master_memory exceeds the node_memory limit"
ERROR_WORKER_IN_MASTER_MEMORY_NOT_SPECIFIED="Parameter worker_in_master_memory is mandatory if worker_in_master_cpus is not 0"
ERROR_LANGUAGE="Value of option --lang must be: java, c or python"

#---------------------------------------------------------------------------------------
# HELPER FUNCTIONS
#---------------------------------------------------------------------------------------

###############################################
# Displays usage
###############################################
usage() {
  local exitValue=$1

  cat <<EOT
Usage: $0 [options]

* Options:
  General:
    --help, -h                              Print this help message

    --opts                                  Show available options

    --version, -v                           Print COMPSs version

    --sc_cfg=<name>                         SuperComputer configuration file to use. Must exist inside queues/cfgs/
                                            Mandatory
                                            Default: ${DEFAULT_SC_CFG}

    --master_node=<string>                  Node where to run the COMPSs Master
                                            Mandatory

    --worker_nodes="<string string...>"     Space separated nodes where to run the COMPSs Workers (Notice the quotes)
                                            Mandatory

  Specific configuration:
EOT

  show_opts "$exitValue"
}

###############################################
# Show Options
###############################################
show_opts() {
  local exitValue=$1

  # Load default CFG for default values
  local defaultSC_cfg=${SCRIPT_DIR}/../queues/cfgs/${DEFAULT_SC_CFG}.cfg
  #shellcheck source=../queues/cfgs/default.cfg
  source "${defaultSC_cfg}"
  local defaultQS_cfg=${SCRIPT_DIR}/../queues/${QUEUE_SYSTEM}/${QUEUE_SYSTEM}.cfg
  #shellcheck source=../queues/slurm/slurm.cfg
  source "${defaultQS_cfg}"

  # Show usage
  cat <<EOT
    --cpus_per_node=<int>                   Available CPU computing units on each node
                                            Default: ${DEFAULT_CPUS_PER_NODE}
    --gpus_per_node=<int>                   Available GPU computing units on each node
                                            Default: ${DEFAULT_GPUS_PER_NODE}
    --fpgas_per_node=<int>                  Available FPGA computing units on each node
                                            Default: ${DEFAULT_FPGAS_PER_NODE}
    --fpga_reprogram="<string>"             Specify the full command that needs to be executed to reprogram the FPGA with
                                            the desired bitstream. The location must be an absolute path.
                                            Default: ${DEFAULT_FPGA_REPROGRAM}
    --max_tasks_per_node=<int>              Maximum number of simultaneous tasks running on a node
                                            Default: ${DEFAULT_MAX_TASKS_PER_NODE}
    --node_memory=<MB>                      Maximum node memory: disabled | <int> (MB)
                                            Default: ${DEFAULT_NODE_MEMORY}
    --network=<name>                        Communication network for transfers: default | ethernet | infiniband | data.
                                            Default: ${DEFAULT_NETWORK}

    --master_working_dir=<path>             Working directory of the application
                                            Default: ${DEFAULT_MASTER_WORKING_DIR}
    --worker_working_dir=<name | path>      Worker directory. Use: scratch | gpfs | <path>
                                            Default: ${DEFAULT_WORKER_WORKING_DIR}

    --worker_in_master_cpus=<int>           Maximum number of CPU computing units that the master node can run as worker. Cannot exceed cpus_per_node.
                                            Default: ${DEFAULT_WORKER_IN_MASTER_CPUS}
    --worker_in_master_memory=<int> MB      Maximum memory in master node assigned to the worker. Cannot exceed the node_memory.
                                            Mandatory if worker_in_master_cpus is specified.
                                            Default: ${DEFAULT_WORKER_IN_MASTER_MEMORY}
    --jvm_worker_in_master_opts="<string>"  Extra options for the JVM of the COMPSs Worker in the Master Node.
                                            Each option separed by "," and without blank spaces (Notice the quotes)
                                            Default: ${DEFAULT_JVM_WORKER_IN_MASTER}
    --container_image=<path>                Runs the application by means of a container engine image
                                            Default: Empty
    --container_compss_path=<path>          Path where compss is installed in the container image
                                            Default: /opt/COMPSs
    --container_opts="<string>"             Options to pass to the container engine
                                            Default: empty
    --elasticity=<max_extra_nodes>          Activate elasticity specifiying the maximum extra nodes (ONLY AVAILABLE FORM SLURM CLUSTERS WITH NIO ADAPTOR)
                                            Default: 0
    --queue=<string>                        Elasticity queue
                                            Default:
    --reservation=<string>                  Elasticity reservation
                                            Default:
    --qos=<string>                          Elasticity QoS
                                            Default:
    --constraints=<string>                  Elasticity constraints
                                            Default:

EOT

  exit "$exitValue"
}

###############################################
# Displays version
###############################################
display_version() {
  local exitValue=$1

  "${SCRIPT_DIR}"/runcompss --version

  exit "$exitValue"
}

###############################################
# Displays errors when treating arguments
###############################################
display_error() {
  local error_msg=$1

  echo "$error_msg"
  echo " "

  usage 1
}


#---------------------------------------------------
# MAIN FUNCTIONS DECLARATION
#---------------------------------------------------

###############################################
# Function to get the arguments
###############################################
get_args() {
  # Avoid enqueue if there is no application
  if [ $# -eq 0 ]; then
    usage 1
  fi

  # Parse COMPSs Options
  while getopts hvgtmd-: flag; do
    # Treat the argument
    case "$flag" in
      h)
        # Display help
        usage 0
        ;;
      v)
        # Display version
        display_version 0
        ;;
      -)
      # Check more complex arguments
      case "$OPTARG" in
        help)
          # Display help
          usage 0
          ;;
        version)
          # Display compss version
          display_version 0
          ;;
        opts)
          # Display options
          show_opts 0
          ;;
        master_node=*)
          master_node=${OPTARG//master_node=/}
          ;;
        worker_nodes=*)
          worker_nodes=${OPTARG//worker_nodes=/}
          ;;
        sc_cfg=*)
          sc_cfg=${OPTARG//sc_cfg=/}
          ;;
        cpus_per_node=*)
          cpus_per_node=${OPTARG//cpus_per_node=/}
          ;;
        gpus_per_node=*)
          gpus_per_node=${OPTARG//gpus_per_node=/}
          ;;
        fpgas_per_node=*)
          fpgas_per_node=${OPTARG//fpgas_per_node=/}
          ;;
        max_tasks_per_node=*)
          max_tasks_per_node=${OPTARG//max_tasks_per_node=}
          ;;
        cpu_affinity=*)
          cpu_affinity=${OPTARG//cpu_affinity=}
          ;;
        fpga_reprogram=*)
          fpga_prog=${OPTARG//fpga_reprogram=/}
          ;;
        master_working_dir=*)
          master_working_dir=${OPTARG//master_working_dir=/}
          ;;
        worker_working_dir=*)
          worker_working_dir=${OPTARG//worker_working_dir=/}
          ;;
        worker_in_master_cpus=*)
          worker_in_master_cpus=${OPTARG//worker_in_master_cpus=/}
          ;;
        worker_in_master_memory=*)
          worker_in_master_memory=${OPTARG//worker_in_master_memory=/}
          ;;
        node_memory=*)
          node_memory=${OPTARG//node_memory=/}
          ;;
        network=*)
          network=${OPTARG//network=/}
          ;;
        lang=*)
          lang=${OPTARG//lang=/}
          ;;
        library_path=*)
          library_path=${OPTARG//library_path=/}
          ;;
        classpath=*)
          cp=${OPTARG//classpath=/}
          ;;
        pythonpath=*)
          pythonpath=${OPTARG//pythonpath=/}
          ;;
	      appdir=*)
	        appdir=${OPTARG//appdir=/}
	        ;;
        jvm_master_opts=*)
          jvm_master_opts=${OPTARG//jvm_master_opts=/}
          ;;
        jvm_workers_opts=*)
          jvm_workers_opts=${OPTARG//jvm_workers_opts=/}
          ;;
        jvm_worker_in_master_opts=*)
          jvm_worker_in_master_opts=${OPTARG//jvm_worker_in_master_opts=/}
          ;;
        master_port=*)
          # Remove from runcompss since launcher will add it
          master_port=${OPTARG//master_port=/}
          ;;
	      container_image=*)
          container_image=${OPTARG//container_image=/}
          ;;
        container_compss_path=*)
          container_compss_path=${OPTARG//container_compss_path=/}
          ;;
	      container_opts=*)
          container_opts=${OPTARG//container_opts=/}
          ;;
        elasticity=*)
          elasticity=${OPTARG//elasticity=/}
          ;;
        queue=*)
          # Added for elasticity
          queue=${OPTARG//queue=/}
          ;;
        reservation=*)
          # Added for elasticity
          reservation=${OPTARG//reservation=/}
          ;;
        qos=*)
          # Added for elasticity
          qos=${OPTARG//qos=/}
          ;;
        constraints=*)
          #Added for elasticity
          constraints=${OPTARG//constraints=/}
          ;;
	      *)
          true  # pass
          ;;
      esac
      ;;
    *)
      true  # pass
      ;;
    esac
  done

}

###############################################
# Function to check the arguments
###############################################
check_args() {
  ###############################################################
  # SC Configuration checks
  ###############################################################
  # Check sc configuration argument
  if [ -z "${sc_cfg}" ]; then
    sc_cfg=${DEFAULT_SC_CFG}
  fi
  if [[ ${sc_cfg} != *cfg ]]; then
    # Add cfg suffix
    sc_cfg=${sc_cfg}.cfg
  fi

  local scCfgFullPath=${SCRIPT_DIR}/../queues/cfgs/${sc_cfg}

  if [ ! -f "${scCfgFullPath}" ]; then
    # CFG file doesn't exist
    display_error "${ERROR_CFG_SC}"
  fi

  # Source SC CFG env
  # shellcheck source=../queues/cfgs/default.cfg
  source "${scCfgFullPath}"

  # Check queue configuration env
  local queueCfgFullPath=${SCRIPT_DIR}/../queues/${QUEUE_SYSTEM}/${QUEUE_SYSTEM}.cfg
  if [ ! -f "${queueCfgFullPath}" ]; then
    # CFG file doesn't exist
    display_error "${ERROR_CFG_Q}"
  fi

  # Source queue system CFG env
  # shellcheck source=../queues/slurm/slurm.cfg
  source "${queueCfgFullPath}"

  ###############################################################
  # Infrastructure checks
  ###############################################################
  if [ -z "${master_node}" ]; then
    display_error "${ERROR_MASTER_NODE}"
  fi

  if [ -z "${worker_nodes}" ] && [ "${DEFAULT_WORKER_IN_MASTER_CPUS}" -eq "0" ]; then
    if [ -z "${worker_in_master_cpus}" ]; then
      display_error "${ERROR_WORKER_NODES}"
    else
      worker_nodes=""
    fi
  fi

  if [ -z "${network}" ]; then
    network=${DEFAULT_NETWORK}
  elif [ "${network}" == "default" ]; then
    network=${DEFAULT_NETWORK}
  elif [ "${network}" != "ethernet" ] && [ "${network}" != "infiniband" ] && [ "${network}" != "data" ]; then
    display_error "${ERROR_NETWORK}"
  fi

  ###############################################################
  # Node checks
  ###############################################################
  if [ -z "${max_tasks_per_node}" ]; then
    max_tasks_per_node=${DEFAULT_MAX_TASKS_PER_NODE}
  fi

  if [ -z "${cpus_per_node}" ]; then
    cpus_per_node=${DEFAULT_CPUS_PER_NODE}
  fi

  if [ "${cpus_per_node}" -lt "${MINIMUM_CPUS_PER_NODE}" ]; then
    display_error "${ERROR_NUM_CPUS}"
  fi

  if [ -z "${gpus_per_node}" ]; then
    gpus_per_node=${DEFAULT_GPUS_PER_NODE}
  fi

  if [ -z "${fpgas_per_node}" ]; then
    fpgas_per_node=${DEFAULT_FPGAS_PER_NODE}
  fi

  if [ -z "${cpu_affinity}" ]; then
    cpu_affinity=${DEFAULT_CPU_AFFINITY}
  fi

  if [ -z "${fpga_prog}" ]; then
    fpga_prog=${DEFAULT_FPGA_REPROGRAM}
  fi

  if [ -z "${node_memory}" ]; then
    node_memory=${DEFAULT_NODE_MEMORY}
  fi

  if [ -z "${worker_in_master_cpus}" ]; then
    worker_in_master_cpus=${DEFAULT_WORKER_IN_MASTER_CPUS}
  fi
  if [ "${worker_in_master_cpus}" -gt "${cpus_per_node}" ]; then
    display_error "${ERROR_WORKER_IN_MASTER_CPUS}"
  fi

  if [ -z "${worker_in_master_memory}" ]; then
    worker_in_master_memory=${DEFAULT_WORKER_IN_MASTER_MEMORY}
  elif [ "${worker_in_master_memory}" != "disabled" ] && ! [[ "${worker_in_master_memory}" =~ ^[0-9]+$ ]]; then
    display_error "${ERROR_WORKER_IN_MASTER_MEMORY}"
  fi
  if [ "${worker_in_master_memory}" != "${DEFAULT_WORKER_IN_MASTER_MEMORY}" ] && [ "${node_memory}" != "${DEFAULT_NODE_MEMORY}" ]; then
    if [ "${worker_in_master_memory}" -gt "${node_memory}" ]; then
      display_error "${ERROR_WORKER_IN_MASTER_MEMORY_TOO_HIGH} ${worker_in_master_memory} < ${node_memory} "
    fi
  fi

  if [ "${worker_in_master_cpus}" -gt 0 ] && [ "${worker_in_master_memory}" -le 0 ]; then
    display_error "${ERROR_WORKER_IN_MASTER_MEMORY_NOT_SPECIFIED}"
  fi

  if [ -z "${master_port}" ]; then
    rand_num=$RANDOM
    offset=$((rand_num % DEFAULT_MASTER_PORT_RAND_RANGE))
    master_port=$((DEFAULT_MASTER_PORT_BASE + offset))
  fi

  while lsof -i :${master_port}; do
    echo "Port ${master_port} is already in use, incrementing port by 1"
    master_port=$((master_port+1))
  done

  ###############################################################
  # Working Directory Checks
  ###############################################################
  if [ -z "${master_working_dir}" ]; then
    master_working_dir=${DEFAULT_MASTER_WORKING_DIR}
  fi

  if [ -z "${worker_working_dir}" ]; then
    worker_working_dir=${DEFAULT_WORKER_WORKING_DIR}
  elif [ "${worker_working_dir}" != "scratch" ] && [ "${worker_working_dir}" != "gpfs" ] && [[ ${worker_working_dir} != /* ]]; then
    display_error "${ERROR_WORKER_WD}"
  fi

  ###############################################################
  # JVM Checks
  ###############################################################
  if [ -z "${jvm_master_opts}" ]; then
    jvm_master_opts=${DEFAULT_JVM_MASTER}
  fi
  jvm_master_opts=${jvm_master_opts//\"/}   # not in xmls, but ketp since is in cfgs

  if [ -z "${jvm_workers_opts}" ]; then
    jvm_workers_opts=${DEFAULT_JVM_WORKERS}
  fi
  jvm_workers_opts=${jvm_workers_opts//\"/}

  if [ -n "${worker_in_master_cpus}" ]; then
    jvm_worker_in_master_opts=${DEFAULT_JVM_WORKER_IN_MASTER}
  fi
  jvm_worker_in_master_opts=${jvm_worker_in_master_opts//\"/}

  ###############################################################
  # Application Checks
  ###############################################################
  # Lang
  if [ "$lang" = "java" ]; then
    lang=java
  elif [ "$lang" = "c" ]; then
    lang=c
  elif [ "$lang" = "python" ]; then
    lang=python
  else
    display_error "${ERROR_LANGUAGE}"
  fi

  # Library path
  if [ -z "${library_path}" ]; then
    library_path=${DEFAULT_LIBRARY_PATH}
  fi

  # Classpath
  if [ -z "$cp" ]; then
    cp=${DEFAULT_CLASSPATH}
  else
    fcp=""
    for currcp in ${cp//:/$'\n'}; do
    if [ ! "${currcp:0:1}" == '/' ]; then            # Relative paths to abs
      if [ -d "$currcp" ] || [ -f "$currcp" ]; then  # If the dir/file exists
        absdir="$(cd "$(dirname "$currcp")" && pwd)" # Get absolute dir
        file="$(basename "$currcp")"
        currcp="$absdir/$file"
      else
        echo "[ WARNING ]: Classpath \"$currcp\" does not exist..."
      fi
    fi
    fcp="${fcp}:$currcp"
    done
    cp="$(echo "$fcp" | cut -c2-)"
  fi

  # Pythonpath
  if [ -z "$pythonpath" ]; then
    pythonpath=${DEFAULT_PYTHONPATH}
  fi

  # AppDir
  if [ -z "$appdir" ]; then
    appdir=${DEFAULT_APPDIR}
  fi
}


###############################################
# Sets job variables
###############################################
set_variables() {
  # Set script variables
  if [ -z "${container_image}" ]; then
    export COMPSS_HOME=${COMPSS_HOME}
    export GAT_LOCATION=${COMPSS_HOME}/Dependencies/JAVA_GAT
    worker_install_dir=${COMPSS_HOME}
  else
    if [ -z "${container_compss_path}" ]; then
      export COMPSS_HOME=/opt/COMPSs/
      export GAT_LOCATION=/opt/COMPSs/Dependencies/JAVA_GAT
      worker_install_dir=/opt/COMPSs
    else
      export COMPSS_HOME=${container_compss_path}/
      export GAT_LOCATION=${container_compss_path}/Dependencies/JAVA_GAT
      worker_install_dir=${container_compss_path}/
    fi
  fi

  # SharedDisk variables
  if [ "${worker_working_dir}" == "gpfs" ]; then
    worker_working_dir=${GPFS_PREFIX}${HOME}
  elif [ "${worker_working_dir}" == "scratch" ]; then
    worker_working_dir=$TMPDIR/${!ENV_VAR_JOB_ID}
  else
    # The working dir is a custom absolute path, create tmp
    worker_working_dir=${worker_working_dir}
  fi

  # Network variables
  if [ "${network}" == "ethernet" ]; then
    network=""
  elif [ "${network}" == "infiniband" ]; then
    network=${NETWORK_INFINIBAND_SUFFIX}
  elif [ "${network}" == "data" ]; then
    network=${NETWORK_DATA_SUFFIX}
  fi

  # Memory variables
  if [ "${node_memory}" == "${DEFAULT_NODE_MEMORY}" ]; then
    # Default value
    node_memory=${DEFAULT_NODE_MEMORY_SIZE}
  else
    # Change from MB to GB
    node_memory=$(( node_memory / 1024 - 4))
  fi

  if [ "${worker_in_master_memory}" == "${DEFAULT_WORKER_IN_MASTER_MEMORY}" ]; then
    # Default value
    worker_in_master_memory=${DEFAULT_NODE_MEMORY_SIZE}
  else
    # Change from MB to GB
    worker_in_master_memory=$(( worker_in_master_memory / 1024 - 4))
  fi

}

###############################################
# Creates XML Files
###############################################
log_variables() {
  echo "-------- Launch arguments --------"
  echo "Master:                    ${master_node}"
  echo "Workers:                   ${worker_nodes}"
  echo "Tasks per Node:            ${max_tasks_per_node}"
  echo "CPUs per Node:             ${cpus_per_node}"
  echo "GPUs per Node:             ${gpus_per_node}"
  echo "FPGAs per Node:            ${fpgas_per_node}"
  echo "CPU Affinity:              ${cpu_affinity}"
  echo "FPGA reprogram command:    ${fpga_prog}"
  echo "Network:                   ${network}"
  echo "Worker in Master CPUs:     ${worker_in_master_cpus}"
  echo "Worker in Master Memory:   ${worker_in_master_memory}"
  echo "Master Port:               ${master_port}"
  echo "Master WD:                 ${master_working_dir}"
  echo "Worker WD:                 ${worker_working_dir}"
  echo "Master JVM Opts:           ${jvm_master_opts}"
  echo "Workers JVM Opts:          ${jvm_workers_opts}"
  echo "Worker in Master JVM Opts: ${jvm_worker_in_master_opts}"
  echo "Library Path:              ${library_path}"
  echo "Classpath:                 ${cp}"
  echo "Pythonpath:                ${pythonpath}"
  echo "Appdir:                    ${appdir}"
  echo "Lang:                      ${lang}"
  echo "-----------------------------------"
  echo " "
}


#---------------------------------------------------------------------------------------
# MAIN EXECUTION
#---------------------------------------------------------------------------------------

  # Set script variables
  if [ -z "$COMPSS_HOME" ]; then
     SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
     COMPSS_HOME=${SCRIPT_DIR}/../../../
  else
     SCRIPT_DIR="${COMPSS_HOME}/Runtime/scripts/user"
  fi

  # Get command args
  get_args "$@"

  # Check other command args
  check_args

  # Set job variables
  set_variables

  # Log variables
  log_variables

  # Create XML files
  source ${SCRIPT_DIR}/../system/xmls/xmls_utils.sh
  CREATE_WORKING_DIRS=false
  create_xml_files
