#!/usr/bin/env bash

#
# Copyright (C) 2010 Charles E. Vejnar
#
# This is free software, licensed under the GNU General Public License v3.
# See /LICENSE for more information.
#
set -u

mu_pid=$$
pr_name=${1%% *}
pr_maxmem=0

to_hms () {
    seconds=$1
    hours=$((seconds / 3600))
    seconds=$((seconds % 3600))
    minutes=$((seconds / 60))
    seconds=$((seconds % 60))
    printf "%02d:%02d:%02d" $hours $minutes $seconds
}

get_last_pid() {
    parent=$1
    dosearch=1
    newid="$parent"
    p=$(pgrep -P $parent -n)
    [ ${#p} -eq 0 ] && dosearch=0
    while [[ $dosearch -eq 1 ]]; do
	    p=$(pgrep -P $parent -n)
	    if [ ${#p} -eq 0 ]; then
	        dosearch=0
	        newid=$parent
	    else
	        parent=$p
	    fi
    done
    echo $newid
}

calc(){ awk "BEGIN{ print $* }" ;}

usage() {
    echo "Usage: $0 [ \"command\" | PID ]"
    echo "If command is given, the memory usage is monitored until the process ends."
    echo "It then prints out stats on memory and time usage."
    echo
    echo "If a PID is given, then the memory usage is output to the stdout every $dt seconds."
    echo
    exit 1
}

get_hugepage() {
    local A=($(grep '^Huge' /proc/meminfo | cut -d ':' -f 2))
    local htotal=${A[0]}
    local hfree=${A[1]}
    local hused=$((htotal-hfree))
    echo $hused
}

# Check command line
[[ $# -ge 1 ]] || { usage; exit 0; }

script="$1"
shift
args="$*"
pr_pid="NONE"
eval_pid=0

date_start=$(date +%s)
argispid=0
uid="$(whoami 2>/dev/null)"
[[ -n $uid ]] || uid="uid$(id -u)"

if [[ ! -x $script ]]; then
    # check if argument is a pid
    printf "%d" "$script" >& /dev/null || { echo "error: invalid argument - must be an executable file or a PID"; usage; exit 255; }
    pr_pid=$script
    ps -p $pr_pid >& /dev/null
    [[ $? -eq 0 ]] || { echo "error: no process with given PID <$pr_pid>"; exit 255; }
    psinfo=$(ps -p $pr_pid --no-header)
    echo "#------------------------------------------------------------------------------"
    echo "# pid     : "$psinfo
    echo "# date    : $(date)"
    echo "# user    : $uid"
    echo "# machine : $(uname -a)"
    echo "#------------------------------------------------------------------------------"
    argispid=1
    eval_pid=$pr_pid
else
    # Launch script/command and get its PID
    echo "#------------------------------------------------------------------------------"
    echo "# command : $script $args"
    echo "# date    : $(date)"
    echo "# user    : $uid"
    echo "# machine : $(uname -a)"
    echo "#------------------------------------------------------------------------------"
    
    
    # Launch the program
    eval $script $args &
    eval_pid=$!

    # eval_pid contains the parent pid
    pr_pid=$(get_last_pid $eval_pid)
fi

[[ "$pr_pid" == "NONE" ]] && { echo "warning: no process to watch - probably already done"; exit 0; }

pr_pid=$(get_last_pid $eval_pid)

echo "monitoring PID = $pr_pid"
# Wait for it to end
let t=0
let dt=2
fmt='MEM: %6s  %12s  %16s  %16s  %16s  %16s %s'
while [[ "$pr_pid" != "NONE" ]]; do
    data=($(ps --no-headers -o comm -o vsz -o rss -o pcpu -p $pr_pid))
    comm=${data[0]}
    vmem=$(printf '%.1f' $(calc "${data[1]}/1024"))
    rmem=$(printf '%.1f' $(calc "${data[2]}/1024"))
    pcpu=${data[3]}
    #label=$(ps -p $pr_pid -o comm= -o pcpu=)
    #[[ $argispid -eq 0 ]] || 
    hpages=$(get_hugepage)
    printf "$fmt\n" $t $pr_pid $vmem $rmem $pcpu $hpages $comm
    #echo -e "MEM: $t\t$pr_pid\t${data[@]}"
    
    ###[[ $mem -gt $pr_maxmem ]] && pr_maxmem=$mem
    # sleep a bit
    sleep $dt
    t=$((t+dt))
    pr_pid=$(get_last_pid $eval_pid)
    #[[ $argispid -eq 0 ]] && pr_pid=$(get_last_pid $eval_pid)
done

# Get exit status
wait $eval_pid
exitstatus=$?
# Output
date_end=$(date +%s)
hms=$(to_hms $(($date_end-$date_start)))
echo "# RESMON: Mem max ${pr_maxmem} MB"
echo "# RESMON: Time ${hms}"
echo "# RESMON: Exit status ${exitstatus}"
