#!/bin/bash
CURRENT_DIR="$( dirname -- "$0"; )";
ROCPROFV2_DIR=$(dirname -- $(realpath ${BASH_SOURCE[0]}));
ROCM_DIR=$( dirname -- "$ROCPROFV2_DIR"; )
PLUGIN_LIST=("ctf" "perfetto" "file" "att")
RUN_FROM_BUILD=0
if [[ $ROCPROFV2_DIR == *"/build"* ]]; then
  RUN_FROM_BUILD=1
  ROCM_DIR=$ROCPROFV2_DIR
fi

export ROCPROFILER_METRICS_PATH=$ROCM_DIR/libexec/rocprofiler/counters/derived_counters.xml
export LD_LIBRARY_PATH=$ROCM_DIR/lib:$LD_LIBRARY_PATH

usage() {
  echo -e "ROCProfilerV2 Run Script Usage:"
  echo -e "-h   | --help               For showing this message"
  echo -e "--list-counters             For showing all available counters for the current GPUs"
  if [ $RUN_FROM_BUILD == 1 ]; then
    echo -e "-t   | --test               For Running the tests"
    echo -e "-mt  | --mem-test           For Running the Memory Leak tests. This run requires building using -acb | --asan-clean-build option"
  fi
  echo -e "--hip-api                    For Collecting HIP API Traces"
  echo -e "--hip-activity | --hip-trace For Collecting HIP API Activities Traces"
  echo -e "--hsa-api                    For Collecting HSA API Traces"
  echo -e "--hsa-activity | --hsa-trace For Collecting HSA API Activities Traces"
  echo -e "--roctx-trace                For Collecting ROCTx Traces"
  echo -e "--kernel-trace               For Collecting Kernel dispatch Traces"
  echo -e "--sys-trace                  For Collecting HIP and HSA APIs and their Activities Traces along ROCTX and Kernel Dispatch traces"
  echo -e "--plugin PLUGIN_NAME         For enabling a plugin (file/perfetto/att/ctf)"
  echo -e "\t# ATT plugin usage: rocprofv2 <rocprofv2_params> --plugin att <ISA_file> <att_parameters> <executable>"
  echo -e "\t# Use \"rocprofv2 --plugin att --help\" for ATT-specific parameters help."
  echo -e "-i   | --input               For adding counters file path (every line in the text file represents a counter)"
  echo -e "-o   | --output-file         For the output file name"
  echo -e "-d   | --output-directory    For adding output path where the output files will be saved"
  echo -e "-fi  | --flush-interval      For adding a flush interval in milliseconds, every \"flush interval\" the buffers will be flushed"
  echo -e "-tp  | --trace-period        Specifies a trace period in milliseconds, with format \"-tp <DELAY>:<ACTIVE_TIME>:<LOOP_INTERVAL_TIME>\"."
  exit 1
}

if [ -z "$1" ] ; then
  usage
  exit 1
fi

while [ 1 ] ; do
  if [[ "$1" = "-h" || "$1" = "--help" ]] ; then
    usage
    exit 1
  elif [[ "$1" = "-t" || "$1" = "--test" ]] ; then
    if [ $RUN_FROM_BUILD == 1 ]; then
      ./run_tests.sh
      exit 1
    fi
  elif [[ "$1" = "-mt" || "$1" = "--mem-test" ]] ; then
    if [ $RUN_FROM_BUILD == 1 ]; then
      $ROCM_DIR/tests-v2/memorytests/run_asan_tests.sh $ROCM_DIR/tests-v2/featuretests/profiler/apps/hip_vectoradd $ROCM_DIR/memleaks.log
      exit 1
    fi
  elif [[ "$1" = "--list-counters" ]] ; then
    export LD_PRELOAD=$LD_PRELOAD:$ROCM_DIR/lib/rocprofiler/librocprofiler_tool.so
    eval $ROCM_DIR/libexec/rocprofiler/ctrl
    exit 1
  elif [[ "$1" = "-i" || "$1" = "--input"  ]] ; then
    if [ $2 ] && [ -n $2 ] && [ -r $2 ] ; then
      export COUNTERS_PATH=$2
    else
      echo -e "Error: \"$2\" doesn't exist!"
      usage
      exit 1
    fi
    shift
    shift
  elif [[ "$1" = "-o" || "$1" = "--output-file-name"  ]] ; then
    if [ $2 ] ; then
      export OUT_FILE_NAME=$2
    else
      usage
      exit 1
    fi
    shift
    shift
  elif [[ "$1" = "-d" || "$1" = "--output-directory"  ]] ; then
    if [ $2 ] ; then
      OUTPUT_PATH_INTERNAL=$2
      MPI_RANK_INTERNAL=""
      if [ -n "$MPI_RANK" ]; then
        MPI_RANK_INTERNAL=$MPI_RANK
      elif [ -n "$OMPI_COMM_WORLD_RANK" ]; then
        MPI_RANK_INTERNAL=$OMPI_COMM_WORLD_RANK
      elif [ -n "$MV2_COMM_WORLD_RANK" ]; then
        MPI_RANK_INTERNAL=$MV2_COMM_WORLD_RANK
      fi
      if [ -n "$MPI_RANK_INTERNAL" ]; then
        OUTPUT_PATH_INTERNAL=${OUTPUT_PATH_INTERNAL//"%rank"/$MPI_RANK_INTERNAL}
      fi

      mkdir -p $OUTPUT_PATH_INTERNAL
      export OUTPUT_PATH=$OUTPUT_PATH_INTERNAL
    else
      usage
      exit 1
    fi
    shift
    shift
  elif [[ "$1" = "-fi" || "$1" = "--flush-interval" ]] ; then
    if [ $2 ] && [ $2 -gt 0 ] ; then
      export ROCPROFILER_FLUSH_INTERVAL=$2
    else
      echo -e "Wrong input \"$2\" for flush interval, it needs to be integer greater than zero!"
      usage
      exit 1
    fi
    shift
    shift
  elif [[ "$1" = "-tp" || "$1" = "--trace-period" ]] ; then
    if [ $2 ] && [[ "$2" == *":"* ]] ; then
      export ROCPROFILER_TRACE_PERIOD=$2
    else
      echo -e "Wrong input \"$2\" for trace period!"
      usage
      exit 1
    fi
    shift
    shift
  elif [ "$1" = "--hip-api" ] ; then
    export ROCPROFILER_HIP_API_TRACE=1
    shift
  elif [[ "$1" = "--hip-activity" || "$1" = "--hip-trace" ]] ; then
    export ROCPROFILER_HIP_API_TRACE=1
    export ROCPROFILER_HIP_ACTIVITY_TRACE=1
    shift
  elif [ "$1" = "--hsa-api" ] ; then
    export ROCPROFILER_HSA_API_TRACE=1
    shift
  elif [[ "$1" = "--hsa-activity" || "$1" = "--hsa-trace" ]] ; then
    export ROCPROFILER_HSA_API_TRACE=1
    export ROCPROFILER_HSA_ACTIVITY_TRACE=1
    shift
  elif [ "$1" = "--roctx-trace" ] ; then
    export ROCPROFILER_ROCTX_TRACE=1
    shift
  elif [ "$1" = "--kernel-trace" ] ; then
    export ROCPROFILER_KERNEL_TRACE=1
    shift
  elif [ "$1" = "--sys-trace" ] ; then
    export ROCPROFILER_HIP_API_TRACE=1
    export ROCPROFILER_HIP_ACTIVITY_TRACE=1
    export ROCPROFILER_HSA_API_TRACE=1
    export ROCPROFILER_HSA_ACTIVITY_TRACE=1
    export ROCPROFILER_ROCTX_TRACE=1
    shift
  elif [ "$1" = "--roc-sys" ] ; then
    export ROCPROFILER_ENABLE_AMDSYS=$2
    shift
    shift
  elif [ "$1" = "--plugin" ] ; then
    if [ -n $2 ] ; then
      PLUGIN=$2
      if [[ ! "${PLUGIN_LIST[*]}" =~ $PLUGIN ]]; then
      echo -e "Wrong input \"$2\" for plugin!"
      usage
      exit 1
      fi
      export ROCPROFILER_PLUGIN_LIB=lib${PLUGIN}_plugin.so
    else
      echo -e "Wrong input \"$2\" for plugin!"
      usage
      exit 1
    fi
    if [ "$2" = "att" ] ; then
      if [ $RUN_FROM_BUILD == 1 ]; then
        ATT_PATH=$ROCM_DIR/plugin/att/att/att.py
      else
        ATT_PATH=$ROCM_DIR/libexec/rocprofiler/att/att.py
        export ROCPROFV2_ATT_LIB_PATH=$ROCM_DIR/lib/hsa-amd-aqlprofile/librocprofv2_att.so
      fi
      ATT_ARGV=$3
      shift

      ATT_OPTIONS="Not done"
      ATT_PYTHON3_ARG="python3 "
      while [ "$ATT_OPTIONS" = "Not done" ]; do
        if [[ "$3" = "--trace_file" ]]; then
          ATT_ARGV="$ATT_ARGV $3 \"$4\""
          shift
          shift
        elif [[ "$3" = "--mpi" ]]; then
          ATT_PYTHON3_ARG="mpirun -np $4 python3 "
          shift
          shift
        elif [[ "$3" = "--mode" || "$3" = "--ports" || "$3" = "--genasm" || "$3" == "--att_kernel"  || "$3" == "--depth" ]]; then
          ATT_ARGV="$ATT_ARGV $3 $4"
          shift
          shift
        else
          ATT_OPTIONS="Done"
        fi
      done
    fi
    shift
    shift
  elif [[ "$1" = "-"* || "$1" = "--"* ]] ; then
    echo -e "Wrong option \"$1\", Please use the following options:\n"
    usage
    exit 1
  else
    break
  fi
done

PMC_LINES=()
if [ -n "$COUNTERS_PATH" ]; then
  input=$COUNTERS_PATH
  while IFS= read -r line || [[ -n "$line" ]]; do
    #skip empty lines
    if [[ -z "$line" ]]; then
         continue
    fi
    # if in att mode, only add the first line
    if [[ ! -n "$PMC_LINES" ]] || [[ ! -n "$ATT_ARGV" ]]; then
    PMC_LINES+=( "$line" )
    fi
  done < $input
fi

COUNTERS_PMC_DIRS=""
if [ -n "$PMC_LINES" ]; then
  COUNTER=1
  for i in ${!PMC_LINES[@]}; do
    export ROCPROFILER_COUNTERS="${PMC_LINES[$i]}"
    #Skipping lines without pmc when not in att mode
    if [[ ! ${PMC_LINES[$i]} =~ "pmc" && ! ${PMC_LINES[$i]} =~ "att" ]]; then
      continue
    fi

    if [ -n "$OUTPUT_PATH" ]; then
      if [ ! -n "$ATT_ARGV" ]; then
        FINAL_PATH="$OUTPUT_PATH_INTERNAL/pmc_$COUNTER"
        COUNTERS_PMC_DIRS="$COUNTERS_PMC_DIRS $FINAL_PATH"
      else
        FINAL_PATH="$OUTPUT_PATH"
      fi
      echo -e "\nThe output path for the following counters: $FINAL_PATH"
      mkdir -p $FINAL_PATH
      echo $ROCPROFILER_COUNTERS > $FINAL_PATH/pmc.txt
      export OUTPUT_PATH=$FINAL_PATH
      let COUNTER=COUNTER+1
    fi
    LD_PRELOAD=$LD_PRELOAD:$ROCM_DIR/lib/rocprofiler/librocprofiler_tool.so $*
  done
elif [ ! -n "$ATT_ARGV" ]; then
  LD_PRELOAD=$LD_PRELOAD:$ROCM_DIR/lib/rocprofiler/librocprofiler_tool.so $*
fi

get_pmc_results_txt_path(){
  for file_name in `ls $1` ; do
    if [[ $file_name == *results.txt ]]; then
      echo "$1/$file_name"
    fi
  done
}

if [ -n "$ATT_PATH" ]; then
  if [ -n "$ATT_ARGV" ]; then
    eval "$ATT_PYTHON3_ARG $ATT_PATH $ATT_ARGV"
  elif [ ! -n "$PMC_LINES" ]; then
    echo "ATT File  is required!"
  fi
fi