#!/bin/bash

# Setting up COMPSs_HOME
if [ -z "${COMPSS_HOME}" ]; then
  COMPSS_HOME="$( cd "$( dirname "${BASH_SOURCE[0]}" )"/../../.. && pwd )/"
fi
if [ ! "${COMPSS_HOME: -1}" = "/" ]; then
  COMPSS_HOME="${COMPSS_HOME}/"
fi
export COMPSS_HOME=${COMPSS_HOME}

# Load auxiliar scripts

# shellcheck source=../system/commons/logger.sh"
# shellcheck disable=SC1091
source "${COMPSS_HOME}Runtime/scripts/system/commons/logger.sh"

# shellcheck source=../system/commons/utils.sh"
# shellcheck disable=SC1091
source "${COMPSS_HOME}Runtime/scripts/system/commons/utils.sh"

# shellcheck source=../system/commons/version.sh"
# shellcheck disable=SC1091
source "${COMPSS_HOME}Runtime/scripts/system/commons/version.sh"

# shellcheck source=../system/runtime/compss_setup.sh"
# shellcheck disable=SC1091
source "${COMPSS_HOME}Runtime/scripts/system/trace/generatePRVs.sh"


# Extrae Environment
if [ -z "${EXTRAE_HOME}" ]; then
  export EXTRAE_HOME="${COMPSS_HOME}/Dependencies/extrae"
fi

if [ -z "${EXTRAE_LIB}" ]; then
  export EXTRAE_LIB="${EXTRAE_HOME}/lib"
fi

###############################################
# SCRIPT CONSTANTS DECLARATION
###############################################

DEFAULT_LOG_LEVEL_ENABLED="${GEN_TRACING_LOG_LEVEL_DEBUG}"
DEFAULT_LOG_LEVEL="${GEN_TRACING_DEFAULT_LOG_LEVEL}"

DEFAULT_CUSTOM_THREADS="true"
DEFAULT_TRACE_NAME="trace"
DEFAULT_KEEP_PACKAGES="false"

###############################################
# ERROR CONSTANTS DECLARATION
###############################################


###############################################
###############################################
# Display functions
###############################################
###############################################

###############################################
# Displays usage
###############################################
usage() {
  cat << EOF
Usage: $0 [OPTION]... [INPUT_DIR]

Merges all the tracing packages generated by a COMPSs execution and merges them in a single trace.

If INPUT_DIR is not indicated, the script tries to merge traces in the current directory.

EOF

  show_gen_trace_opts

  cat << EOF
General options:
    --help, -h                              Prints this message

    --version, -v                           Prints COMPSs version
EOF
}

###############################################
# Display Arguments & Environment error
###############################################
arguments_error() {
  local error_msg=$1
  display_error "${error_msg}"
  usage
  exit 1
}

###############################################
# Display Invocation details
###############################################
display_configuration() {
  display_info \
"COMPSs Paraver trace generation. 
Traces:
  Input folder: ${in_dir}
  Output folder: ${out_dir}
  Trace name: ${trace_name}

Options:
  Custom threads: ${custom_threads}
  Keep packages: ${keep_packages}
  
Logging:
  Level: ${gen_tracing_log_level}
  Folder: ${gen_tracing_log_dir}"
  
}

###############################################
###############################################
# Option management functions
###############################################
###############################################

###############################################
# Display All Script Options
############################################### 
show_gen_trace_opts() {
  cat <<EOF
Debug options:
    -d,--debug                              Enables log at level ${DEFAULT_LOG_LEVEL_ENABLED}. 
                                            Default level: ${DEFAULT_LOG_LEVEL}

    --log_dir=<path>                        Path where to leave the log information
                                            Default: same as output directory
                                            
    --log_level=<log_level>                 Sets the log level: : ${GEN_TRACING_LOG_LEVEL_OFF} | ${GEN_TRACING_LOG_LEVEL_DEBUG} 
                                            Default level: ${DEFAULT_LOG_LEVEL}

Trace Merging options:
    --keep_packages=<bool>                  If true, the script mantains the packages after merging; otherwise, input packages are deleted.
                                            Default: ${DEFAULT_KEEP_PACKAGES}

    --custom_threads=<bool>                 Disables the thread re-organization and labeling indicating what each thread is
                                            Default: ${DEFAULT_CUSTOM_THREADS}

    --out_dir=<path>                        Directory where to save the output trace
                                            Default: same as input directory

    --trace_name=<string>                   Name of the trace gathering all the events.
                                            Default: ${DEFAULT_TRACE_NAME}
EOF
  
}

###############################################
# Parses the options from the commandline and updates the current option values
###############################################
get_args() {
  # Parse Script Options
  while getopts hvd-: flag; do
    # Treat the argument
    case "${flag}" in
      h)
        # Display help
        usage
        exit 0
        ;;
      d)
        # Enable debug in log level
        gen_tracing_log_level=${DEFAULT_LOG_LEVEL_ENABLED}
        ;;
      v)
        # Display version
        show_version
        exit 0
        ;;
      -)
      # Check more complex arguments
        case "$OPTARG" in
          # Options Description options
          help)
            # Display help
            usage
            exit 0
            ;;
          opts)
            # Display options
            show_gen_trace_opts
            exit 0
            ;;  
          # Version Options
          flower)
            # Display flower
            show_flower
            exit 0
            ;;
          recipe)
            # Display recipe
            show_recipe
            exit 0
            ;;
          version)
            # Show version
            show_full_version
            exit 0
            ;;

          # Trae Generation options
          keep_packages=*)
            keep_packages=${OPTARG//keep_packages=/}
            ;;
          custom_threads=*)
            custom_threads=${OPTARG//custom_threads=/}
            ;;
          debug*)
            # Enable debug in log level
            gen_tracing_log_level=${DEFAULT_LOG_LEVEL_ARGUMENT}
            ;;
          log_dir=*)
            # Tracing system
            gen_tracing_log_dir=${OPTARG//log_dir=/}
            ;;
          log_level=*)
            # Tracing system
            gen_tracing_log_level=${OPTARG//log_level=/}
            ;;
          out_dir=*)
            # Register desired output directory
            out_dir=${OPTARG//out_dir=/}
            ;;
          trace_name=*)
            # Register desired trace name
            trace_name=${OPTARG//trace_name=/}
            ;;
          *)
            arguments_error "Bad argument: ${OPTARG}"
            ;;
        esac
        ;;
      *)
        # Flag didn't match any patern. Raise exception
        arguments_error "Bad argument: $flag"
        ;;
    esac
  done

  # Shift COMPSs arguments
  shift $((OPTIND-1))

  # Parse input dir
  in_dir=${1}
}


###############################################
# Validates the current script configuration
################################################
check_args() {
  if [ -z "${gen_tracing_log_level}" ]; then
    gen_tracing_log_level=${DEFAULT_LOG_LEVEL}
  fi

  if [ -z "${in_dir}" ]; then
    in_dir=$(pwd)
  fi
  in_dir=$(eval "readlink -f ${in_dir}")

  
  if [ -z "${out_dir}" ]; then
    out_dir=${in_dir}
  fi
  out_dir=$(eval "readlink -f ${out_dir}")

  if [ -z "${gen_tracing_log_dir}" ]; then
    gen_tracing_log_dir=${out_dir}
  fi
  gen_tracing_log_dir=$(eval "readlink -f ${gen_tracing_log_dir}")

  if [ -z "${trace_name}" ]; then
    trace_name=${DEFAULT_TRACE_NAME}
  fi
  
  if [ -z "${custom_threads}" ]; then
    custom_threads=${DEFAULT_CUSTOM_THREADS}
  fi

  if [ -z "${keep_packages}" ]; then
    keep_packages=${DEFAULT_KEEP_PACKAGES}
  fi
}

#######################################
# Main script
#######################################

# Parse script options
get_args "$@"
check_args

display_configuration

# Generate all traces from packages
packages=$(find "${in_dir}" -name "*.tar.gz")
gen_traces "${out_dir}" "${trace_name}" "1" ${packages}
if [[ ${endCode} -ne 0 ]]; then
  echo "FAILED gen_traces"
  exit "${endCode}"
else
  echo "SUCCESS gen_traces"
fi

# Merge python traces into main
python_traces=$(find "${out_dir}/python" -name "*.prv")
merge_python_traces "${out_dir}" "${trace_name}" ${python_traces}
if [[ ${endCode} -ne 0 ]]; then
  echo "FAILED merge_python_traces"
  exit "${endCode}"
else
  echo "SUCCESS merge_python_traces"
fi
rm -rf "${out_dir}/python"

# Re-arranging PRV trace with custom threads
if [ "${custom_threads}" == "true" ]; then
  rearrange_trace_threads "${out_dir}" "${trace_name}"
  if [[ ${endCode} -ne 0 ]]; then
    echo "FAILED rearrange_trace_threads"
    exit "${endCode}"
  else
    echo "SUCCESS rearrange_trace_threads"
  fi
fi

# Clear packages if required
if [ "${keep_packages}" == "false" ]; then
  rm -rf ${packages}
fi
