#!/bin/bash
set -o nounset
set -o errexit
set -o pipefail
#set -o xtrace

BASE=$(cd `dirname $0`; cd ../; pwd)
DEFAULT_CLS=${CLUSTER:-cluster01}

source $BASE/utils/logger
source $BASE/utils/env

function docker_exec_kubeasz() {
    local tty_arg="-it"
    if [[ "${KUBEMORE_DOCKER_EXEC_TTY:-1}" = "0" ]]; then
        tty_arg=""
    fi
    docker exec ${tty_arg} kubeasz "$@"
}

MASTER_ROLE_LAB=kubemore.io/role.master
ETCD_ROLE_LAB=kubemore.io/role.etcd
WORKER_ROLE_LAB=kubemore.io/role.worker
MASTERS=
ETCDS=
WORKERS=
PASSWD=
BASE_BIN=/opt/kube/bin
SSH_USER=root
SSH_PORT=22
XPU_MODE=
XPU_TYPE=
WORKER_XPU_MODE_LAB=kubemore_xpu_device_mode
WORKER_XPU_TYPE_LAB=kubemore_xpu_type


function print_help(){
    echo "./kubemore add --masters 192.168.1.3,192.168.1.4 \\"
    echo "               --etcd 192.168.1.3,192.168.1.4 \\"
    echo "               --workers 192.168.1.5,192.168.1.6"
    echo " Options:"
    echo "   m|masters, required:  cluster master nodes"
    echo "   e|etcds, required:    cluster etcd nodes"
    echo "   w|workers, requilred: cluster worker nodes"
    echo "   password:             password for all nodes"
    echo "   ssh-user:             ssh user for all nodes, default root"
    echo "   ssh-port:             ssh port for all nodes, default 22"
    echo "   xpu-mode:             GPU/DPU mode, available mode [nvidia|ascend910|enflame|biren], "
    echo "                         only set to worker node"
    echo "   xpu-type:             GPU/DPU type, available type, "
    echo "                         for nvidia: [nvidia-v10032g|nvidia-v10016g|nvidia-t4|"
    echo "                                      nvidia-a10040g|nvidia-a10080g|nvidia-h200141g],"
    echo "                         for ascend910: [ascend-ascend910a32g|ascend-ascend910b64g]"
    echo "                         for enflame: [enflame-s60]"
    echo "                         for biren: [biren-106m]"
    echo "   h|help:               print help usage"
    echo ""
}


TEMP=`getopt -o m:e:w:h -l masters:,etcds:,workers:,password:,ssh-user:,ssh-port:,xpu-mode:,xpu-type:,help -- "$@"`
if [[ ${?} != 0 ]] ; then error "Bad command" >&2 ; print_help ; exit 1 ; fi

eval set -- "${TEMP}"

while true; do
  case "$1" in
  -m|--masters)
    MASTERS=$2
    check_empty "${MASTERS}" "master nodes"
    shift 2
    ;;
  -e|--etcds)
    ETCDS=$2
    check_empty "${ETCDS}" "etcd nodes"
    shift 2
    ;;
  -w|--workers)
    WORKERS=$2
    check_empty "${WORKERS}" "worker nodes"
    shift 2
    ;;
  --password)
    PASSWD=$2
    check_empty "${PASSWD}" "password for nodes"
    shift 2
    ;;
  --ssh-user)
    SSH_USER=$2
    check_empty "${SSH_USER}" "ssh user for nodes"
    shift 2
    ;;
  --ssh-port)
    SSH_PORT=$2
    check_empty "${SSH_PORT}" "ssh port for nodes"
    shift 2
    ;;
  --xpu-mode)
    XPU_MODE=$2
    check_empty "${XPU_MODE}" "xpu mode"
    shift 2
    ;;
  --xpu-type)
    XPU_TYPE=$2
    check_empty "${XPU_TYPE}" "xpu type"
    shift 2
    ;;
  -h|--help)
    print_help
    shift
    exit 0
    ;;
  --)
    shift
    break
    ;;
  *)
    error "Bad command $1"
    print_help
    exit 1
    ;;
  esac
done

# Add etcd node
if [[ "${SSH_USER}" != "root" ]]; then
    logger info "Enable ansible sudo mode for ssh user: ${SSH_USER}"
    cp $BASE/ansible.cfg.not_root $BASE/ansible.cfg
fi
for i in $(echo $ETCDS| sed "s/,/ /g"); do
    item="$i ansible_ssh_user=${SSH_USER} ansible_ssh_port=${SSH_PORT}"
    if [[ "$PASSWD" != "" ]]; then
        item="$item ansible_ssh_pass=$PASSWD"
    fi
    if [[ "$SSH_USER" != "root" && "$PASSWD" != "" ]]; then
        item="$item ansible_become_pass=${PASSWD}"
    fi
    docker_exec_kubeasz morectl add-etcd ${DEFAULT_CLS} $item
done

# Add master node
for i in $(echo $MASTERS| sed "s/,/ /g"); do
    item="$i ansible_ssh_user=${SSH_USER} ansible_ssh_port=${SSH_PORT}"
    if [[ "$PASSWD" != "" ]]; then
        item="$item ansible_ssh_pass=$PASSWD"
    fi
    if [[ "$SSH_USER" != "root" && "$PASSWD" != "" ]]; then
        item="$item ansible_become_pass=${PASSWD}"
    fi
    docker_exec_kubeasz morectl add-master ${DEFAULT_CLS} $item
done

# Add worker node
for i in $(echo $WORKERS| sed "s/,/ /g"); do
    item="$i ansible_ssh_user=${SSH_USER} ansible_ssh_port=${SSH_PORT}"
    if [[ "$PASSWD" != "" ]]; then
        item="$item ansible_ssh_pass=$PASSWD"
    fi
    if [[ "$SSH_USER" != "root" && "$PASSWD" != "" ]]; then
        item="$item ansible_become_pass=${PASSWD}"
    fi
    docker_exec_kubeasz morectl add-node ${DEFAULT_CLS} $item
done
if [[ "$PASSWD" != "" && -f $BASE/clusters/${DEFAULT_CLS}/hosts ]]; then
    chmod 0600 $BASE/clusters/${DEFAULT_CLS}/hosts
fi

logger info "Set node roles"
for m in $(echo "${MASTERS}" | awk -F ',' '{for(i=1; i<=NF; i++){ print $i }}');do
    ${BASE_BIN}/kubectl label node $m "${MASTER_ROLE_LAB}=1"
done
for e in $(echo "${ETCDS}" | awk -F ',' '{for(i=1; i<=NF; i++){ print $i }}');do
    ${BASE_BIN}/kubectl label node $e "${ETCD_ROLE_LAB}=1"
done
for w in $(echo "${WORKERS}" | awk -F ',' '{for(i=1; i<=NF; i++){ print $i }}');do
    ${BASE_BIN}/kubectl label node $w "${WORKER_ROLE_LAB}=1"
done

logger info "Set node xpu mode label"
if [[ "$XPU_MODE" != "" ]]; then
    for w in $(echo "${WORKERS}" | awk -F ',' '{for(i=1; i<=NF; i++){ print $i }}');do
        ${BASE_BIN}/kubectl label node $w "${WORKER_XPU_MODE_LAB}=${XPU_MODE}"
    done
fi

logger info "Set node xpu type label"
if [[ "$XPU_TYPE" != "" ]]; then
    for w in $(echo "${WORKERS}" | awk -F ',' '{for(i=1; i<=NF; i++){ print $i }}');do
        ${BASE_BIN}/kubectl label node $w "${WORKER_XPU_TYPE_LAB}=${XPU_TYPE}"
    done
fi
