#! /bin/bash
ASCEND_SECLOG="/var/log/ascend_seclog"
logFile="${ASCEND_SECLOG}/ascend_install.log"
runLogFile="${ASCEND_SECLOG}/ascend_run_servers.log"
installInfo="/etc/ascend_install.info"
ASCEND_DRIVER_CFG_FILE="/etc/ascend_driver.conf"
KMSAGENT_CONFIG_ITEM="kmsagent"
KMSAGENT_WORK_DIR="/var/kmsagentd"
KMSAGENT_CONF_MAIN_KSF="/var/kmsagentd/kmsconf.ksf"
KMSAGENT_SERVICE_DIR="/lib/systemd/system"
KMSAGENT_USER_SERVICE_DIR="/usr/lib/systemd/system"
KMSAGENT_SERVICE_NAME="ascend-kmsagent.service"
ASCEND_DRIVER_SETUP_SCRIPT="ascend_driver_config.sh"
username=HwHiAiUser
usergroup=HwHiAiUser
hotreset_status_file="/var/log/hotreset_status.log"
preinst_flag_first_file="/var/log/ascend_preinst_first_flag"
davinci_num_file="/var/log/davinci_num.log"
export feature_hot_reset=n
export feature_crl_check=n
export feature_nvcnt_compile=n
export feature_dkms_compile=n
export input_install_for_all=n
export feature_virt_scene=n
export feature_no_device_kernel=n
export feature_ver_compatible_check=n

OPERATION_LOGDIR="${ASCEND_SECLOG}"
OPERATION_LOGPATH="${OPERATION_LOGDIR}/operation.log"
LOG_OPERATION_INSTALL="Install"
LOG_OPERATION_UPGRADE="Upgrade"
LOG_LEVEL_SUGGESTION="SUGGESTION"
LOG_LEVEL_MINOR="MINOR"
LOG_LEVEL_UNKNOWN="UNKNOWN"
LOG_RESULT_SUCCESS="success"
LOG_RESULT_FAILED="failed"

install_info_key_array=("Driver_Install_Type" "UserName" "UserGroup" "Driver_Install_Path_Param" "Firmware_Install_Path_Param" "Firmware_Install_Type"
                        "Driver_Install_For_All" "Driver_Install_Mode")

get_kmsagent_config() {
    if [ $kmsagent_is_enable = y ];then
        return 0
    fi

    if [ -f $ASCEND_DRIVER_CFG_FILE ];then
        local kms_config_txt=`grep -r "${KMSAGENT_CONFIG_ITEM}=" "${ASCEND_DRIVER_CFG_FILE}" | cut -d"=" -f2-`
        if [ "${kms_config_txt}X" = "enableX" ];then
            kmsagent_is_enable=y
        fi
    fi
}

log() {
    cur_date=`date +"%Y-%m-%d %H:%M:%S"`
    user_id=`id | awk '{printf $1}'`
    echo "[Driver] [$cur_date] [$user_id] "$1 >> $logFile
    return 0
}

drvEcho() {
    cur_date=`date +"%Y-%m-%d %H:%M:%S"`
    echo "[Driver] [$cur_date] $1"
}

drvColorEcho() {
    cur_date=`date +"%Y-%m-%d %H:%M:%S"`
    echo -e  "[Driver] [$cur_date] $1"

}

exitLog() {
    cur_date=`date +"%Y-%m-%d %H:%M:%S"`
    drvEcho "[INFO]End time: ${cur_date}"
    drvEcho "[INFO]End time: ${cur_date}" >> $logFile
    exit $1
}

logOperation() {
    local operation="$1"
    local start_time="$2"
    local runfilename="$3"
    local result="$4"
    local install_type="$5"
    local cmdlist="$6"
    local level

    if [ "${operation}" = "${LOG_OPERATION_INSTALL}" ]; then
        level="${LOG_LEVEL_SUGGESTION}"
    elif [ "${operation}" = "${LOG_OPERATION_UPGRADE}" ]; then
        level="${LOG_LEVEL_MINOR}"
    else
        level="${LOG_LEVEL_UNKNOWN}"
    fi

    if [ ! -d "${OPERATION_LOGDIR}" ]; then
        mkdir -p ${OPERATION_LOGDIR}
        setFileChmod -f 750 ${OPERATION_LOGDIR}
    fi

    if [ ! -f "${OPERATION_LOGPATH}" ]; then
        touch ${OPERATION_LOGPATH}
        setFileChmod -f 640 ${OPERATION_LOGPATH}
    fi

    if [ $upgrade = y ]; then
        echo "${operation} ${level} root ${start_time} 127.0.0.1 ${runfilename} ${result}"\
            "cmdlist=${cmdlist}." >> ${OPERATION_LOGPATH}
    else
        echo "${operation} ${level} root ${start_time} 127.0.0.1 ${runfilename} ${result}"\
            "install_type=${install_type}; cmdlist=${cmdlist}." >> ${OPERATION_LOGPATH}
    fi
}

getInstallParam() {
    local _key="$1"
    local _file="$2"
    local _param

    if [ ! -f "${_file}" ];then
        exit 1
    fi

    # deal with abnormal install-cfg file.
    # when install-cfg file is abnormal, command 'grep' can not be used but 'sed' does.
    local normal_args
    grep -r "${_key}" "${_file}" | grep -q '='
    if [ $? -ne 0 ]; then
        normal_args=$(sed -n '/=/p' ${_file})
        echo "${normal_args}" > ${_file}
    fi

    for key_param in "${install_info_key_array[@]}"; do
        if [ ${key_param} == ${_key} ]; then
            _param=`grep -r "${_key}=" "${_file}" | cut -d"=" -f2-`
            break;
        fi
    done
    echo "${_param}"
}

if [ -z $Driver_Install_Path_Param ];then
    Driver_Install_Path_Param=/usr/local/Ascend
fi

source ${Driver_Install_Path_Param}/driver/script/specific_func.inc


checkDriverProduct() {
    local bdfs=($(lspci -nn | grep -o '19e5:d[0-9]0[0-9]\|20c6:d[0-9]0[0-9]\|20e9:d[0-9]0[0-9]'))
    first_bdf=$(echo "$device_bdf" | cut -d'|' -f1)
    log "[INFO] Driver  expected bdfs (${first_bdf//0x/}), device actual bdfs is (${bdfs[*]})"
    if [ -n "$device_bdf" ] && [ ${#bdfs[@]} -ne 0 ] && [[ ! ${bdfs[*]} =~ ${first_bdf//0x/} ]]; then
        # does not match the device_bdf configured in the driver package.
        drvEcho "[ERROR]The current installation package and environment is not match, please check."
        exit 1
    fi
}


get_feature_config() {
    local config
    config=$(awk -F= '/FEATURE_HOT_RESET/{print $2}' "${sourcedir}/script/feature.conf")
    log "[INFO]FEATURE_HOT_RESET is : $config"
    if [ "$config"x = "y"x ];then
        feature_hot_reset=y
        log "[INFO]set feature_hot_reset=y"
    fi

    config=$(awk -F= '/FEATURE_CRL_CHECK/{print $2}' "${sourcedir}/script/feature.conf")
    log "[INFO]FEATURE_CRL_CHECK is : $config"
    if [ "$config"x = "y"x ];then
        feature_crl_check=y
        log "[INFO]set feature_crl_check=y"
    fi

    config=$(awk -F= '/FEATURE_NVCNT_CHECK/{print $2}' "${sourcedir}/script/feature.conf")
    log "[INFO]FEATURE_NVCNT_CHECK is : $config"
    if [ "$config"x = "y"x ];then
        feature_nvcnt_check=y
        log "[INFO]set feature_nvcnt_check=y"
    fi

    config=$(awk -F= '/FEATURE_DKMS_COMPILE/{print $2}' "${sourcedir}/script/feature.conf")
    log "[INFO]FEATURE_DKMS_COMPILE is : $config"
    if [ "$config"x = "y"x ];then
        feature_dkms_compile=y
        log "[INFO]set feature_dkms_compile=y"
    fi

    config=$(awk -F= '/FEATURE_VIRT_SCENE/{print $2}' "${sourcedir}/script/feature.conf")
    log "[INFO]FEATURE_VIRT_SCENE is : $config"
    if [ "$config"x = "y"x ];then
        feature_virt_scene=y
        log "[INFO]set feature_virt_scene=y"
    fi

    config=$(awk -F= '/FEATURE_NO_DEVICE_KERNEL/{print $2}' "${sourcedir}/script/feature.conf")
    log "[INFO]FEATURE_NO_DEVICE_KERNEL is : $config"
    if [ "$config"x = "y"x ];then
        feature_no_device_kernel=y
        log "[INFO]set feature_no_device_kernel=y"
    fi

    config=$(awk -F= '/FEATURE_VER_COMPATIBLE_CHECK/{print $2}' "${sourcedir}/script/feature.conf")
    log "[INFO]FEATURE_VER_COMPATIBLE_CHECK is : $config"
    if [ "$config"x = "y"x ];then
        feature_ver_compatible_check=y
        log "[INFO]set feature_ver_compatible_check=y"
    fi

    return 0
}

initKmsagentService(){
    local new_path="${Install_Path_Param%/}"
    local kms_library_path="${new_path}/driver/lib64/inner:${new_path}/driver/lib64:${new_path}/driver/lib64/driver:${new_path}/driver/lib64/common"
    local kms_restore_path="${new_path}/driver/script/kmsagent_restore_conf.sh"

    if [ $docker != y ] && [ $devel != y ]; then

        if [ ! -f "$Driver_Install_Path_Param/driver/script/kmsagent.service" ];then
            # Not need init kmsagent service.
            return 0
        fi
        if [ $1 == "false" ] && [ -f "$Driver_Install_Path_Param/driver/script/kmsagent_install.sh" ]; then
            $Driver_Install_Path_Param/driver/script/kmsagent_install.sh ${new_path}/driver/tools/kmsagent ${new_path}/driver/script/kmsagent.conf $KMSAGENT_WORK_DIR $username $usergroup $kms_library_path $kms_restore_path "false" >>$logFile 2>&1
            return $?
        fi
        # 1. Create driver config file
        if [ -d $KMSAGENT_SERVICE_DIR ]; then
            KMSAGENT_SERVICES="${KMSAGENT_SERVICE_DIR}/${KMSAGENT_SERVICE_NAME}"
            source $Driver_Install_Path_Param/driver/script/kmsagent.service
            systemctl daemon-reload 1>/dev/null 2>&1
        elif [ -d $KMSAGENT_USER_SERVICE_DIR ]; then
            KMSAGENT_SERVICES="${KMSAGENT_USER_SERVICE_DIR}/${KMSAGENT_SERVICE_NAME}"
            source $Driver_Install_Path_Param/driver/script/kmsagent.service
            systemctl daemon-reload 1>/dev/null 2>&1
        else
            log "[WARNING]The current environment does not support kmsagent services and can't use the model protection function."
            drvEcho "[WARNING]The current environment does not support kmsagent services and can't use the model protection function."
        fi

        createFile "$ASCEND_DRIVER_CFG_FILE" "root":"root" 644
        sed -i "/${KMSAGENT_CONFIG_ITEM}=/d" "${ASCEND_DRIVER_CFG_FILE}"
        if [ $kmsagent_is_enable = y ];then
            echo "${KMSAGENT_CONFIG_ITEM}=enable" >> "${ASCEND_DRIVER_CFG_FILE}"
        else
            echo "${KMSAGENT_CONFIG_ITEM}=disable" >> "${ASCEND_DRIVER_CFG_FILE}"
        fi
        setFileChmod -f 644 $ASCEND_DRIVER_CFG_FILE

        # Init kmsagent log file.
        local user_home_path=`eval echo "~${username}"`
        chown -fh $username:$usergroup ${user_home_path}/kmsagent_*.log >& /dev/null
        log "[INFO]Switch to the file owner to config the kmsagent_log in the home dir: $?"

        chown -fh $username:$usergroup ${KMSAGENT_WORK_DIR}/kmsagent_*.log >& /dev/null
        log "[INFO]Switch to the file owner to config the kmsagent_log in the work dir: $?"

        # 2. Create kmsagent work dir and restore backup data.
        if [ -f "$Driver_Install_Path_Param/driver/script/kmsagent_install.sh" ]; then
            $Driver_Install_Path_Param/driver/script/kmsagent_install.sh ${new_path}/driver/tools/kmsagent ${new_path}/driver/script/kmsagent.conf $KMSAGENT_WORK_DIR $username $usergroup $kms_library_path $kms_restore_path "true" >>$logFile 2>&1
        fi
        # 3. Auto start kms service.
        if [ $kmsagent_is_enable = y ]; then
            ${new_path}/driver/script/${ASCEND_DRIVER_SETUP_SCRIPT} $KMSAGENT_CONFIG_ITEM start >> $runLogFile 2>&1 &
        fi
    fi
}

installationCompletionMessage() {
    local run_install_type="$1"
    local new_path

    new_path="${Install_Path_Param%/}"
    if [ -e "$hotreset_status_file" ]; then
        hotreset_status=`cat "$hotreset_status_file"`
        rm -f "$hotreset_status_file"
    else
        hotreset_status="unknown"
    fi

    # Convert it to lowercase.
    run_install_type="${run_install_type,,}"
    # delete 'e' from install type.
    run_install_type="${run_install_type/e/}ed"

    local start_time=$(date +%s)
    local cnt=0
    local tool_path="$new_path/driver/tools/upgrade-tool"
    if [[ -e ${tool_path} && "${hotreset_status}" == "success" ]] || \
        [[ "$first_time_install_drv_flag" == "Y" && "$PCIE_INSMOD_STATUS_BEFORE_INSTALL" == "n" && "$force" == "n" ]]; then
        log "[INFO]Hotreset success or first to install, start to check boot status."
        hotreset_status="unknown"
        while (($(date +%s) - start_time < 60)) || (( cnt < 3 ))
        do
            (( cnt += 1 ))
            output=$(timeout 20s ${tool_path} --device_index -1 --boot_status 2>&1)
            if (( $? != 124 )); then  # Error code 124 indicates timeout
                boot_status_cnt=$(echo "${output}" | grep -c 'boot status')
                boot_status_oks=$(echo "${output}" | grep -c 'boot status:16')
                if (( boot_status_cnt > 0 && boot_status_cnt == boot_status_oks )); then
                    # Got boot status and all boot status is 16
                    "${Driver_Install_Path_Param}"/driver/script/run_driver_upgrade_check.sh -s &>/dev/null &
                    hotreset_status="success"
                    break
                fi
            fi
            sleep 10
        done
        if [ -e $davinci_num_file ]; then
            source $davinci_num_file
        fi
        if [ "$hotreset_status"x = "unknown"x ] || [ -e $davinci_num_file -a $scan_davinci_num -ne $all_davinci_num ]; then
            drvEcho "[WARNING]There are some chips not ready, please check after reboot"
            log "[WARNING]There are some chips not ready, please check after reboot"
            rm -f  $davinci_num_file
        fi
    fi

    if [ "$hotreset_status"x = "success"x ] || [ "$installType"x = "docker"x ] || [ "$installType"x = "devel"x ] ; then
        drvColorEcho  "[INFO]\033[32mDriver package ${run_install_type} successfully! The new version takes effect immediately. \033[0m"
        if [ "$installType"x = "docker"x ];then
            # lib-put-path is different from the others.
            if [ "${lib_put_path}"x = "specific"x ]; then
                echo -e "Please make sure that\n    - LD_LIBRARY_PATH includes $new_path/driver/lib64\n"\
                        "   - Please refer to the instruction manual for specific details"
            else
                echo -e "Please make sure that\n    - LD_LIBRARY_PATH includes $new_path/driver/lib64/common:$new_path/driver/lib64/driver\n"\
                        "   - Please refer to the instruction manual for specific details"
            fi
        fi
    else
        if [ "$hotreset_status"x = "ko_abort"x ]; then
            drvColorEcho "[INFO]\033[32mDriver package ${run_install_type} finished! \033[0m"
            drvColorEcho "[WARNING]Kernel modules can not be removed, reboot needed for installation/upgrade to take effect!"
        else
            log "[INFO]hotreset_status : $hotreset_status, need to reboot to take effect."
            drvColorEcho "[INFO]\033[32mDriver package ${run_install_type} successfully! \033[0m\033[31mReboot needed for installation/upgrade to take effect! \033[0m"
        fi
    fi
}
statusUpdate() {
    echo "Driver_Install_Status=complete" >> $installInfo
    log "Installation status updated successfully"
}
start_time=$(date +"%Y-%m-%d %H:%M:%S")
if [ -e "${preinst_flag_first_file}" ]; then
    source $preinst_flag_first_file
    rm -f $preinst_flag_first_file
fi
Driver_Install_Type=$(getInstallParam "Driver_Install_Type" "${installInfo}")
Driver_Install_Path_Param=$(getInstallParam "Driver_Install_Path_Param" "${installInfo}")
if [ -z $Driver_Install_Path_Param ];then
    Driver_Install_Path_Param=/usr/local/Ascend
fi
Install_Path_Param=$Driver_Install_Path_Param
Driver_Install_For_All="$(getInstallParam "Driver_Install_For_All" "${installInfo}")"
UserName=$(getInstallParam "UserName" "${installInfo}")
UserGroup=$(getInstallParam "UserGroup" "${installInfo}")
username=$UserName
usergroup=$UserGroup
force=n
input_install_for_all=Driver_Install_For_All
sourcedir=$Driver_Install_Path_Param/driver
devel=n
docker=n
upgrade=n
uninstall=n
kmsagent_is_enable=n

# load common.sh, get install.info
COMMON_SHELL="$Driver_Install_Path_Param/driver/script/common.sh"
source "${COMMON_SHELL}"

get_feature_config
get_kmsagent_config

check=n
quiet=n

if [ -f "/etc/uvp_version" ] && [ ! -h "/etc/uvp_version" ];then
    systemctl status sshd 2>&1 | grep chroot | grep -v "grep" &>/dev/null
    if [ $? -eq 0 ];then
        force=y
    fi
fi

export LD_LIBRARY_PATH="$Driver_Install_Path_Param"/driver/lib64/common:"$Driver_Install_Path_Param"/driver/lib64/driver:"${Driver_Install_Path_Param}"/driver/lib64:${LD_LIBRARY_PATH}
cp -rf $Driver_Install_Path_Param/driver/host_rpm $Driver_Install_Path_Param/driver/host
cp -rf $Driver_Install_Path_Param/driver/host_driver_binary_rpm $Driver_Install_Path_Param/driver/host_driver_binary
cp -rf $Driver_Install_Path_Param/driver/device_rpm $Driver_Install_Path_Param/driver/device
if [ $1 -eq 2 ]; then
    if [ -f $Driver_Install_Path_Param/driver/script/ver_check.sh ]; then
        $Driver_Install_Path_Param/driver/script/ver_check.sh $check $quiet || exitLog 1
    fi
fi
checkDriverProduct
bash $Driver_Install_Path_Param/driver/script/run_driver_install.sh "$Driver_Install_Path_Param" $Driver_Install_Type $force
driver_install_status=$?
operation="${LOG_OPERATION_INSTALL}"
if [ $driver_install_status -eq 0 ];then
    # Init kmsagent service.
    initKmsagentService "true"
    if [ $1 -eq 1 ]; then
        installationCompletionMessage "Install"
    else
        installationCompletionMessage "Upgrade"
    fi
else
    initKmsagentService "false"
    drvColorEcho "[INFO]Failed to install driver package, please retry after uninstall and reboot!"
fi

installType="full"
installed_version=`cat "$Driver_Install_Path_Param"/driver/version.info | awk -F '=' '$1=="Version" {print $2}'`
runfilename="Ascend-driver-"$installed_version
all_parma=none
statusUpdate
if [ $driver_install_status -eq 0 ];then
    logOperation "${operation}" "${start_time}" "${runfilename}" "${LOG_RESULT_SUCCESS}" "${installType}" "${all_parma}"
    exitLog 0
else
    logOperation "${operation}" "${start_time}" "${runfilename}" "${LOG_RESULT_FAILED}" "${installType}" "${all_parma}"
    exitLog 1
fi