#!/usr/bin/env bash
# Copyright 2024, 2026 The Torch-Spyre Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

top="$(dirname $(readlink -m ${BASH_SOURCE[0]}))"
iam=$(basename $0)

usage()
{
   echo "Usage: $iam [cmd] [options]"
   echo ""
   echo "Commands:"
   echo "  dmon - Monitor device metrics (default)"
   echo ""
   usage_dmon 1

   exit 2
}

# Supported python versions
pythonVerMax=${PYTHON_VER_MAX:-15}  # 'Feature' release as of June, 2025
pythonVerMin=${PYTHON_VER_MIN:-10}  # Required for 'match' statement

if (command -v "${PYTHON_CMD}" > /dev/null); then
   pythonCmd=$(command -v "${PYTHON_CMD}")  # Assume python 3.10 or higher
else
   pythonCmd=python3
fi
pythonCfg="${pythonCmd}-config"

libDir=$top/../lib
instDirs=( /opt/ibm/spyre/aiu-monitor \
           /opt/ibm/spyre/runtime )

cmdDefault=dmon
#metricFile=${AIUSMI_METRIC_FILE:-'/tmp/metrics'}

if [[ -n "$AIUSMI_METRIC_FILE" ]]; then
       metricFile="$AIUSMI_METRIC_FILE"
elif [[ -n "$SPYRE_METRIC_PATH" ]]; then
       metricFile="$SPYRE_METRIC_PATH"
else
       metricFile='/tmp/metrics.%BUSID'
fi

usage_dmon()
{
   echo "Usage: $iam [dmon] [-s] [-i <device-ids>] [-g <metric-list>] [-d <delay>] [-f <output-file>]"
   echo "                   [--idle-power <watts>] [--llm-type <model-type>] [-v] [<metric-files>...]"
   echo ""
   echo "options:"
   echo "  -s | --csv:           Output csv format"
   echo "  -i | --id:            Comma-separated device IDs to monitor (e.g., 0,1,2) [default=all]"
   echo "  -g | --metric-groups: Select metric groups to print [default="DP"]:"
   echo "                          D=device Data rate, U=actual and peak device memory Usage (torch-spyre only),"
   echo "                          P=Pt_act, A=All metrics."
   echo "                        Always print dev. ID, timestamp, host cpu/mem, power, temperature, and busy."
   echo "  -d | --delay:         Console summary interval in seconds [default=1sec]. Accept floating-point value."
   echo "                        Minimum is 0.1 sec on s390x console or VF mode, otherwise 0.01 sec."
   echo "  -f | --filename:      Log to a specified file, rather than to stdout"
   echo "  --idle-power          Idle power baseline (W) for pt_act estimation [default=20.0]"
   echo "  --llm-type            LLM model type for pt_act estimation: encoder|decoder [default=decoder]"
#   echo "  -r | --remove-old:    Remove old metric file before reading metrics"
#   echo "                        Note: the metric file is also removed in the beginning of this command"
#   echo "                              if there is no process that opens the metric file"
   echo "  -v | --version:       Show version"
   echo "  -h | -? | --help:     Show help message"
   [[ $# -ge 1 ]] || exit 1
}


check_extension()
{
   if [[ -x "$pythonCmd" ]]; then
      pythonCfg="${pythonCmd}-config"
      if [[ -x "$pythonCfg" ]]; then
         xsuf=$( ${pythonCfg} --extension-suffix )
         if [[ -f "$libDir/$1$xsuf" ]]; then
            return
         fi
      fi
   fi
   if [[ -n "$PYTHON_CMD" && -x "$pythonCfg" ]]; then
      if (command -v "${PYTHON_CMD}" > /dev/null); then
         echo "Warning: No $1 extension available for $PYTHON_CMD. Ignored"
      else
         echo "Warning: $PYTHON_CMD is not available. Ignored"
      fi
   fi

   # Search python 3.10 or higher whose extention is available
   pythonCmd=python3
   for r in $(seq $pythonVerMax -1 $pythonVerMin ); do
      if (command -v python3.${r} > /dev/null); then
         pythonCmd=$(command -v python3.${r})
         pythonCfg="${pythonCmd}-config"
         if [[ -x "$pythonCfg" ]]; then
            xsuf=$( ${pythonCfg} --extension-suffix 2>/dev/null )
            if [[ -f "$libDir/$1$xsuf" ]]; then
               return
            fi
         fi
      fi
   done

   if (command -v "${PYTHON_CMD}" > /dev/null); then
      pythonCmd=$(command -v "${PYTHON_CMD}")
   else
      pythonCmd=$(command -v python3 2>/dev/null)
   fi
   pythonCfg="${pythonCmd}-config"

   if [[ ! -x "$pythonCmd" ]]; then
       echo "ERROR: No python3 command found"
       exit 1
   fi
}

do_dmon()
{
   dcr_opts=''
   while [[ $# -gt 0 ]]; do
      case "$1" in
      # Compatible to nvidia-smi
      '-i'|'--id')            [[ $# -ge 2 ]] || usage_dmon; dcr_opts+=" --device-ids $2"; shift;;
      '-d'|'--delay')         [[ $# -ge 2 ]] || usage_dmon; dcr_opts+=" -i $2"; shift;;
      '-f'|'--filename')      [[ $# -ge 2 ]] || usage_dmon; dcr_opts+=" -f $2"; shift;;

      # aiu_smi_main options usable in aiu-smi
      '-r'|'--remove-old')    dcr_opts+=' --remove-old';;
      '-s'|'--csv')           dcr_opts+=' --csv';;
      '-g'|'--metric-groups') [[ $# -ge 2 ]] || usage_dmon; dcr_opts+=" -g $2"; shift;;

      # aiu_smi_main options usable in aiu-smi (hidden from the help message)
      '--idle-power')               [[ $# -ge 2 ]] || usage_dmon; dcr_opts+=" --idle-power $2"; shift;;
      '--llm-type')                 [[ $# -ge 2 ]] || usage_dmon; dcr_opts+=" --llm-type $2"; shift;; # <-- NEW: Pass LLM Type
      '--metric-file-stale-delay')  [[ $# -ge 2 ]] || usage_dmon; dcr_opts+=" --metric-file-stale-delay $2"; shift;;

      # show version and exit. Print python pacakge versoin
      '-v'|'--version')  dcr_opts+=" --version";;

      # usage and other stuff
      '-h'|'-?'|'--help')  usage_dmon;;
      '--')                shift; break;;

      # Unsupported nvidia-smi options
      -*)  echo "ERROR: unknown option: $1"
           usage_dmon;;
      *)   break;;
      esac
      shift
   done

   # Detect environment:
   if [[ -f "$top/../lib/aiu_smi_main.py" ]]; then
      # Installed from Wheel or Development environment: files are in ../lib
      parserDir="$top/../lib"
   elif [[ -f "$top/aiu_smi_main.py" ]]; then
      # Installed from this repo: files are in the same directory
      parserDir="$top"
   else
      echo "ERROR: Cannot find 'aiu_smi_main.py'"
      exit 1
   fi

   if [[ "x$metricFile" = x*'%BUSID'* ]]; then
      if [[ -n "$PCIDEVICE_IBM_COM_AIU_PF" ]]; then
         declare -a busids=(${PCIDEVICE_IBM_COM_AIU_PF//,/ })
      elif [[ -n "$AIU_WORLD_SIZE" ]]; then
         declare -a busids
         for i in $( seq 0 $(($AIU_WORLD_SIZE - 1)) ); do
            busids+=($(eval echo "\$AIU_WORLD_RANK_${i}"))
         done
      fi
      declare -a files
      for id in "${busids[@]}"; do
         func=$(echo ${id##*:}. | cut -d. -f2)
         files+=(${metricFile//%BUSID/${id}})
      done
      if [[ "${#files[@]}" -eq 0 && " $dcr_opts " != *" --version "* ]]; then
         echo "ERROR: Failed to expand '%BUSID' to actual bus IDs"
         exit 1
      fi
      metricFile="${files[*]}"
   fi
   for f in $metricFile; do
      if [[ -f "$f" ]]; then
         if !(find /proc/*/fd -type l | xargs ls -l | grep -q " $f"'$') 2>/dev/null; then
            rm -f "$f"
         fi
      fi
   done

   # Check if extension is available, and set 'pythonCmd' and 'pythonCfg'
   check_extension libaiumonitor

   # Lookup candidates of installation target in the system
   libs=""
   if [[ -x "${pythonCfg}" ]]; then
      xsuf=$( ${pythonCfg} --extension-suffix )
      for d in "${instDirs[@]}"; do
         [[ -f "${d}/lib/libaiumonitor${xsuf}" ]] && libs="$libs:$d/lib"
      done
   fi
   libs="${libDir}${libs}"

   env PYTHONPATH="${libs}:$PYTHONPATH" \
      ${pythonCmd:-python3} -u $parserDir/aiu_smi_main.py $dcr_opts $metricFile "$@"
}


##
## Main starts here
##
declare -a opts
while [[ $# -gt 0 ]]; do
   case "$1" in
   dmon)  cmd="$1"; shift; break;;
   -*)    opts+=("$1");;
   *)     break;;
   esac
   shift
done

cmd="${cmd:-$cmdDefault}"

case "$cmd" in
dmon)   do_dmon $AIUSMI_OPTS "${opts[@]}" "$@";;
*)      echo "ERROR: Unknown command: $cmd"
        usage;;
esac
