#!/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=
SERVICE_CIDR=10.68.0.0/16
CLUSTER_CIDR=172.20.0.0/16
NTP_SERVER=
PORT_RANGE=30000-32767
MAX_PODS=110
PASSWD=
EXTER_EIP=agione.com
DNS_DOMAIN=cluster.local
SSH_USER=root
SSH_PORT=22
BASE_BIN=/opt/kube/bin
XPU_MODE=
XPU_TYPE=
WORKER_XPU_MODE_LAB=kubemore_xpu_device_mode
WORKER_XPU_TYPE_LAB=kubemore_xpu_type
KUBELET_ROOT_PATH=/var/lib/kubelet
CONTAINERD_ROOT_PATH=/var/lib/containerd

function print_help(){
    echo "./morecli install --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 "   service-cidr:         service cidr, default: 10.68.0.0/16"
    echo "   cluster-cidr:         cluster cidr, default: 172.20.0.0/16"
    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 "   ntp-server:           ntp server for cluster"
    echo "   p|port-range:         node port range, default: 30000-32767"
    echo "   max-pods:             max ports per node, default: 110"
    echo "   exter-eip:            external eip for k8s cluster"
    echo "   d|dns-domain:         cluster dns domain, default: cluster.local"
    echo "   xpu-mode:             GPU/DPU mode, available mode [nvidia|ascend910|enflame|biren|hygon], "
    echo "                         only set to worker node"
    echo "   xpu-type:             GPU/DPU type, available type, "
    echo "                         for nvidia: [nvidia-v10032g|nvidia-v10016g|nvidia-t416g|"
    echo "                                      nvidia-a10040g|nvidia-a10080g|nvidia-h200141g],"
    echo "                         for ascend910: [ascend-ascend910a32g|ascend-ascend910b64g]"
    echo "                         for enflame: [enflame-s6032g]"
    echo "                         for biren: [biren-106m32g]"
    echo "                         for hygon: [hygon-bw20064g]"
    echo "                         only set to worker node"
    echo "   kubelet-root-data:    root data path for kubelet, default '/var/lib/kubelet'. "
    echo "   containerd-root-data: root data path for containerd, default '/var/lib/containerd'. "
    echo "   h|help:               print help usage"
    echo ""
}


TEMP=`getopt -o m:e:w:p:d:h -l masters:,etcds:,workers:,service-cidr:,cluster-cidr:,password:,ssh-user:,ssh-port:,port-range:,ntp-server:,max-pods:,exter-eip:,dns-domain:,xpu-mode:,xpu-type:,kubelet-root-data:,containerd-root-data:,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
    ;;
  -p|--port-range)
    PORTS=$2
    check_empty "${PORTS}" "port range"
    shift 2
    ;;
  --ntp-server)
    NTP_SERVER=$2
    check_empty "${NTP_SERVER}" "ntp server"
    shift 2
    ;;
  --max-pods)
    MAX_PODS=$2
    check_empty "${MAX_PODS}" "max pods"
    shift 2
    ;;
  --exter-eip)
    EXTER_EIP=$2
    check_empty "${EXTER_EIP}" "external eip"
    shift 2
    ;;
  -d|--dns-domain)
    DNS_DOMAIN=$2
    check_empty "${DNS_DOMAIN}" "dns domain"
    shift 2
    ;;
  --service-cidr)
    SERVICE_CIDR=$2
    check_empty "${SERVICE_CIDR}" "service cidir"
    shift 2
    ;;
  --cluster-cidr)
    CLUSTER_CIDR=$2
    check_empty "${CLUSTER_CIDR}" "cluster cidr"
    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
    ;;
  --kubelet-root-data)
    KUBELET_ROOT_PATH=$2
    check_empty "${KUBELET_ROOT_PATH}" "kubelet root data path"
    shift 2
    ;;
  --containerd-root-data)
    CONTAINERD_ROOT_PATH=$2
    check_empty "${CONTAINERD_ROOT_PATH}" "containerd root data path"
    shift 2
    ;;
  -h|--help)
    print_help
    shift
    exit 0
    ;;
  --)
    shift
    break
    ;;
  *)
    error "Bad command $1"
    print_help
    exit 1
    ;;
  esac
done

logger info "[Pre setup] change config files"
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
if [[ -f $BASE/example/hosts.multi-node ]]; then
  rm -rf $BASE/example/hosts.multi-node
fi
cp $BASE/example/hosts.multi-node.ex $BASE/example/hosts.multi-node
if [ -f $BASE/example/config.yml ]; then
  rm -rf $BASE/example/config.yml
fi
cp $BASE/example/config.yml.ex $BASE/example/config.yml
check_empty "${ETCDS}" "etcd nodes"
check_empty "${MASTERS}" "master nodes"
# check_empty "${WORKERS}" "worker nodes"

# Generate etcd nodes config
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
    sed -i "/\[etcd\]/a $item" $BASE/example/hosts.multi-node
done

# Generate master nodes config
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

    sed -i "/\[kube_master\]/a $item" $BASE/example/hosts.multi-node
done

# Generate worker node config
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
    sed -i "/\[kube_node\]/a $item" $BASE/example/hosts.multi-node
done
if [[ "$PASSWD" != "" ]]; then
    chmod 0600 $BASE/example/hosts.multi-node
fi

# Generate ntp server config
if [[ "$NTP_SERVER" != "" ]];then
    sed -i "/\[chrony\]/a $NTP_SERVER" $BASE/example/hosts.multi-node
fi

# Generate service cidr config
sed -i "s#__SERVICE_CIDR__#$SERVICE_CIDR#g" $BASE/example/hosts.multi-node

# Generate cluster cidr config
sed -i "s#__CLUSTER_CIDR__#$CLUSTER_CIDR#g" $BASE/example/hosts.multi-node

# Generate port range config
sed -i "s#__PORT_RNAGE__#$PORT_RANGE#g" ${BASE}/example/hosts.multi-node

# Generate dns domain config
sed -i "s#__DNS_DOMAIN__#$DNS_DOMAIN#g" ${BASE}/example/hosts.multi-node

OS=`check_os`
# Generate arm config
sed -i "s#NVIDIA_RUNTIME_INSTALL_PKG=.*#NVIDIA_RUNTIME_INSTALL_PKG=${OS}#g" ${BASE}/example/hosts.multi-node

if [[ "$OS" = "ubuntu-22.04" ]]; then
  sed -i "s/calicoVer=v3.19.4/calicoVer=v3.24.5/g" ${BASE}/predown
fi

if [[ "$OS" != "centos7" && "$XPU_MODE" = "ascend910" ]]; then
  sed -i "s#RUNTIME_XPU=.*#RUNTIME_XPU=\"no\"#g" ${BASE}/example/hosts.multi-node
else
  sed -i "s#RUNTIME_XPU=.*#RUNTIME_XPU=\"yes\"#g" ${BASE}/example/hosts.multi-node
fi

# Generate external EIP config
if [[ "$EXTER_EIP" != "" ]]; then
    sed -i "s#__EXTER_EIP__#- $EXTER_EIP#g" ${BASE}/example/config.yml
else
    sed -i "s#__EXTER_EIP__##g" ${BASE}/example/config.yml
fi

# Generate max pods config
sed -i "s/__MAX_PODS__/$MAX_PODS/g" ${BASE}/example/config.yml

# Generate kubelet root dir
sed -i "s#KUBELET_ROOT_DIR:.*#KUBELET_ROOT_DIR: \"$KUBELET_ROOT_PATH\"#g" ${BASE}/example/config.yml

# Generate kubelet root dir
sed -i "s#CONTAINERD_STORAGE_DIR:.*#CONTAINERD_STORAGE_DIR: \"$CONTAINERD_ROOT_PATH\"#g" ${BASE}/example/config.yml


logger info "Prepare images and packages"
curr_arch=`check_arch`
if [[ -d $BASE/resources/$curr_arch ]]; then
    if [[ -d $BASE/down ]]; then
        rm -rf $BASE/down
    fi
    ln -s $BASE/resources/$curr_arch/down $BASE/down
    if [[ -d $BASE/bin ]]; then
        rm -rf $BASE/bin
    fi
    ln -s $BASE/resources/$curr_arch/bin $BASE/bin
else
    logger error "Not supported arch: `arch`"
    exit 1
fi
sh -c "bash $BASE/predown -D"
logger info "Prepare ansible service"
sh -c "bash $BASE/predown -S"

logger info "Start to create cluster"
if [[ -d $BASE/clusters ]]; then
    rm -rf $BASE/clusters
fi
docker_exec_kubeasz morectl new ${DEFAULT_CLS}
docker_exec_kubeasz morectl setup ${DEFAULT_CLS} all

if [[ ! -f ${BASE_BIN}/kubectl ]]; then
    cp ${BASE}/bin/kubectl ${BASE_BIN}/kubectl
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

${BASE_BIN}/kubectl label ns kube-system ns=kube-system
