getusage

#!/bin/bash

show_usage() {
  echo "Usage:"
  echo "  getusage --account <slurm_account_name>"
  echo "  getusage --pi <VT_PI_pid> (PID/username of the VT PI)"
  echo "  getusage --user <VT_pid> (PID/username of the VT user)"
  echo "Optional:"
  echo "  --cluster <cluster_name> (tinkercliffs, owl, falcon)"
  exit 0
}

# Parse arguments
ACCOUNT=""
PI=""
USER=""
CLUSTER=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --account)
      if [[ -n "$2" && "$2" != --* ]]; then
        ACCOUNT="$2"
        shift 2
      else
        show_usage
      fi
      ;;
    --pi)
      if [[ -n "$2" && "$2" != --* ]]; then
        PI="$2"
        shift 2
      else
        show_usage
      fi
      ;;
    --user)
      if [[ -n "$2" && "$2" != --* ]]; then
        USER="$2"
        shift 2
      else
        show_usage
      fi
      ;;
    --cluster)
      if [[ -n "$2" && "$2" != --* ]]; then
        case "$2" in
          tinkercliffs|owl|falcon)
            CLUSTER="$2"
            shift 2
            ;;
          *)
            echo "ERROR: Invalid cluster name '$2'. Must be one of: tinkercliffs, owl, falcon."
            exit 1
            ;;
        esac
      else
        show_usage
      fi
      ;;
    *)
      show_usage
      ;;
  esac
done

# Validate that only one of --account, --pi, or --user is provided
num_args_set=0
[[ -n "$ACCOUNT" ]] && ((num_args_set++))
[[ -n "$PI" ]] && ((num_args_set++))
[[ -n "$USER" ]] && ((num_args_set++))

if [[ $num_args_set -ne 1 ]]; then
  show_usage
fi

getParentAccountUsage() {  
  local sshare_output
  sshare_output=$(sshare -a -A "$1" -M "$2" --long --format=Account,User,GrpTRESMins,RawUsage,TRESRunMins -nP 2>/dev/null)

  # Split into lines and filter out "CLUSTER:" lines
  IFS=$'\n' read -rd '' -a raw_lines <<< "$sshare_output"
  lines=()
  for line in "${raw_lines[@]}"; do
    [[ "$line" =~ ^CLUSTER:\  ]] && continue
    lines+=("$line")
  done

  # Initialize totals
  local grp_mins=0
  local used_mins=0
  local run_mins=0

  # Aggregate data
  for line in "${lines[@]}"; do
    this_grp_mins=$(echo "$line" | cut -d'|' -f3 | sed -n 's/.*billing=\([0-9]*\).*/\1/p')
    this_used_mins=$(echo "$line" | cut -d'|' -f4)
    this_run_mins=$(echo "$line" | cut -d'|' -f5 | sed -n 's/.*billing=\([0-9]*\).*/\1/p')

    grp_mins=$((${this_grp_mins:-0}))  # Overwrites with latest total
    used_mins=$((used_mins + ${this_used_mins:-0}))
    run_mins=$((run_mins + ${this_run_mins:-0}))
  done

  local total_hrs used_hrs run_hrs
  total_hrs=$(awk "BEGIN {printf \"%.2f\", $grp_mins/60}")
  used_hrs=$(awk "BEGIN {printf \"%.2f\", $used_mins/3600}")
  run_hrs=$(awk "BEGIN {printf \"%.2f\", $run_mins/60}")

  echo "$total_hrs $used_hrs $run_hrs"
}

getAccountUsage() {
  # Only run if $1 does NOT end with '-paid'
  if [[ "$1" != *-paid ]]; then
    # Parent account
    parent=$(sacctmgr show assoc where account="$1" format=parentname -np cluster="$2" | grep -v "^|" | tr -d "|")
   
    if [[ "$parent" == "" ]]; then
      echo "ERROR: Account $1 account doesn't exist"
      exit 0
    fi
    
    read parent_total_hrs parent_used_hrs parent_run_hrs <<< "$(getParentAccountUsage $parent $2)"
    parent_avail_hrs=$(awk "BEGIN {diff = $parent_total_hrs - $parent_used_hrs - $parent_run_hrs; printf \"%.2f\", (diff > 0 ? diff : 0)}")
    parent_avail_pct=$(awk "BEGIN { 
        if ($parent_total_hrs == 0) 
            printf \"0.0\"; 
        else 
            printf \"%.1f\", ($parent_avail_hrs/$parent_total_hrs)*100 
    }")
  fi

  # Get users to exclude (those with '0' in the second field)
  local exclude_users=()
  while IFS='|' read -r user field2 _; do
    if [[ -z "$USER" ]]; then
      # Original logic
      if [[ "$user" == "" && "$field2" == "0" ]]; then
        echo "ERROR: The account $ACCOUNT is not active"
        return
      elif [[ "$field2" == "0" ]]; then
        exclude_users+=("$user")
      fi
    else
      # Exclude everyone except $USER
      if [[ "$user" != "$USER" ]]; then
        exclude_users+=("$user")
      fi
    fi
  done < <(sacctmgr show assoc where account=$ACCOUNT format=user,grpsubmi -np | sort | uniq)

  # Optional: build a map of excluded users (safe against empty values)
  declare -A exclude_map
  for u in "${exclude_users[@]}"; do
    [[ -n "$u" ]] && exclude_map["$u"]=1
  done

  local sshare_output
  sshare_output=$(sshare -a -A "$1" -M "$2" --long --format=Account,User,GrpTRESMins,RawUsage,TRESRunMins -nP 2>/dev/null)

  # Split into lines and filter out "CLUSTER:" lines
  local raw_lines lines line
  IFS=$'\n' read -rd '' -a raw_lines <<< "$sshare_output"
  lines=()
  for line in "${raw_lines[@]}"; do
    [[ "$line" =~ ^CLUSTER:\  ]] && continue
    lines+=("$line")
  done

  # Initialize totals
  local acct_name=""
  local grp_mins=0
  local used_mins=0
  local run_mins=0

  declare -A user_used_mins
  declare -A user_run_mins
 
  # Aggregate data
  for line in "${lines[@]}"; do
    local acct_field user_field this_grp_mins this_used_mins this_run_mins
    acct_field=$(echo "$line" | cut -d'|' -f1)
    user_field=$(echo "$line" | cut -d'|' -f2)

    if [[ -z "$user_field" ]]; then
      acct_name="$acct_field"
      this_grp_mins=$(echo "$line" | cut -d'|' -f3 | sed -n 's/.*billing=\([0-9]*\).*/\1/p')
      this_used_mins=$(echo "$line" | cut -d'|' -f4)
      this_run_mins=$(echo "$line" | cut -d'|' -f5 | sed -n 's/.*billing=\([0-9]*\).*/\1/p')

      grp_mins=$((${this_grp_mins:-0}))
      used_mins=$((used_mins + ${this_used_mins:-0}))
      run_mins=$((run_mins + ${this_run_mins:-0}))
    else
      local user="$user_field"
      local u_used_mins u_run_mins
      u_used_mins=$(echo "$line" | cut -d'|' -f4)
      u_run_mins=$(echo "$line" | cut -d'|' -f5 | sed -n 's/.*billing=\([0-9]*\).*/\1/p')

      user_used_mins["$user"]=$(( ${user_used_mins["$user"]:-0} + ${u_used_mins:-0} ))
      user_run_mins["$user"]=$(( ${user_run_mins["$user"]:-0} + ${u_run_mins:-0} ))
    fi
  done

  # Compute hours as floats  
  local total_hrs used_hrs run_hrs avail_hrs
  total_hrs=$(awk "BEGIN {printf \"%.2f\", $grp_mins/60}")
  used_hrs=$(awk "BEGIN {printf \"%.2f\", $used_mins/3600}")
  run_hrs=$(awk "BEGIN {printf \"%.2f\", $run_mins/60}")
  avail_hrs=$(awk "BEGIN {diff = $total_hrs - $used_hrs - $run_hrs; printf \"%.2f\", (diff > 0 ? diff : 0)}")

  # Compute percentages
  local used_pct avail_pct
  used_pct=$(awk -v t="$total_hrs" -v u="$used_hrs" 'BEGIN {
      if (t == 0) printf "0.0";
      else printf "%.1f", (u/t)*100
  }')

  avail_pct=$(awk -v t="$total_hrs" -v a="$avail_hrs" 'BEGIN {
      if (t == 0) printf "0.0";
      else printf "%.1f", (a/t)*100
  }')

  echo "  Account: $acct_name"
  printf "    Account total:     %'12.2f\n" "$total_hrs"
  printf "    Account running:   %'12.2f --> Service Units committed for completing running jobs\n" "$run_hrs"
  printf "    Account used:      %'12.2f (%s%%)\n" "$used_hrs" "$used_pct"
  printf "    Account available: %'12.2f (%s%%)\n" "$avail_hrs" "$avail_pct"
  if [[ "$1" != *-paid ]]; then
    actual_available=$(awk "BEGIN {
    parent = $parent_total_hrs - $parent_used_hrs - $parent_run_hrs;
    child = $total_hrs - $used_hrs - $run_hrs;
    printf \"%.2f\", (parent < child) ? $parent_avail_hrs : $avail_hrs;
    }")

    if [[ "$avail_hrs" != "$actual_available" ]]; then
      availability_warning="Warning: availability of Service Units in this account is limited by the availability of the PI's"
    else
      availability_warning=""
    fi

    printf "    PI available:      %'12.2f (%s%%) out of %'.2f Service Units shared among all PI's accounts\n" "$parent_avail_hrs" "$parent_avail_pct" "$parent_total_hrs"
    printf "    Actual available:  %'12.2f %s\n" "$actual_available" "$availability_warning"
  fi
  echo  

  # Per-user breakdown
  echo "    User breakdown:"
  for user in "${!user_used_mins[@]}"; do
    if [[ -n "${exclude_map["$user"]}" ]]; then
      continue
    fi

    local u_used_mins u_run_mins u_used_hrs u_run_hrs
    u_used_mins=${user_used_mins["$user"]}
    u_run_mins=${user_run_mins["$user"]:-0}

    u_used_hrs=$(awk "BEGIN {printf \"%.2f\", $u_used_mins/3600}")
    u_run_hrs=$(awk "BEGIN {printf \"%.2f\", $u_run_mins/60}")

    printf "      %-15s Used: %'12.2f h | Running: %'12.2f h\n" "$user" "$u_used_hrs" "$u_run_hrs"
  done
}

# Perform action based on input
if [[ -n "$ACCOUNT" ]]; then
  if [ "$ACCOUNT" == "personal" ]; then
    echo "Personal accounts are restricted for testing during the first 90 days. Users need to be added to an allocation owned by a PI"
    exit 0
  fi
      
  for cluster in tinkercliffs owl falcon; do
      if [[ -n "$CLUSTER" && "$CLUSTER" != "$cluster" ]]; then
          continue
      fi
      echo "Cluster: $cluster"
      getAccountUsage "${ACCOUNT}" "$cluster"
      echo
  done
elif [[ -n "$USER" ]]; then
  accounts=$(sacctmgr show assoc where user=${USER} format=account -nP cluster=tinkercliffs)
  for cluster in tinkercliffs owl falcon; do
      if [[ -n "$CLUSTER" && "$CLUSTER" != "$cluster" ]]; then
          continue
      fi
      echo "Cluster: $cluster"
      for ACCOUNT in $accounts; do
        if [ "$ACCOUNT" == "personal" ]; then
          continue
        fi
        getAccountUsage "$ACCOUNT" $cluster
        echo
      done
      echo
  done
elif [[ -n "$PI" ]]; then
  for cluster in tinkercliffs owl falcon; do
    if [[ -n "$CLUSTER" && "$CLUSTER" != "$cluster" ]]; then
            continue
    fi
    echo "Cluster: $cluster"
    ## Free accounts
    read total_hrs used_hrs run_hrs <<< "$(getParentAccountUsage __${PI}__ $cluster)"  
    
    if [[ "$total_hrs" == "0.00" ]]; then
        echo "ERROR: PI ${PI} doesn't exist"
        exit 0
    fi

    avail_hrs=$(awk "BEGIN {diff = $total_hrs - $used_hrs - $run_hrs; printf \"%.2f\", (diff > 0 ? diff : 0)}")
    used_pct=$(awk "BEGIN {
        if ($total_hrs == 0)
            printf \"0.0\";
        else
            printf \"%.1f\", ($used_hrs/$total_hrs)*100
    }")

    avail_pct=$(awk "BEGIN {
        if ($total_hrs == 0)
            printf \"0.0\";
        else
            printf \"%.1f\", ($avail_hrs/$total_hrs)*100
    }")


    echo "  PI: $PI"
    printf "    Free-tier allocation\n"
    printf "    PI total:     %'12.2f\n" "$total_hrs"
    printf "    PI running:   %'12.2f --> Service Units committed for completing running jobs\n" "$run_hrs"
    printf "    PI used:      %'12.2f (%s%%)\n" "$used_hrs" "$used_pct"
    printf "    PI available: %'12.2f (%s%%) --> Service Units shared among all PI's accounts\n" "$avail_hrs" "$avail_pct"
    echo

    children=$(sacctmgr show assoc format=account,parentname,grpsubmit -np cluster=tinkercliffs | awk -F'|' -v parentname="__${PI}__" '$2 == parentname && $3 != 0 {print $1}')

    for ACCOUNT in $children; do
        if [ "$ACCOUNT" == "personal" ]; then
          continue
        fi
        getAccountUsage "$ACCOUNT" $cluster
        echo
    done
    echo
  done
fi