#!/bin/bash

#---------------------------------------------------
# SCRIPT CONSTANTS DECLARATION
#---------------------------------------------------
SHARED_DISK_NAME="shared_disk"
# shellcheck source=../comons/common.sh
# shellcheck disable=SC1091
source "${COMPSS_HOME}Runtime/scripts/system/commons/logger.sh"

###############################################
# Identifies other agents which it may interact
###############################################
get_topology_plain() {

  local MASTER="${1}"
  shift 1
  local workers=${@}

  if [ "${HOSTNAME}" == "${MASTER}" ]; then
    CHILD_NODES=$(
      for blade in ${workers}; do
        if [ "${blade}" != "${MASTER}" ]; then
          echo "${blade} ${cpus_per_node} ${gpus_per_node} ${fpgas_per_node} ${memory_per_node} ${storage_bw_per_node} ${max_tasks_per_node}";
        fi
      done
    )
  else
    CHILD_NODES=""
  fi
}

get_topology_tree() {
  MASTER="${1}"

  local HOST_RACK;
  local MASTER_RACK;
  HOST_RACK=$(echo "${HOSTNAME}" | cut -d 'b' -f 1)
  MASTER_RACK=$(echo "${MASTER}" | cut -d 'b' -f 1)

  if [ "${HOSTNAME}" == "${MASTER}" ]; then
      # IS THE MASTER
      local RACK_OTHER_BLADES 
      local RACKS
      local WORKER_RACKS
      RACK_OTHER_BLADES=$(echo "${@}"| tr " " "\\n" | grep "${HOST_RACK}" | grep -v "${HOSTNAME}")
      RACKS=$(echo "${@}"| tr " " "\\n" | cut -d 'b' -f 1 | uniq )
      WORKER_RACKS=$(echo "${RACKS}"| grep -v "${MASTER_RACK}" )
      
      CHILD_NODES=$(
        for blade in ${RACK_OTHER_BLADES}; do
          echo "${blade} ${cpus_per_node} ${gpus_per_node} ${fpgas_per_node} ${memory_per_node} ${storage_bw_per_node} ${max_tasks_per_node}"
        done
        for rack in ${WORKER_RACKS}; do
          WORKER_RACK_BLADES=$(echo "${@:2}"| tr " " "\\n" | grep "${rack}")
          WORKER_RACK_REP=$(echo "${WORKER_RACK_BLADES}" | sort | head -n 1)
          WORKER_RACK_COUNT=$(echo "${WORKER_RACK_BLADES}" | wc -l)
          echo "${WORKER_RACK_REP} $(( WORKER_RACK_COUNT * cpus_per_node )) $(( WORKER_RACK_COUNT * gpus_per_node )) $(( WORKER_RACK_COUNT * fpgas_per_node )) $(( WORKER_RACK_COUNT * memory_per_node )) $(( WORKER_RACK_COUNT * storage_bw_per_node )) $(( WORKER_RACK_COUNT * max_tasks_per_node ))"
        done
      )
  else
    if [ ! "${HOST_RACK}" == "${MASTER_RACK}" ]; then

      # IS A NODE IN A RACK OTHER THAN MASTER
      local RACK_BLADES
      local MAIN_RACK_BLADE
      RACK_BLADES=$(echo "${@}"| tr " " "\\n" | grep "${HOST_RACK}" | sort)
      MAIN_RACK_BLADE=$(echo "${RACK_BLADES}" | head -n 1)

      if [ "${MAIN_RACK_BLADE}" == "${HOSTNAME}" ]; then
        # IS REPRESENTATIVE ON REMOTE RACK
        local RACK_OTHER_BLADES
        RACK_OTHER_BLADES=$(echo "${RACK_BLADES}" | grep -v "${HOSTNAME}")

        CHILD_NODES=$(
         for blade in ${RACK_OTHER_BLADES}; do
           echo "${blade} ${cpus_per_node} ${gpus_per_node} ${fpgas_per_node} ${memory_per_node} ${storage_bw_per_node} ${max_tasks_per_node}";
         done
        )
      #else
      #  IS A LEAF ON A REMOTE RACK. DO NOTHING
      fi
    #else
    # IS A LEAF ON THE MASTER. DO NOTHING
    fi
  fi
}


#---------------------------------------------------------------------------------------
# XML HELPER FUNCTIONS (PROJECT.XML)
#---------------------------------------------------------------------------------------

###############################################
# Add agents' own resources to project.xml
###############################################
xml_project_add_own_resources() {
  local cpu=${cpus_per_node}
  local mem=${memory_per_node};

  if [ "${HOSTNAME}" == "${MASTER}" ]; then
    if [ -n "${cpus_in_master}" ]; then
      cpu=${cpus_in_master}
    fi

    if [ -n "${memory_in_master}" ]; then
      mem=${memory_in_master};
    fi
  fi
  
  add_master_node "${cpu}" "${gpus_per_node}" "${fpgas_per_node}" "${mem}" "${shared_disks_info}"
}

###############################################
# Add remote agents to project.xml
###############################################
xml_project_add_child_agents() {
  # Add workers
  if [ ! -z "${CHILD_NODES}" ]; then
    while read -r worker_node cpus gpus fpgas max_tasks ignored_values; do
      local worker_node_name
      if [ ! -z "${NODE_NAME_XML}" ]; then
        worker_node_name=$(${NODE_NAME_XML} "${worker_node}" "${network}")
      elif [ -n "${NODE_NAME_QUEUE}" ]; then
        worker_node_name=$(${NODE_NAME_QUEUE} "${worker_node}")
      else
        worker_node_name=${worker_node}
      fi
      if [ "${max_tasks}" -lt "-1" ]; then
        max_tasks=-1
      fi
      # Add compute node
      worker_node_name=${worker_node_name}${network}
      add_compute_node "${worker_node_name}" "${worker_install_dir}" "${worker_working_dir}" "NULL" "NULL" "NULL" "NULL" "NULL" "NULL" "${max_tasks}"      
    done <<< "${CHILD_NODES}"
  fi
}

###############################################
# Generate project.xml
###############################################
generate_project() {
  # shellcheck source=../xmls/generate_project.sh
  # shellcheck disable=SC1091
  source "${COMPSS_HOME}Runtime/scripts/system/xmls/generate_project.sh"
  echo "[XML] Generating project at ${PROJECT_FILE}"
  xml_project_setup

  # Add header (from generate_project.sh)
  add_header

  xml_project_add_own_resources

  xml_project_add_child_agents

  # Close project (from generate_project.sh)
  add_footer
}

###############################################
# Setup to generate project.xml
###############################################
xml_project_setup() {
  # Shared disks information
  shared_disks_names=""
  shared_disks_info=""
  if [ -n "${SHARED_DISK_PREFIX}" ]; then
    shared_disks_names="${SHARED_DISK_NAME}"
    shared_disks_info="${SHARED_DISK_NAME}=${SHARED_DISK_PREFIX}"
  fi
  if [ -n "${SHARED_DISK_2_PREFIX}" ] && [ "${SHARED_DISK_2_PREFIX}" != "${SHARED_DISK_PREFIX}" ]; then
    shared_disks_names="${shared_disks_names} ${SHARED_DISK_NAME}2"
    shared_disks_info="${shared_disks_info} ${SHARED_DISK_NAME}2=${SHARED_DISK_2_PREFIX}"
  fi

  # Init project file
  init "${PROJECT_FILE}"
}



#---------------------------------------------------------------------------------------
# XML HELPER FUNCTIONS (RESOURCES.XML)
#---------------------------------------------------------------------------------------
###############################################
# Generate resources.xml
###############################################
generate_resources() {
  # shellcheck source=../xmls/generate_resources.sh
  # shellcheck disable=SC1091
  source "${COMPSS_HOME}Runtime/scripts/system/xmls/generate_resources.sh"
  
  echo "[XML] Generating resources at ${RESOURCES_FILE}"
  xml_resources_setup

  # Add header (from generate_project.sh)
  add_header

  # Add shared disks (from generate_resources.sh)
  add_shared_disks "${shared_disks_names}"

   # Add workers
  xml_resources_add_workers

  # Add elasticity
  xml_resources_add_elasticity

  # Close project (from generate_resources.sh)
  add_footer
}


###############################################
# Setup to generate resources.xml
###############################################
xml_resources_setup() {
  # Shared disks information
  shared_disks_names=""
  shared_disks_info=""
  if [ -n "${SHARED_DISK_PREFIX}" ]; then
    shared_disks_names="${SHARED_DISK_NAME}"
    shared_disks_info="${SHARED_DISK_NAME}=${SHARED_DISK_PREFIX}"
  fi
  if [ -n "${SHARED_DISK_2_PREFIX}" ] && [ "${SHARED_DISK_2_PREFIX}" != "${SHARED_DISK_PREFIX}" ]; then
    shared_disks_names="${shared_disks_names} ${SHARED_DISK_NAME}2"
    shared_disks_info="${shared_disks_info} ${SHARED_DISK_NAME}2=${SHARED_DISK_2_PREFIX}"
  fi

  # Init resources file
  init "${RESOURCES_FILE}"
}

###############################################
# Add workers to resources.xml
###############################################
xml_resources_add_workers() {
  # Add workers
  if [ ! -z "${CHILD_NODES}" ]; then
    while read -r worker_node cpus gpus fpgas memory storage_bw max_tasks_per_node; do
      local worker_node_name
      if [ ! -z "${NODE_NAME_XML}" ]; then
        worker_node_name=$(${NODE_NAME_XML} "${worker_node}" "${network}")
      elif [ -n "${NODE_NAME_QUEUE}" ]; then
        worker_node_name=$(${NODE_NAME_QUEUE} "${worker_node}")
      else
        worker_node_name=${worker_node}
      fi

      if [ -z "${max_tasks_per_node}" ] || [ "${max_tasks_per_node}" -lt "-1" ]; then
        max_tasks_per_node="-1"
      fi
      
      # Add compute node
      worker_node_name=${worker_node_name}${network}
      add_compute_node "${worker_node_name}" "${cpus}" "${gpus}" "${fpgas}" "${memory}" "${storage_bw}" "43001" "43002" "${REMOTE_EXECUTOR:-NULL}" "${shared_disks_info}"
    done <<< "${CHILD_NODES}"
  fi
}

###############################################
# Add heterogeneity to resources.xml
###############################################
xml_resources_add_elasticity() {
  # Add elasticity if defined
  if [ -n "${elasticity}" ]; then
    local instance_types="default:${cpus_per_node}:${gpus_per_node}:${fpgas_per_node}:${node_memory}:1:0.085"
    add_cloud "SLURM-Cluster" "NULL" "slurm-conn.jar" "es.bsc.conn.slurm.SlurmConnector" "${container_image}" "${shared_disks_info:-NULL}" "10" "43001" "43002" "${REMOTE_EXECUTOR:-NULL}" "${instance_types}"
  fi
}

#
# MAIN (when script is called directly)
#
# Expected parameters:
HOSTNAME="${1}"
MASTER="${2}"
worker_nodes_size="${3}"
shift 3

worker_nodes=""
for (( worker_idx=1; worker_idx<=${worker_nodes_size}; worker_idx++ ))
do
  worker_nodes="${worker_nodes} ${1}"
  shift 1
done

network="${1}" 
network=$(echo ${network} | cut -d '=' -f 2)
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
agents_hierarchy="${2}"
log_dir="${3}"
mkdir -p "${log_dir}"
trace_label="${4}"
shift 4

cpus_in_master=${1}
cpus_per_node=${2}
cpu_affinity=${3}

gpus_per_node=${4}
gpu_affinity=${5}

fpgas_per_node=${6}
fpga_affinity=${7}

memory_in_master=${8}
memory_per_node=${9}

storage_bw_per_node=${10}
shift 10

worker_working_dir=${1}
worker_install_dir=${2}
max_tasks_per_node=${3}
shift 3


INFO_FILE="${log_dir}/../${HOSTNAME}${network}.INFO"
OUT_FILE="${log_dir}/../${HOSTNAME}${network}.OUT"
ERR_FILE="${log_dir}/../${HOSTNAME}${network}.ERR"

# shellcheck disable=SC2086
"get_topology_${agents_hierarchy}" "${MASTER}" ${worker_nodes}

echo "--------${HOSTNAME}--------" >> "${INFO_FILE}"
echo "VARIABLES" >>"${INFO_FILE}"

echo "hostname:       ${HOSTNAME}" >>"${INFO_FILE}"
echo "Network         ${network}" >>"${INFO_FILE}"

echo "Topology:       ${agents_hierarchy}" >>"${INFO_FILE}"
echo "Master name:    ${MASTER}" >>"${INFO_FILE}"
echo "Workers' names: ${worker_nodes}" >>"${INFO_FILE}"

echo "Working Dir:    ${worker_working_dir}" >>"${INFO_FILE}"
echo "Install Dir:    ${worker_install_dir}" >>"${INFO_FILE}"
echo "Log Dir:        ${log_dir}" >>"${INFO_FILE}"

echo "CPUs:           ${cpus_per_node} ${cpu_affinity}" >>"${INFO_FILE}"
echo "GPUs:           ${gpus_per_node} ${gpu_affinity}" >>"${INFO_FILE}"
echo "FPGAs:          ${fpgas_per_node} ${fpga_affinity}" >>"${INFO_FILE}"
echo "Storage BW:     ${storage_bw_per_node}" >>"${INFO_FILE}"
echo "Max Tasks:      ${max_tasks_per_node}" >>"${INFO_FILE}"


echo "Child nodes:" >>"${INFO_FILE}"

if [ ! -z "${CHILD_NODES}" ]; then
  while read -r blade cpu gpu fpgas memory storage_bw limit_of_tasks; do
    echo "  * name ${blade} cpus ${cpu} gpus ${gpu} fpgas ${fpgas} memory ${memory} storage_bw ${storage_bw} limit of tasks ${limit_of_tasks}" >>"${INFO_FILE}"
  done <<< "${CHILD_NODES}"
fi


PROJECT_FILE="${worker_working_dir}/project_${HOSTNAME}.xml"
export PROJECT_FILE
generate_project


RESOURCES_FILE="${worker_working_dir}/resources_$HOSTNAME.xml"
export RESOURCES_FILE
generate_resources
echo "--------${HOSTNAME}-OUT----" >"${OUT_FILE}"
echo "--------${HOSTNAME}-ERR----" >"${ERR_FILE}"

echo "${worker_install_dir}/Runtime/scripts/user/compss_agent_start" \
"--hostname=${HOSTNAME}${network}" \
"--rest_port=46101" \
"--comm_port=46102" \
"--project=${PROJECT_FILE}" \
"--resources=${RESOURCES_FILE}" \
"--trace_label=${trace_label}" \
${@} \
>> "${OUT_FILE}"

"${worker_install_dir}/Runtime/scripts/user/compss_agent_start" \
"--hostname=${HOSTNAME}${network}" \
"--rest_port=46101" \
"--comm_port=46102" \
"--project=${PROJECT_FILE}" \
"--resources=${RESOURCES_FILE}" \
"--trace_label=${trace_label}" \
${@} \
>> "${OUT_FILE}" 2>>"${ERR_FILE}"



echo "Agent ${HOSTNAME}${network}'s process has died" >> "${OUT_FILE}"