#!/bin/bash
INSTALLER_HELPER=/usr/local/qcloud/nv//nv_driver_install_helper.sh
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
INSTALL_LOG=/tmp/nv_driver_install.log
INSTALL_CFG=/tmp/nv_driver_install.ini
PROFILE=/etc/profile

DRIVER_INSTALL_FINISH_LOG="install nvidia driver finished"
CUDA_INSTALL_FINISH_LOG="install nvidia cuda finished"
CUDNN_INSTALL_FINISH_LOG="install nvidia cudnn finished"

install_prompt(){
    echo ''
    echo 'Hello, This script will download and install the GPU driver, CUDA, CUDNN library automatically, you can not suspend or stop it until completed.'
    echo '1. The whole process will take about 15 to 25 minutes. During this time, please do not operate the GPU or install any GPU related software.'
    echo ''
    echo ''
}

install_over(){
    echo ''
    echo 'All install OK! Enjoy it!'
    echo ''
}

sudo sed -i '/progress_bar/d' ${PROFILE} >/dev/null 2>/dev/null

if [ ! -f ${INSTALL_CFG} ]; then
    exit 0
fi

driver_version=`awk -F '=' '/driver_version/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`
cuda_version=`awk -F '=' '/cuda_version/{print $2}' $INSTALL_CFG | awk '{gsub(/^\s+|\s+$/, "");print}'`
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}'`

if [ -z "$driver_version" ]; then
    exit 0
fi

p_num=`ps -ef | grep ${INSTALLER_HELPER} | grep -v grep | wc -l`
if [ $p_num -eq 0 ]; then
    exit 0
fi

install_prompt

is_ubuntu=0
if [ `grep -i "Ubuntu" /etc/os-release >/dev/null 2>/dev/null; echo $?` -eq 0 ]; then
    is_ubuntu=1
fi 

check_download() {
    filename=$1
    arry=("\\" "|" "/" "-")
    persent=0
    i=0

    totalsize=`curl -sI "http://mirrors.tencentyun.com/install/GPU/$filename"  | grep -i Content-Length | awk '{print int($2)}'`

    echo 

    sleep 0.5
    while true
    do
        p_num=`ps -ef | grep wget |grep ${filename} |grep -v grep | wc -l`
        if [ $p_num -eq 0 ]; then
            break;
        fi

        str=""
        for ii in `seq $persent`
        do
            str+="#"
        done

        let index=i%4
        let i=i+1

        printf "\033[1A"
        printf "[%-100s] %d%% %c     \n" "$str" "$persent" "${arry[$index]}"
        
        cursize=`ls -l "/tmp/$filename" | awk '{print $5}'`
        persent=`expr $cursize \* 100 / $totalsize `

        sleep 0.5
    done

    str=$(printf "%-100s" "#")
    str2=$(echo "$str" | sed 's/ /#/g') 
    printf "\033[1A"   
    printf "[%-100s] %d%% %c     \n\n" "$str2" "100" "/"
    printf ""
}

check_install() {
    arry=("\\" "|" "/" "-")
    persent=0
    i=0
    t=1

    if [ -n "`echo $1 | grep '^cuda'`" ]; then
        t=3
    fi

    echo 

    sleep 0.5
    while true
    do
        num=`ps -ef | grep /tmp/$1 | grep -v wget | grep -v grep | wc -l`
        if [ $num -eq 0 ] && [ ! -f "/tmp/$1" ]; then
            break;
        fi

        str=""
        for ii in `seq $persent`
        do
            str+="#"
        done

        let index=i%4
        let i=i+1

        let persent=persent+1
        if [ $persent -ge 99 ];then
            persent=99
        fi
        printf "\033[1A"
        printf "[%-100s] %d%% %c     \n" "$str" "$persent" "${arry[$index]}"

        sleep $t
    done

    str=$(printf "%-100s" "#")
    str2=$(echo "$str" | sed 's/ /#/g')
    printf "\033[1A"
    printf "[%-100s] %d%% %c     \n\n" "$str2" "100" "/"
    printf ""
}


check_install_result() {
    if [ ! -f "$INSTALL_LOG" ];then
        echo "No installation log file "$INSTALL_LOG", please check the basic plugin log"
        exit -1
    fi

    if [ "$1" = "driver" ]; then
        install_finish_log=${DRIVER_INSTALL_FINISH_LOG}
    elif [ "$1" = "cuda" ]; then
        install_finish_log=${CUDA_INSTALL_FINISH_LOG}
    elif [ "$1" = "cudnn" ]; then
        install_finish_log=${CUDNN_INSTALL_FINISH_LOG}
    fi
    
    found=`cat $INSTALL_LOG | grep "${install_finish_log}"`
    if [ -z "$found" ] ;then
        echo "There maybe some ERROR during installation $1, please check the install log $INSTALL_LOG."
        return
    fi
}


install_process() {

    while true
    do
        p_num=`ps -ef | grep ${INSTALLER_HELPER} | grep -v grep | wc -l`
        if [ $p_num -eq 0 ]; then
            break
        else
            p_num=`ps -ef | grep ${driver_version} | grep wget | grep -v grep | wc -l`
            if [ $p_num -gt 0 ]; then
                echo "Downloading driver: ${driver_version}, it will take about 2 minutes. the remaining installation needs about 15-20 minutes."
                check_download ${driver_version}
            fi

            p_num=`ps -ef | grep "/tmp/${driver_version}" | grep -v wget | grep -v grep | wc -l`
            if [ $p_num -gt 0 ]; then
                echo "Installing driver: ${driver_version}, it will take about 3 minutes. the remaining installation needs about 14-19 minutes."
                check_install ${driver_version}
                check_install_result "driver"
            fi

            if [ -n "${cuda_version}" ]; then
                p_num=`ps -ef | grep ${cuda_version} | grep wget | grep -v grep | wc -l`
                if [ $p_num -gt 0 ]; then
                    echo "Downloading cuda: ${cuda_version}, it will take about 3 minutes. the remaining installation needs about 13-18 minutes."
                    check_download ${cuda_version}
                fi

                p_num=`ps -ef | grep "/tmp/${cuda_version}" | grep -v wget | grep -v grep | wc -l`
                if [ $p_num -gt 0 ]; then
                    echo "Installing cuda: ${cuda_version}, it will take about 5 minutes. the remaining installation needs about 12-16 minutes."
                    check_install ${cuda_version}
                    check_install_result "cuda"
                fi
            fi

            if [ -n "${cudnn_version1}" ]; then
                p_num=`ps -ef | grep ${cudnn_version1} | grep wget | grep -v grep | wc -l`
                if [ $p_num -gt 0 ]; then
                    echo "Downloading cudnn: ${cudnn_version1}, it will take about 2 minutes. the remaining installation needs about 4-7 minutes."
                    check_download ${cudnn_version1}
                fi

                p_num=`ps -ef | grep "/tmp/${cudnn_version1}" | grep -v wget | grep -v grep | wc -l`
                if [ $p_num -gt 0 ]; then
                    echo "Installing cudnn: ${cudnn_version1}, it will take about 2 minutes. the remaining installation needs about 2-3 minutes."
                    check_install ${cudnn_version1}
                    if [ "$is_ubuntu"x == "0"x ]; then
                        check_install_result "cudnn"
                    fi
                fi
            fi

            if [ "$is_ubuntu"x == "1"x ] && [ -n "${cudnn_version2}" ] && [ -n "${cudnn_version3}" ]; then
                p_num=`ps -ef | grep ${cudnn_version2} | grep wget | grep -v grep | wc -l`
                if [ $p_num -gt 0 ]; then
                    echo "Downloading cudnn: ${cudnn_version2}, it will take about 1 minutes. the remaining installation needs about 3-6 minutes."
                    check_download ${cudnn_version2}
                fi

                p_num=`ps -ef | grep ${cudnn_version3} | grep wget | grep -v grep | wc -l`
                if [ $p_num -gt 0 ]; then
                    echo "Downloading cudnn: ${cudnn_version3}, it will take about 1 minutes. the remaining installation needs about 2-5 minutes."
                    check_download ${cudnn_version3}
                fi

                p_num_v2=`ps -ef | grep "/tmp/${cudnn_version2}" | grep -v wget | grep -v grep | wc -l`
                p_num_v3=`ps -ef | grep "/tmp/${cudnn_version3}" | grep -v wget | grep -v grep | wc -l`
                if [ $p_num_v2 -gt 0 ] || [ $p_num_v3 -gt 0 ]; then
                    echo "Installing cudnn: ${cudnn_version2}, it will take about 1 minutes. the remaining installation needs about 1-2 minutes."
                    check_install ${cudnn_version2}

                    echo "Installing cudnn: ${cudnn_version3}, it will take about 1 minutes. the remaining installation needs about 1-2 minutes."
                    check_install ${cudnn_version3}
                    check_install_result "cudnn"
                fi
            fi
        fi

    done

}

installed_info() {
    found=`cat $INSTALL_LOG | grep "$DRIVER_INSTALL_FINISH_LOG"`
    if [ -n "$found" ] ;then
        echo "Driver install finished: ${driver_version}, the remaining installation needs about 14-19 minutes."
        echo ""
    fi

    found=`cat $INSTALL_LOG | grep "$CUDA_INSTALL_FINISH_LOG"`
    if [ -n "$found" ] ;then
        echo "Cuda install finished: ${cuda_version}, the remaining installation needs about 12-16 minutes."
        echo ""
    fi
}

installed_info

install_process

install_over