#!/bin/bash -e

set -eo pipefail

ROCPROFV3_DIR=$(dirname -- "$(realpath "${BASH_SOURCE[0]}")")
ROCM_DIR=$(dirname -- "${ROCPROFV3_DIR}")

: ${HSA_TOOLS_LIB:="${ROCM_DIR}/lib/librocprofiler-sdk.so"}

export HSA_TOOLS_LIB

# Define color codes
GREEN='\033[0;32m'
GREY='\033[0;90m'
RESET='\033[0m'

usage() {
  local EC=${1}
  if [ -z "${EC}" ]; then EC=1; fi
  echo -e "${RESET}ROCProfilerV3 Run Script Usage:"
  echo -e "${GREEN}-h   | --help ${RESET}                For showing this message"
  #   echo -e "${GREEN}-i   | --input ${RESET}               For adding counters file path (every line in the text file represents a counter)"
  echo -e "${GREEN}--hsa-trace ${RESET}                  For Collecting HSA API Traces"
  echo -e "${GREEN}--kernel-trace ${RESET}               For Collecting Kernel Dispatch Traces"
  echo -e "${GREEN}--memory-copy-trace ${RESET}          For Collecting Memory Copy Traces"
  echo -e "${GREEN}--marker-trace ${RESET}               For Collecting Marker (ROCTx) Traces"
  echo -e "${GREEN}--hip-trace ${RESET}                  For Collecting HIP Runtime Traces"
  echo -e "${GREEN}--hip-compiler-trace ${RESET}         For Collecting HIP Compiler generated code Traces"
  echo -e "${GREEN}--sys-trace ${RESET}                  For Collecting HIP,HSA, Memory Copy, (marker)ROCTx, Kernel dispatch and HIP compiler traces\n"
  echo -e "${GREEN}-o   | --output-file ${RESET}         For the output file name"
  echo -e "\t#${GREY} usage e.g:(with current dir): rocprofv3 --hsa-trace -o <file_name> <executable>"
  echo -e "\t#${GREY} usage e.g:(with custom dir):  rocprofv3 --hsa-trace -d <out_dir> -o <file_name> <executable>${RESET}\n"
  echo -e "${GREEN}-d   | --output-directory ${RESET}    For adding output path where the output files will be saved"
  echo -e "\t#${GREY} usage e.g:(with custom dir):  rocprofv3 --hsa-trace -d <out_dir> <executable>${RESET}"
  echo -e "${GREEN}-M   | --mangled-kernels ${RESET}     Do not demangle the kernel names"
  echo -e "${GREEN}-T   | --truncate-kernels ${RESET}    Truncate the demangled kernel names"
  echo -e ""
  exit ${EC}
}

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

: ${ROCPROF_OUTPUT_PATH:="."}
: ${ROCPROF_OUTPUT_PATH_INTERNAL:="."}
: ${ROCPROF_OUTPUT_FILE_NAME:=""}
: ${ROCPROF_COUNTERS_PATH:=""}
: ${ROCPROF_PRELOAD:=""}

while true; do
  if [[ "$1" == "-h" || "$1" == "--help" ]]; then
    usage 0
  elif [[ "$1" == "-M" || "$1" == "--mangled-kernels" ]]; then
    export ROCPROF_DEMANGLE_KERNELS=0
    shift
  elif [[ "$1" == "-T" || "$1" == "--truncate-kernels" ]]; then
    export ROCPROF_TRUNCATE_KERNELS=1
    shift
  elif [[ "$1" == "-i" || "$1" == "--input" ]]; then
    if [ "$2" ] && [ -n "$2" ] && [ -r "$2" ]; then
      export ROCPROF_COUNTERS_PATH=$2
      export ROCPROF_COUNTER_COLLECTION=1
    else
      echo -e "Error: input file \"$2\" doesn't exist!"
      usage 1
    fi
    shift
    shift
  elif [[ "$1" == "-o" || "$1" == "--output-file-name" ]]; then
    if [ "$2" ]; then
      export ROCPROF_OUTPUT_FILE_NAME=$2
    else
      usage 1
    fi
    shift
    shift
  elif [[ "$1" == "-d" || "$1" == "--output-directory" ]]; then
    if [ "$2" ]; then
      ROCPROF_OUTPUT_PATH_INTERNAL=$2
      export ROCPROF_OUTPUT_PATH=$ROCPROF_OUTPUT_PATH_INTERNAL
    else
      usage 1
    fi
    shift
    shift
  elif [ "$1" == "--hsa-trace" ]; then
    export ROCPROF_HSA_API_TRACE=1
    shift
  elif [ "$1" == "--kernel-trace" ]; then
    export ROCPROF_KERNEL_TRACE=1
    shift
  elif [ "$1" == "--memory-copy-trace" ]; then
    export ROCPROF_MEMORY_COPY_TRACE=1
    shift
  elif [ "$1" == "--marker-trace" ]; then
    export ROCPROF_MARKER_API_TRACE=1
    shift
  elif [ "$1" == "--hip-trace" ]; then
    export ROCPROF_HIP_API_TRACE=1
    shift
  elif [ "$1" == "--hip-compiler-trace" ]; then
    export ROCPROF_HIP_COMPILER_API_TRACE=1
    shift
  elif [ "$1" == "--sys-trace" ]; then
    export ROCPROF_HSA_API_TRACE=1
    export ROCPROF_KERNEL_TRACE=1
    export ROCPROF_MEMORY_COPY_TRACE=1
    export ROCPROF_MARKER_API_TRACE=1
    export ROCPROF_HIP_API_TRACE=1
    shift
  elif [ "$1" == "--" ]; then
    shift
    break
  elif [[ "$1" == "-"* || "$1" == "--"* ]]; then
    echo -e "Wrong option \"$1\", Please use the following options:\n"
    usage 1
  else
    break
  fi
done

# read input counter file
PMC_LINES=()
if [ -n "$ROCPROF_COUNTERS_PATH" ]; then
  input=$ROCPROF_COUNTERS_PATH
  while IFS= read -r line || [[ -n "$line" ]]; do
    #skip empty lines
    if [[ -z "$line" ]]; then
      continue
    fi
    PMC_LINES+=("$line")
  done <"$input"
fi

if [ -n "${PMC_LINES:-}" ]; then
  #for counter collection
  COUNTER=1
  for i in "${!PMC_LINES[@]}"; do
    export ROCPROF_COUNTERS="${PMC_LINES[$i]}"
    if [[ ! ${PMC_LINES[$i]} =~ "pmc" ]]; then
      continue
    fi

    RESULT_PATH="$ROCPROF_OUTPUT_PATH_INTERNAL/pmc_$COUNTER"
    if [ -n "$ROCPROF_OUTPUT_FILE_NAME" ] || [ -n "$ROCPROF_OUTPUT_PATH" ]; then
      export ROCPROF_OUTPUT_PATH=$RESULT_PATH
    fi
    ((COUNTER++))
    ROCP_TOOL_LIBRARIES="${ROCM_DIR}/lib/rocprofiler-sdk/librocprofiler-sdk-tool.so" \
      LD_LIBRARY_PATH=${ROCM_DIR}/lib:${LD_LIBRARY_PATH} \
      LD_PRELOAD="${ROCPROF_PRELOAD}:${ROCM_DIR}/lib/librocprofiler-sdk.so" \
      "${@}"
    if [ -n "$ROCPROF_OUTPUT_PATH" ]; then
      echo -e "\nThe output path for the following counters: $ROCPROF_OUTPUT_PATH"
    fi
  done
else
  # for non counter collection. e.g: tracing
  ROCP_TOOL_LIBRARIES="${ROCM_DIR}/lib/rocprofiler-sdk/librocprofiler-sdk-tool.so" \
    LD_LIBRARY_PATH=${ROCM_DIR}/lib:${LD_LIBRARY_PATH} \
    LD_PRELOAD="${ROCPROF_PRELOAD}:${ROCM_DIR}/lib/librocprofiler-sdk.so" \
    "${@}"
fi