#!/bin/bash

DRIVER_INSTALLER=/usr/local/qcloud/nv/nv_driver_install.sh
CUDA_INSTALLER=/usr/local/qcloud/nv/nv_cuda_install.sh
CUDNN_INSTALLER=/usr/local/qcloud/nv/nv_cudnn_install.sh
FABRICMGR_INSTALLER=/usr/local/qcloud/nv/nv_fabricmgr_install.sh
FABRICMGR_INSTALLER_FILENAME=nv_fabricmgr_install.sh
TACO_SETUP=/usr/local/qcloud/nv/taco_setup.sh
PROGRESS_BAR=/usr/local/qcloud/nv/progress_bar.sh
PROFILE=/etc/profile

INSTALL_CFG=/tmp/nv_driver_install.ini
GPUINFO_CFG=/tmp/gpu_info.ini
GPUINFO_FILENAME=gpu_info.ini

wait_time=0

while [ ! -f ${INSTALL_CFG} ]; do
    if [ $wait_time -gt 300 ]; then
        echo "Can not find " ${INSTALL_CFG}
        echo "We will do not install nvidia drivers"
        echo "### end ###"
        exit
    fi
    sleep 1
    wait_time=$(($wait_time+1))
done

cat ${INSTALL_CFG} | xargs echo

wget_file() {
    wait_info_time=0
    get_gpu_info=0
    while [ $wait_info_time -lt 10 ]; do
        wget "http://mirrors.tencentyun.com/install/GPU/$1" -O $2
        if [ $? == 0 ]; then
            get_gpu_info=1
            break;
        fi
        wait_info_time=$(($wait_info_time+1))
        sleep 1
    done
    if [ x$get_gpu_info != x1 ]; then
        echo "download file $1 failed, please check the input,url and network"
        return 1
    fi
    return 0
}

get_gpuinfo() {
    # get driver name from driver version
    gpu_version=$1
    cuda_version=$2
    cudnn_version1=$3
    if [ ! -z "$4" ]; then
        cudnn_version2=$4
        cudnn_version3=$5
    fi

    array_tmp=`eval cat $GPUINFO_CFG`
    IFS=";"
    array=($array_tmp)

    for(( i=0; i<${#array[@]}; i++))
    do
        tmp=`eval echo ${array[i]}`
        tmp_driver=`eval echo $tmp | awk -F ',' '{print $1}'`
        tmp_driver_version=`eval echo $tmp_driver | awk -F ':' '{print $2}'`
        driver_version=`eval echo $tmp_driver_version`
        tmp_name=`eval echo $tmp | awk -F ',' '{print $2}'`
        tmp_driver_name=`eval echo $tmp_name | awk -F ':' '{print $2}'`
        driver_name=`eval echo $tmp_driver_name`
        if [ ! -z "$cuda_version" ]; then
            tmp_cuda=`eval echo $tmp | awk -F ',' '{print $3}'`
            tmp_cuda_version=`eval echo $tmp_cuda | awk -F ':' '{print $2}'`
            tmp_cuda_version=`eval echo $tmp_cuda_version`

            tmp_name=`eval echo $tmp | awk -F ',' '{print $4}'`
            tmp_cuda_name=`eval echo $tmp_name | awk -F ':' '{print $2}'`
            cuda_name=`eval echo $tmp_cuda_name`
        elif [ $gpu_version == $driver_version ]; then
            echo "install driver only, goto install $driver_name"
            rm -f ${INSTALL_CFG}
            rm -f ${GPUINFO_CFG}
            $DRIVER_INSTALLER $driver_name
            exit 0
        fi

        if [ ! -z "$cudnn_version1" ]; then
            tmp_cudnn=`eval echo $tmp | awk -F ',' '{print $5}'`
            tmp_cudnn_version=`eval echo $tmp_cudnn | awk -F ':' '{print $2}'`
            tmp_cudnn_version=`eval echo $tmp_cudnn_version`

            #ubuntu has 3 cudnn names
            lsb_release -a 2>&1 | grep "Ubuntu" &> /dev/null
            if [ $? == 0 ]; then
                tmp_cudnn_name1=`eval echo $tmp | awk -F ',' '{print $6}'`
                tmp_cudnn_name2=`eval echo $tmp | awk -F ',' '{print $7}'`
                tmp_cudnn_name3=`eval echo $tmp | awk -F ',' '{print $8}'`

                cudnn_name1=`eval echo $tmp_cudnn_name1 | awk -F ':' '{print $2}'`
                cudnn_name2=`eval echo $tmp_cudnn_name2 | awk -F ':' '{print $2}'`
                cudnn_name3=`eval echo $tmp_cudnn_name3 | awk -F ':' '{print $2}'`
                cudnn_name1=`eval echo $cudnn_name1`
                cudnn_name2=`eval echo $cudnn_name2`
                cudnn_name3=`eval echo $cudnn_name3`
                if [ -z "$cudnn_name2" ] || [ -z "$cudnn_name3" ]; then
                    continue
                fi
                if [ $gpu_version == $driver_version ] && [ $cuda_version == $tmp_cuda_version ] && [ $cudnn_version1 == $tmp_cudnn_version ]; then
                    echo "install driver, cuda and cudnn, goto install $driver_name $cuda_name $cudnn_name1 $cudnn_name2 $cudnn_name3"
                    rm -f ${INSTALL_CFG}
                    rm -f ${GPUINFO_CFG}
                    $DRIVER_INSTALLER $driver_name
                    $CUDA_INSTALLER $cuda_name
                    $CUDNN_INSTALLER $cudnn_name
                    $CUDNN_INSTALLER $cudnn_name1 $cudnn_name2 $cudnn_name3
                    exit 0
                fi
            else
                tmp_cudnn_name=`eval echo $tmp | awk -F ',' '{print $6}'`
                cudnn_name=`eval echo $tmp_cudnn_name | awk -F ':' '{print $2}'`
                cudnn_name1=`eval echo $cudnn_name`

                if [ $gpu_version == $driver_version ] && [ $cuda_version == $tmp_cuda_version ] && [ $cudnn_version1 == $tmp_cudnn_version ]; then
                    echo "install driver, cuda and cudnn, goto install $driver_name $cuda_name $cudnn_name"
                    rm -f ${INSTALL_CFG}
                    rm -f ${GPUINFO_CFG}
                    $DRIVER_INSTALLER $driver_name
                    $CUDA_INSTALLER $cuda_name
                    $CUDNN_INSTALLER $cudnn_name
                    exit 0
                fi
            fi
        elif [ $gpu_version == $driver_version ] && [ $cuda_version == $tmp_cuda_version ]; then
            echo "install driver and cuda, goto install $driver_name $cuda_name"
            rm -f ${INSTALL_CFG}
            rm -f ${GPUINFO_CFG}
            $DRIVER_INSTALLER $driver_name
            $CUDA_INSTALLER $cuda_name
            exit 0
        fi
    done
}

taco_setup=`awk -F '=' '/taco_setup/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`
if [ "$taco_setup"x = "1"x ]; then
    echo "setup taco"
    nohup $TACO_SETUP > /dev/null 2>&1 &
fi

install_driver=`awk -F '=' '/install_driver/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`

if [ "$install_driver"x = "1"x ]; then
    driver_version=''
    cuda_version=''
    cudnn_version1=''
    cudnn_version2=''
    cudnn_version3=''

    driver_version=`awk -F '=' '/driver_version/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`
    echo "install nvidia driver $driver_version"

    install_fabric_mgr=`awk -F '=' '/install_fabric_manager/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`
    echo "install fabric mgr $install_fabric_mgr"

    install_cuda=`awk -F '=' '/install_cuda/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`
    if [ "$install_cuda"x = "1"x ]; then
        cuda_version=`awk -F '=' '/cuda_version/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`
        echo "install nvidia cuda $cuda_version"

        install_cudnn=`awk -F '=' '/install_cudnn/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`
        if [ "$install_cudnn"x = "1"x ]; then
            cudnn_version1=`awk -F '=' '/cudnn_version1/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`
            cudnn_version2=`awk -F '=' '/cudnn_version2/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`
            cudnn_version3=`awk -F '=' '/cudnn_version3/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`
            echo "install nvidia cudnn $cudnn_version1 $cudnn_version2 $cudnn_version3"
        else
            echo "### not install cudnn ###"
        fi
    else
        echo "### not install cuda ###"
    fi

    #check input args and install driver, cuda, cudnn
    echo "start here... $driver_version"
    if [[ $driver_version = *NVIDIA* ]]
    then
        echo "$PROGRESS_BAR" >> $PROFILE
        
        #install driver, cuda, and cudnn with their names
        $DRIVER_INSTALLER $driver_version
        if [ "$install_cuda"x = "1"x ]; then
            $CUDA_INSTALLER $cuda_version
        fi
        if [ "$install_cudnn"x = "1"x ]; then
            $CUDNN_INSTALLER $cudnn_version1 $cudnn_version2 $cudnn_version3
        fi
        if [ "$install_fabric_mgr"x = "1"x ]; then
            wget_file $FABRICMGR_INSTALLER_FILENAME $FABRICMGR_INSTALLER
            [ $? -eq 0 ] && chmod u+x $FABRICMGR_INSTALLER && $FABRICMGR_INSTALLER $driver_version
        fi
        sed -i '/progress_bar/d' ${PROFILE} >/dev/null 2>/dev/null
        rm -f ${INSTALL_CFG}
        rm -f ${GPUINFO_CFG}
    else
        echo "get driver name from gpu_info config file"
        wget_file $GPUINFO_FILENAME $GPUINFO_CFG
        [ $? -eq 0 ] && get_gpuinfo $driver_version $cuda_version $cudnn_version1 $cudnn_version2 $cudnn_version3
    fi

else
    rm -f ${INSTALL_CFG}
    rm -f ${GPUINFO_CFG}
    echo "### not install driver ###"
fi