##############################################################################
# Copyright (c) Members of the EGEE Collaboration. 2007.
# See http://www.eu-egee.org/partners/ for details on the copyright
# holders.
#
# 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.
##############################################################################
#
# NAME :        config_mpi_wn
#
# DESCRIPTION : This function configures the WN for MPI.
#               
#
# AUTHORS :     Stephen.Childs@cs.tcd.ie, John.Walsh@cs.tcd.ie,
#               enolfc@ifca.unican.es
#
# NOTES :       
#
# YAIM MODULE:  glite-yaim-mpi
#
##############################################################################

mpi_find_file_dir_rpm() {
    pkgname=$1
    archs=`rpm -q $1 | wc -l`
    if [ $archs -gt 1 ] 
        then
            pkgname=$1.`uname -i`
    fi
    dir=`rpm -ql $pkgname | awk -F/bin/$2 "/${2}\$/ { print \\\$1}" | tail -1`
    if [ "x$dir" != "x" ]
        then
        if [ -d "$dir" -a -f "$dir/bin/$2" ]
            then
            echo $dir
            return 0
        fi
    fi
    return 1
}

mpi_find_dir_rpm() {
    # mpicc may not be available in non-devel packages, look for mpirun or mpiexec 
    mpidir=`mpi_find_file_dir_rpm $1 mpiexec`
    if [ $? -ne 0 ]
        then
        mpidir=`mpi_find_file_dir_rpm $1 mpirun`
        if [ $? -eq 0 ]
            then
            echo $mpidir
            return 0
        fi
    else
        echo $mpidir
        return 0
    fi
    return 1
}

mpi_find_file_dir_deb() {
    # tries to find the directory where the binaries are
    pkg=$1
    # two lines instead of one to have it working
    dir=`dpkg -L $pkg | egrep "(${3}|${3}.${2})\$" | head -1`
    dir=`echo $dir | awk -F/bin/${3} '{print $1}'`
    if [ "x$dir" = "x" ]
        then
        return 1
    fi
    echo $dir
    return 0
}

mpi_find_dir_deb() {
    FLAVOUR_UPPER=`echo ${1} | tr '[:lower:]' '[:upper:]' | tr '-' '_'`
    pkg=`eval echo \\$MPI_${FLAVOUR_UPPER}_DEBIAN_PKG`
    if [ "x$pkg" = "x" ]
        then
        return 1
    fi
    mpidir=`mpi_find_file_dir_deb $pkg $1 mpiexec`
    if [ $? -ne 0 ]
        then
        mpidir=`mpi_find_file_dir_deb $pkg $1 mpirun`
        if [ $? -eq 0 ]
            then
            echo $mpidir
            return 0
        fi
    else
        echo $mpidir
        return 0
    fi
    return 1
}

mpi_find_version_deb() {
    FLAVOUR_UPPER=`echo ${1} | tr '[:lower:]' '[:upper:]' | tr '-' '_'`
    pkg=`eval echo \\$MPI_${FLAVOUR_UPPER}_DEBIAN_PKG`
    if [ "x$pkg" = "x" ] 
        then
        return
    fi
    v=`dpkg-query -W $pkg`
    if [ $? -ne 0 ] 
        then
        return
    fi
    ver=`echo $v | awk '{print $2}' | awk -F- '{print $1}'`
    if [ "x$ver" != "x" ]
        then
        echo $ver
    fi
}

mpi_find_version_rpm() {
    version=`rpm -q $1 | grep -v "is not installed" | awk -F- '{print $2}' | tail -1`
    if [ "x$version" != "x" ]
        then
        echo $version
    fi
}


mpi_find_dir() {
    if [ "x$MPI_DEBIAN_SUPPORT" = "x1" ] 
        then
        mpi_find_dir_deb $*
    else
        mpi_find_dir_rpm $*
    fi
    return $?
}

mpi_find_version() {
    if [ "x$MPI_DEBIAN_SUPPORT" = "x1" ] 
        then
        mpi_find_version_deb $*
    else
        mpi_find_version_rpm $*
    fi
    return $?
}

function config_mpi_wn_check() {
    yaimlog DEBUG "This function doesn't currently require any variables."
}

function config_mpi_wn_setenv() {
# Figure out if we are in deb or rpm based distro
    MPI_DEBIAN_SUPPORT=0
    if [ "x$OS_TYPE" == "xdebian" ]; then
        MPI_DEBIAN_SUPPORT=1
    fi
    export MPI_DEBIAN_SUPPORT
# These are the flavours of MPI currently supported. If you want to add
# a new one just append it to this variable, default flavour if not
# defined by admin is the first one found.
    MPI_FLAVOURS="OPENMPI MPICH2 MPICH LAM"
    if [ "x$MPI_DEFAULT_FLAVOUR" != "x" ]
        then
        MPI_DEFAULT_FLAVOUR=`echo ${MPI_DEFAULT_FLAVOUR} | tr '[:upper:]' '[:lower:]'`
    fi
    SOME_MPI_IS_SELECTED="false"
####@ Set environment variables describing the installed MPI implementations
    for FLAVOUR in ${MPI_FLAVOURS}; do
    #----
        FLAVOUR_LOWER=`echo ${FLAVOUR} | tr '[:upper:]' '[:lower:]'`

        eval FLAVOUR_ENABLED="\${MPI_${FLAVOUR}_ENABLE}"
        eval FLAVOUR_VERSION="\${MPI_${FLAVOUR}_VERSION}"
        eval FLAVOUR_PATH="\${MPI_${FLAVOUR}_PATH}"
        eval FLAVOUR_MPIEXEC="\${MPI_${FLAVOUR}_MPIEXEC}"

        if [ "x${FLAVOUR_ENABLED}" == "xyes" ]
            then
            yaimlog DEBUG "Configuring ${FLAVOUR} (Path \"${FLAVOUR_PATH}\" Version \"${FLAVOUR_VERSION}\")"
            SOME_MPI_IS_SELECTED="true"
            if [ "x$FLAVOUR_PATH" == "x" ]
                then
                yaimlog DEBUG "Searching for MPI flavour ${FLAVOUR_LOWER}"
                FLAVOUR_PATH=`mpi_find_dir ${FLAVOUR_LOWER}`
                yaimlog INFO "Found ${FLAVOUR} at ${FLAVOUR_PATH}"
                if [ "x$FLAVOUR_PATH" == "x" ]
                    then
                    yaimlog ERROR "${FLAVOUR}_PATH could not be determined. Check definition in site-info.def"
                    yaimlog ERROR "MPI configuration is not complete. Breaking..."
                    exit 1
                fi
            fi

            if [ -e "${FLAVOUR_PATH}" ]
                then
                yaimgridenv_set MPI_${FLAVOUR}_PATH $FLAVOUR_PATH
            else
                yaimlog ERROR "${FLAVOUR}_PATH not found. Check definition in site-info.def"
                yaimlog ERROR "MPI configuration is not complete. Breaking..."
                exit 1
            fi

            if [ "x${FLAVOUR_VERSION}" == "x" ]
                then
                yaimlog DEBUG "Trying to detect version of ${FLAVOUR} installed."
                FLAVOUR_VERSION=`mpi_find_version ${FLAVOUR_LOWER}`
                if [ "x$FLAVOUR_VERSION" == "x" ]
                    then
                    yaimlog ERROR "${FLAVOUR}_VERSION could not be determined. Check definition in site-info.def"
                    yaimlog ERROR "MPI configuration is not complete. Breaking..."
                    exit 1
                else
                    yaimlog INFO "Installed version of ${FLAVOUR} is ${FLAVOUR_VERSION}"  
                fi
            fi
            yaimgridenv_set MPI_${FLAVOUR}_VERSION $FLAVOUR_VERSION

            if [ "x$FLAVOUR_MPIEXEC" != "x" ]
                then
                yaimgridenv_set MPI_${FLAVOUR}_MPIEXEC $FLAVOUR_MPIEXEC
            fi

            if [ "x$MPI_DEFAULT_FLAVOUR" = "x" ]
                then
                MPI_DEFAULT_FLAVOUR=${FLAVOUR_LOWER}
            fi
        fi
        
    done

####@ Set environment variables describing shared area.
#### if MPI_SHARED_HOME is defined as a PATH then set this in environment
    if [ "x${MPI_SHARED_HOME}" != "x" ] ; then
        if [ "$MPI_SHARED_HOME" = "no" ] ; then
            yaimgridenv_set MPI_SHARED_HOME $MPI_SHARED_HOME
        elif [ "$MPI_SHARED_HOME" = "yes" ] ; then
            # is a location provided? if not then, assume that the admin wanted to
            # just let us know that home is shared
            if [ "x$MPI_SHARED_HOME_PATH" = "x" ] ; then
                yaimgridenv_set MPI_START_SHARED_FS 1
            else
                yaimgridenv_set MPI_SHARED_HOME $MPI_SHARED_HOME
                yaimgridenv_set MPI_SHARED_HOME_PATH $MPI_SHARED_HOME_PATH
            fi
        else
            yaimlog INFO "MPI_SHARED_HOME value should be yes/no instead of $MPI_SHARED_HOME"
        fi
    fi

####@ Set environment variables common to all MPIs
    if [ "x$SOME_MPI_IS_SELECTED" = "xtrue" ]; then
        if [ "x$MPI_SSH_HOST_BASED_AUTH" != "x" ] ; then
            yaimgridenv_set MPI_SSH_HOST_BASED_AUTH $MPI_SSH_HOST_BASED_AUTH
        fi
        
        if [ "x$MPI_MPI_START" != "x" ]; then
            yaimlog INFO "Setting I2G_MPI_START env variable to $MPI_MPI_START"
            yaimgridenv_set I2G_MPI_START $MPI_MPI_START
        else
            # old mpi-start path, keep it just in case
            if [ -f /opt/i2g/bin/mpi-start ] 
                then
                yaimlog INFO "Setting I2G_MPI_START env variable to /opt/i2g/bin/mpi-start"
                yaimgridenv_set I2G_MPI_START /opt/i2g/bin/mpi-start
            else
                # new path
                yaimlog INFO "Setting I2G_MPI_START env variable to /usr/bin/mpi-start"
                yaimgridenv_set I2G_MPI_START /usr/bin/mpi-start
            fi
        fi

        yaimgridenv_set MPI_DEFAULT_FLAVOUR $MPI_DEFAULT_FLAVOUR
    else
        yaimlog INFO "No MPI flavours enabled."
    fi
    return $YEX_OK
}

config_mpi_wn() {
    #-------------------------------------------------------------------------
    # WN-related configuration for MPI
    #-------------------------------------------------------------------------
####@ If the EDG_WL_SCRATCH directory is used, then create
####@ .mpi directory world-writeable with sticky bit so
####@  that all users can write to it.
    if [ "x${EDG_WL_SCRATCH}" != "x" ]; then
        if [ ! -d ${EDG_WL_SCRATCH}/.mpi ]; then
            mkdir ${EDG_WL_SCRATCH}/.mpi
        fi
        chmod 1777 ${EDG_WL_SCRATCH}/.mpi
    fi
    return $YEX_OK
}
