#!/bin/bash

function print_help() {
echo "\
${1:-TiDB cluster runmode switch script.}
Usage:
  ./runmode.sh [OPTIONS] [ARG]

Options:
   -m              The mode tidb app to convert,{debug|normal},default:debug
   -n              Namespace of the TidbCluster
   -a              The tidb app name {tikv|tidb|pd}
   -t              Resource type {pod|tidbcluster}，if choose 'pod' only change the mode of the pod, if you choose 'tidbcluster', change all pods of the specified app under the tidb cluster
   -h              Display the help message

ARG:
    The pod name or tidb cluster name

Example:
    1. change a tikv pod's mode to debug                        $./runmode.sh -n tidb-cluster -m debug -t pod -a tikv {tikvPodName}
    2. change  a tidb cluster all tikv pods's mode to debug     $./runmode.sh -n tidb-cluster -m debug -t tidbcluster -a tikv {tidbclusterName}
"
}

if [[ $# -eq 0 ]]
then
    print_help
    exit 0
fi

optstring=":m:n:a:t:h"

while getopts "$optstring" opt; do
	case $opt in
		m)
			MODE=${OPTARG:-debug}
			;;
		n)
			NAMESPACE=${OPTARG:-default}
			;;
		a)
		    APP=${OPTARG:-tikv}
		    ;;
		t)
			TYPE=${OPTARG:-pod}
			;;
		h)
			print_help
			exit 0
			;;
		\?)
			"Invalid option: -$OPTARG" >&2
			exit 1
			;;
		:)
			"Option -$OPTARG requires an argument" >&2
			exit 1
			;;
	esac
done

shift $(($OPTIND - 1))

if [[ $# -ne 1 ]]
then
    echo "Please specify a resource name,Usage:runmode.sh -h"
    exit 1
fi

TARGET=$1

if [ X${TYPE} == X"tidbcluster" ];then
    info="the cluster:[${NAMESPACE}/${TARGET}]'s pods"
else
    info="pod:[${NAMESPACE}/${TARGET}]"
fi

read -p "Do you want to change ${info} to ${MODE} mode. please enter yes or no: "
if [ X${REPLY} == X"no" ];then
   exit 0
fi

if [ X${MODE} == X"debug" ]
then
	patch='{"metadata":{"labels":{"runmode":"debug"}}}'
    case ${APP} in
        tikv)
	        current_process="tikv-server"
	        ;;
        tidb)
	        current_process="tidb-server"
	        ;;
	    pd)
	        current_process="pd-server"
	        ;;
	    *)
	        echo $"Usage: -a {tikv|tidb|pd}"
            exit 1
    esac
else
	patch='{"metadata":{"labels":{"runmode":"normal"}}}'
	current_process="tail"
fi

if [ X${TYPE} == X"tidbcluster" ];then
    pods=(`kubectl get po --no-headers -n ${NAMESPACE} -l cluster.pingcap.com/tidbCluster=${TARGET} -l cluster.pingcap.com/app=${APP}|cut -d' ' -f1`)
    if [[ $? -ne 0 ]];then
        echo "kubectl get ${NAMESPACE}/${TARGET}'s pods failed"
        exit 1
    fi
else
    pods=("${TARGET}")
fi

echo "/******** 1.patch pods labels **********/"
success=0
for pod in ${pods[@]}
do
    kubectl patch po/${pod} -n ${NAMESPACE} --patch ${patch}
    if [[ $? -ne 0 ]];then
        echo "patch namespace: ${NAMESPACE}, pod: ${pod},patch: ${patch} failed" >&2
    else
        echo "patch namespace: ${NAMESPACE}, pod: ${pod},patch: ${patch} successfully"
        let success++
    fi
done
if [[ ${success} -ne ${#pods[@]} ]];then
    echo "patch all pods's labels does not completely successful" >&2
    exit 1
fi


echo "/******** 2.check pod's file refresh to ${MODE} **********/"
count=0
while true
do
    success=0
    for pod in ${pods[@]}
    do
        current_mode=`kubectl exec ${pod} -n ${NAMESPACE} -c ${APP} -- cat /etc/k8s-pod-info/labels|grep runmode |awk -F\" '{print $2}'`
        if [[ ${current_mode} == ${MODE} ]];then
            echo "pod: ${NAMESPACE}/${pod}'s mode refresh to ${MODE} successfully."
            let success++
        fi
    done

    if [[ ${success} -eq ${#pods[@]} ]];then
        echo "all pods's mode is refresh to ${MODE}" >&2
        break
    fi

    if [[ ${count} -eq 30 ]];then
        echo "pods's mode refresh failed." >&2
        exit 1
    fi

    echo "only ${success} pods refresh to ${MODE} mode,please wait"
    sleep 5
    let count++
done

echo "/******** 3.kill the  ${current_process} process **********/"
for pod in ${pods[@]}
do
    kubectl -n ${NAMESPACE} exec -it ${pod} -- killall ${current_process}
    if [[ $? -ne 0 ]];then
        echo "kill the ${current_process} process failed, use command: kubectl -n ${NAMESPACE} exec -it ${pod} -- killall ${current_process}" >&2
    else
        echo "kill the ${current_process} process successfully, use command: kubectl -n ${NAMESPACE} exec -it ${pod} -- killall ${current_process}"
    fi
done