#! /bin/bash

HIVE_CUSTOM_LIBS=/usr/lib/customhivelibs
ACTIVEAMBARIHOST=headnodehost
AMBARICONFIGS_SH=/var/lib/ambari-server/resources/scripts/configs.py
PORT=8080

# Import the helper method module.
wget -O /tmp/HDInsightUtilities-v01.sh -q https://hdiconfigactions.blob.core.windows.net/linuxconfigactionmodulev01/HDInsightUtilities-v01.sh && source /tmp/HDInsightUtilities-v01.sh && rm -f /tmp/HDInsightUtilities-v01.sh
########################################
usage() {
    echo ""
    echo "Usage: sudo bash setupcustomhivelibs.sh <HIVE_LIBS_WASB_DIR>";
    echo "       [HIVE_LIBS_WASB_DIR]: Mandatory WASB directory where Hive additional jars are stored. This directory has to be accessible from the cluster. e.g. wasb://hivecontainer@mystorage.blob.core.windows.net/jars/"
    exit 132;
}

checkHostNameAndSetClusterName() {
    fullHostName=$(hostname -f)
    echo "fullHostName=$fullHostName"
    if [ `test_is_zookeepernode` == 1 ]; then
        echo  "Setting up Hive libraries only needs to run on headnode and workernode, exiting ..."
        exit 0
    fi
    CLUSTERNAME=$(echo -e "import hdinsight_common.ClusterManifestParser as ClusterManifestParser\nprint ClusterManifestParser.parse_local_manifest().deployment.cluster_name" | python)
    if [ $? -ne 0 ]; then
        echo "[ERROR] Cannot determine cluster name. Exiting!"
        exit 133
    fi
    echo "Cluster Name=$CLUSTERNAME"
}

validateUsernameAndPassword() {
    coreSiteContent=$(python $AMBARICONFIGS_SH --user=$USERID --password=$PASSWD --action=get --host=$ACTIVEAMBARIHOST --cluster=$CLUSTERNAME --config-type=core-site)
    if [[ $coreSiteContent == *"[ERROR]"* && $coreSiteContent == *"Bad credentials"* ]]; then
        echo "[ERROR] Username and password are invalid. Cannot connect to Ambari Server. Exiting!"
        exit 134
    fi
}

updateAmbariConfigs() {
    echo "Updating Ambari configurations"
    currentHiveAuxJarsPath=$(python $AMBARICONFIGS_SH --user=$USERID --password=$PASSWD --action=get --host=$ACTIVEAMBARIHOST --cluster=$CLUSTERNAME --config-type=hive-site | grep 'hive.aux.jars.path' | sed -n -e 's/.*: "\([^"]*\)".*/\1/p')
    
    if [ -z "$currentHiveAuxJarsPath" ]; then
        newJars=$JARS
    else
        newJars=$currentHiveAuxJarsPath,$JARS
    fi
    
    echo "newJars=$newJars"
    updateResult=$(python $AMBARICONFIGS_SH --user=$USERID --password=$PASSWD --action=set --host=$ACTIVEAMBARIHOST --cluster=$CLUSTERNAME --config-type=hive-site -k "hive.aux.jars.path" -v "$newJars")
    
    if [[ $updateResult != *"Tag:version"* ]] && [[ $updateResult == *"[ERROR]"* ]]; then
        echo "[ERROR] Failed to update hive-site. Exiting!"
        echo $updateResult
        exit 135
    fi
    
    currentHiveEnvContent=$(python $AMBARICONFIGS_SH --user=$USERID --password=$PASSWD --action=get --host=$ACTIVEAMBARIHOST --cluster=$CLUSTERNAME --config-type=hive-env | grep '"content"' | cut -c 17- | rev | cut -c 4- | rev)
    
    exportAuxJars="\nexport HIVE_AUX_JARS_PATH=\$HIVE_AUX_JARS_PATH:$JARS_COLON_SEPARATED"
    newHiveEnvContent=$(echo "$currentHiveEnvContent$exportAuxJars" | sed 's/\\n/\n/g' | sed 's/\\//g')
    
    echo "newHiveEnvContent=$newHiveEnvContent"
    
    updateResult=$(python $AMBARICONFIGS_SH --user=$USERID --password=$PASSWD --action=set --host=$ACTIVEAMBARIHOST --cluster=$CLUSTERNAME --config-type=hive-env -k "content" -v "${newHiveEnvContent}")
    
    if [[ $updateResult != *"Tag:version"* ]] && [[ $updateResult == *"[ERROR]"* ]]; then
        echo "[ERROR] Failed to update hive-env. Exiting!"
        echo $updateResult
        exit 136
    fi
    
    currentWebHCatEnvContent=$(python $AMBARICONFIGS_SH --user=$USERID --password=$PASSWD --action=get --host=$ACTIVEAMBARIHOST --cluster=$CLUSTERNAME --config-type=webhcat-env | grep '"content"' | cut -c 17- | rev | cut -c 2- | rev)
    newWebHCatEnvContent=$(echo "$currentWebHCatEnvContent$exportAuxJars" | sed 's/\\n/\n/g' | sed 's/\\//g')
    
    echo "newWebHCatEnvContent=$newWebHCatEnvContent"
    
    updateResult=$(python $AMBARICONFIGS_SH --user=$USERID --password=$PASSWD --action=set --host=$ACTIVEAMBARIHOST --cluster=$CLUSTERNAME --config-type=webhcat-env -k "content" -v "${newWebHCatEnvContent}")
    
    if [[ $updateResult != *"Tag:version"* ]] && [[ $updateResult == *"[ERROR]"* ]]; then
        echo "[ERROR] Failed to update webhcat-env. Exiting!"
        echo $updateResult
        exit 143
    fi
}

stopServiceViaRest() {
    if [ -z "$1" ]; then
        echo "Need service name to stop service"
        exit 137
    fi
    SERVICENAME=$1
    echo "Stopping $SERVICENAME"
    curl -u $USERID:$PASSWD -i -H 'X-Requested-By: ambari' -X PUT -d '{"RequestInfo": {"context" :"Stop Service for updating Hive libraries"}, "Body": {"ServiceInfo": {"state": "INSTALLED"}}}' http://$ACTIVEAMBARIHOST:$PORT/api/v1/clusters/$CLUSTERNAME/services/$SERVICENAME
}

startServiceViaRest() {
    if [ -z "$1" ]; then
        echo "Need service name to start service"
        exit 138
    fi
    sleep 2
    SERVICENAME=$1
    echo "Starting $SERVICENAME using a background process."
    nohup bash -c "sleep 90; curl -u $USERID:'$PASSWD' -i -H 'X-Requested-By: ambari' -X PUT -d '{\"RequestInfo\": {\"context\" :\"Start Service for updating Hive libraries\"}, \"Body\": {\"ServiceInfo\": {\"state\": \"STARTED\"}}}' http://$ACTIVEAMBARIHOST:$PORT/api/v1/clusters/$CLUSTERNAME/services/$SERVICENAME" > /tmp/Start$SERVICENAME.out 2> /tmp/Start$SERVICENAME.err < /dev/null &
}

########################################
if [ "$(id -u)" != "0" ]; then
    echo "[ERROR] The script has to be run as root."
    usage
fi

USERID=$(echo -e "import hdinsight_common.Constants as Constants\nprint Constants.AMBARI_WATCHDOG_USERNAME" | python)

echo "USERID=$USERID"

PASSWD=$(echo -e "import hdinsight_common.ClusterManifestParser as ClusterManifestParser\nimport hdinsight_common.Constants as Constants\nimport base64\nbase64pwd = ClusterManifestParser.parse_local_manifest().ambari_users.usersmap[Constants.AMBARI_WATCHDOG_USERNAME].password\nprint base64.b64decode(base64pwd)" | python)

HIVE_LIBS_WASB_DIR=$1

if [ -z "$HIVE_LIBS_WASB_DIR" ]; then
    echo "[ERROR] No hive lib wasb directory specified. Exiting!"
    usage
fi

if [ -e $HIVE_CUSTOM_LIBS ]; then
    echo "Custom Hive Libs directory already exists. Exiting..."
    exit 0
fi

checkHostNameAndSetClusterName
validateUsernameAndPassword

echo mkdir -p $HIVE_CUSTOM_LIBS
mkdir -p $HIVE_CUSTOM_LIBS

echo hadoop fs -copyToLocal $HIVE_LIBS_WASB_DIR/*.* $HIVE_CUSTOM_LIBS
hadoop fs -copyToLocal $HIVE_LIBS_WASB_DIR/*.* $HIVE_CUSTOM_LIBS

if [ $? -ne 0 ]; then
    echo "[ERROR] Cannot copy libraries from $HIVE_LIBS_WASB_DIR to local directory. Cleaning up $HIVE_CUSTOM_LIBS and exiting!"
    rm -rf $HIVE_CUSTOM_LIBS
    exit 139
fi

JARS=`find $HIVE_CUSTOM_LIBS -iname '*.jar' -type f | sed -e 's|^|file://|' | paste -d, -s`
JARS_COLON_SEPARATED=`find $HIVE_CUSTOM_LIBS -iname '*.jar' -type f | paste -d: -s`

if [ -z "$JARS" ]; then
    echo "[ERROR] No jars file found in $HIVE_CUSTOM_LIBS after copying. Listing current content of $HIVE_CUSTOM_LIBS before cleaning up $HIVE_CUSTOM_LIBS and exiting!"
    find $HIVE_CUSTOM_LIBS
    rm -rf $HIVE_CUSTOM_LIBS
    exit 140
fi
echo "Copy completed."

echo "Updating Ambari configs and restarting services from primary headnode"
PRIMARYHEADNODE=`get_primary_headnode`
PRIMARY_HN_NUM=`get_primary_headnode_number`

#Check if values retrieved are empty, if yes, exit with error
if [[ -z $PRIMARYHEADNODE ]]; then
	echo "Could not determine primary headnode."
	exit 141
fi

if [[ -z "$PRIMARY_HN_NUM" ]]; then
	echo "Could not determine primary headnode number."
	exit 142
fi

fullHostName=$(hostname -f)
echo "fullHostName=$fullHostName. Lower case: ${fullHostName,,}"
echo "primary headnode=$PRIMARYHEADNODE. Lower case: ${PRIMARYHEADNODE,,}"
if [ "${fullHostName,,}" == "${PRIMARYHEADNODE,,}" ]; then
    updateAmbariConfigs
    stopServiceViaRest HIVE
    stopServiceViaRest OOZIE
    startServiceViaRest HIVE
    startServiceViaRest OOZIE
fi