#!/bin/bash
# Detect resources of a list of hosts and print OAR commands to create them in OAR database.

while getopts "c:e:p:n:o:HTh" opt; do
    case $opt in
    c)
        OARCONFFILE=$OPTARG
        ;;
    e)
        OPENSSH_CMD=$OPTARG
        ;;
    n)
        OARNODESETTING_CMD=$OPTARG
        ;;
    p)
        OARPROPERTY_CMD=$OPTARG
        ;;
    o)
        GENERATED_COMMANDS_FILE=$OPTARG
        ;;
    H)
        WITH_HT=1
        ;;
    T)
        WITH_HT=1
        CREATE_THREAD_PROPERTY=1
        ;;
    h)
        cat <<EOF
Usage:
    $0 <hosts file> [options]
Options:
    -o  output file for generated commands
    -e  ssh command to connect nodes
    -H  handle HyperThreading
    -T  handle HyperThreading + create thread property
    -n  path to oarnodesetting
    -p  path to oarproperty
    -c  path to oar.conf
    -h  display this message

EOF
        exit 0
        ;;
    \?)
        echo "Invalid option: -$OPTARG" >&2
        exit 1
      ;;
  esac
done
shift $((OPTIND-1))
OARCONFFILE=${OARCONFFILE:-/etc/oar/oar.conf}
OPENSSH_CMD=${OPENSSH_CMD:-/usr/bin/ssh}
OARPROPERTY_CMD=${OARPROPERTY_CMD:-oarproperty}
OARNODESETTING_CMD=${OARNODESETTING_CMD:-oarnodesetting}
GENERATED_COMMANDS_FILE=${GENERATED_COMMANDS_FILE:-/tmp/oar_resources_init.cmd}

. "$OARCONFFILE" || exit 20

# check the file with the nodes to add
if [ -z "$1" ] || [ ! -r "$1" ]; then
    echo "Please provide as argument a file containing the list of nodes to add to OAR"
    echo "resources (this file must be readable by user oar)"
    exit 1
fi

echo "Did you configured the OAR SSH key on all the nodes? [yes/NO] "
read l ; [ "$l" != "yes" -a "$l" != "YES" ] && exit 2

# Check output file
if [ -e "$GENERATED_COMMANDS_FILE" ]; then
    echo "Commands to run will be written into '$GENERATED_COMMANDS_FILE'."
    echo "Do you want to override it? [yes/NO]"
    read l ; [ "$l" != "yes" -a "$l" != "YES" ] && exit 3
    rm $GENERATED_COMMANDS_FILE || exit 4
fi

# Add properties
echo "$OARPROPERTY_CMD -a cpu" > $GENERATED_COMMANDS_FILE
echo "$OARPROPERTY_CMD -a core" >> $GENERATED_COMMANDS_FILE
[ -n "$WITH_HT" -a -n "$CREATE_THREAD_PROPERTY" ] && echo "$OARPROPERTY_CMD -a thread" >> $GENERATED_COMMANDS_FILE
echo "$OARPROPERTY_CMD -c -a host" >> $GENERATED_COMMANDS_FILE
echo "$OARPROPERTY_CMD -c -a cpuset" >> $GENERATED_COMMANDS_FILE
echo "$OARPROPERTY_CMD -a mem" >> $GENERATED_COMMANDS_FILE

# Init CPU id
CPU=$($OARNODESETTING_CMD --last-property-value cpu 2> /dev/null)
[ -n "$CPU" ] && ((CPU++)) || CPU=0 
# Init CORE id
CORE=$($OARNODESETTING_CMD --last-property-value core 2> /dev/null)
[ -n "$CORE" ] && ((CORE++)) || CORE=0
# Init THREAD id
if [ -n "$WITH_HT" -a -n "$CREATE_THREAD_PROPERTY" ]; then
    THREAD=$($OARNODESETTING_CMD --last-property-value thread 2> /dev/null)
fi
[ -n "$THREAD" ] && ((THREAD++)) || THREAD=0

echo
for host in $(< $1); do
    echo -n "Checking $host ... "

    # Store stderr and redirect it to /dev/null
    exec 3>&2
    exec 2> /dev/null

    mem=$($OPENSSH_CMD $host cat /proc/meminfo | grep ^MemTotal | awk '{print $2}')

    # Get CORE info in the form:
    # cpu:core
    # cpu:core
    # cpu:core
    # ...
    # of if WITH_TH:
    # cpu:core:thread
    # cpu:core:thread
    # cpu:core:thread
    # ...
    if [ -z "$WITH_HT" ]; then
        if [ "`$OPENSSH_CMD $host cat /proc/cpuinfo|grep 'physical id'`" != "" ]; then
            DATA=$($OPENSSH_CMD $host cat /proc/cpuinfo| awk '{if ($1 == "processor") core=$3; if ($1 == "physical" && $2 =="id") print $4 ":" core}'|sort -n)
        else
            DATA=$($OPENSSH_CMD $host cat /proc/cpuinfo| awk '{if ($1 == "processor") print $3 ":" $3}'|sort -n)
        fi
    else
        DATA=$($OPENSSH_CMD $host cat /proc/cpuinfo| awk '{if ($1 == "processor") core=$3; if ($1 == "physical" && $2 =="id") socket=$4; if ($1 == "core" && $2 =="id") print socket ":" $4 ":" core}'|sort -n)
    fi

    # Restore stderr
    exec 2>&3

    if [ -z "$mem" ] || [ -z "$DATA" ]; then
        echo "*** KO ***"
        rm $GENERATED_COMMANDS_FILE
        echo 
        echo "ERROR with node $host. Please check the ssh connection with oar user"
        echo "($OPENSSH_CMD $host)..."
        exit 5
    fi

    ((mem = $mem / 1024 / 1024 + 1))
    if [ -z "$WITH_HT" ]; then
        prev_cpu="-1"
        for i in $DATA; do
            cpu=$(echo $i|awk -F: '{print $1}')
            core=$(echo $i|awk -F: '{print $2}')
            thread=$(echo $i|awk -F: '{print $3}')
            [ "$cpu" != "$prev_cpu" ] && ((CPU++))
            ((CORE++))
            prev_cpu=$cpu
            echo "$OARNODESETTING_CMD -a -h $host -p host=$host -p cpu=$CPU -p core=$CORE -p cpuset=$core -p mem=$mem" >> $GENERATED_COMMANDS_FILE
        done
    else
        prev_cpu="-1"
        prev_core="-1"
        tmp_cpuset=""
        cpu=""
        core=""
        for i in $DATA; do
            cpu=$(echo $i|awk -F: '{print $1}')
            core=$(echo $i|awk -F: '{print $2}')
            thread=$(echo $i|awk -F: '{print $3}')
            if [ -z "$CREATE_THREAD_PROPERTY" ]; then
                if [ "$core" != "$prev_core" ]; then
                    if [ -n "$tmp_cpuset" ]; then
                        echo "$OARNODESETTING_CMD -a -h $host -p host=$host -p cpu=$CPU -p core=$CORE -p cpuset=$tmp_cpuset -p mem=$mem" >> $GENERATED_COMMANDS_FILE
                    fi
                    tmp_cpuset="$thread"
                else
                    tmp_cpuset="$tmp_cpuset+$thread"
                fi
            fi
            [ "$cpu" != "$prev_cpu" ] && ((CPU++))
            [ "$core" != "$prev_core" ] && ((CORE++))
            ((THREAD++))
            if [ -n "$CREATE_THREAD_PROPERTY" ]; then
                echo "$OARNODESETTING_CMD -a -h $host -p host=$host -p cpu=$CPU -p core=$CORE -p thread=$THREAD -p cpuset=$thread -p mem=$mem" >> $GENERATED_COMMANDS_FILE
            fi
            prev_cpu=$cpu
            prev_core=$core
        done
        if [ -z "$CREATE_THREAD_PROPERTY" ]; then
            echo "$OARNODESETTING_CMD -a -h $host -p host=$host -p cpu=$CPU -p core=$CORE -p cpuset=$tmp_cpuset -p mem=$mem" >> $GENERATED_COMMANDS_FILE
        fi
    fi
    echo "OK"
done



echo
echo "If the content of '$GENERATED_COMMANDS_FILE' is OK for you then you just need to execute:"
echo "  source $GENERATED_COMMANDS_FILE"

