#!/bin/bash
#set -e

function logger() {
  TIMESTAMP=$(date +'%Y-%m-%d %H:%M:%S')
  case "$1" in
    debug)
      echo -e "$TIMESTAMP \033[36mDEBUG\033[0m $2"
      ;;
    info)
      echo -e "$TIMESTAMP \033[32mINFO\033[0m $2"
      ;;
    warn)
      echo -e "$TIMESTAMP \033[33mWARN\033[0m $2"
      ;;
    error)
      echo -e "$TIMESTAMP \033[31mERROR\033[0m $2"
      ;;
    *)
      ;;
  esac
}

function 01_label() {
  logger debug "01_为主机打标签"
  #重启fluentd
  $BASE/bin/kubectl delete pods -n kube-system  -l k8s-app=fluentd-es --force
  master_ips=$(grep -E 'MGR_NODE_IP' $BASE/ai.config|awk -F " " '{print $1}'|awk -F "=" '!($2 in IP){IP[$2];print $2}')
  for master_ip in $master_ips
  do
    $BASE/bin/kubectl label node  $master_ip masterselector=dls-master-node  --overwrite
  done

  #为x86主机打标签
  x86_ips=$(grep -E 'MGR_NODE_IP|COMPUTE_NODE_IP' $BASE/ai.config|awk -F " " '{print $1}'|awk -F "=" '!($2 in IP){IP[$2];print $2}')
  for x86_ip in $x86_ips
  do
    $BASE/bin/kubectl label node  $x86_ip node-role.kubernetes.io/worker=worker workerselector=dls-worker-node host-arch=huawei-arm  accelerator-type=card ai.machine.role=plugins  machine=x86 app=explorer --overwrite
  done

  #为gpu主机打标签
  gpu_ips=$(grep -E 'NPU_TYPE' $BASE/ai.config|grep -v '^#'|awk -F " " '{print $1}'|awk -F "=" '!($2 in IP){IP[$2];print $2}')
  for gpu_ip in $gpu_ips
  do
    $BASE/bin/kubectl label node $gpu_ip type=gpu --overwrite
    gpu_share=$(grep -E 'NPU_TYPE' $BASE/ai.config|grep -v '^#'|grep -E $gpu_ip|awk -F " " '{print $2}'|awk -F "=" '!($2 in gpu_share){gpu_share[$2];print $2}')
    $BASE/bin/kubectl label node $gpu_ip gpushare=$gpu_share --overwrite
    npu_type=$(grep -E 'NPU_TYPE' $BASE/ai.config|grep -v '^#'|grep -E $gpu_ip|awk -F " " '{print $3}'|awk -F "=" '!($2 in gpu_type){gpu_type[$2];print $2}')
    $BASE/bin/kubectl label node $gpu_ip gpu=Ascend$npu_type --overwrite
    #如果npu_type包含910，则打上标签huawei-Ascend910 如果包含310，则打上标签huawei-Ascend310
    if [[ $npu_type == *910* ]];then
      sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-device-plugin_5.0.0/device-plugin-910-v5.0.0.yaml
      $BASE/bin/kubectl label node $gpu_ip accelerator=huawei-Ascend910 --overwrite
      $BASE/bin/kubectl apply -f $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-device-plugin_5.0.0/device-plugin-910-v5.0.0.yaml
    elif [[ $npu_type == *310* ]];then
      sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-device-plugin_5.0.0/device-plugin-310-v5.0.0.yaml 
      $BASE/bin/kubectl label node $gpu_ip accelerator=huawei-Ascend310 --overwrite
      $BASE/bin/kubectl apply -f $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-device-plugin_5.0.0/device-plugin-310-v5.0.0.yaml
    fi
  done
}

function 02_mindl-dl_true() {
  edition=$(grep -r 'EDITION' $BASE/ai.config |cut -d'=' -f2)
  if [[ $edition == professional ]];then
    logger debug "02_安装mindx_dl插件"
    $BASE/bin/kubectl  create ns mindx-dl
    echo "REGISTRY_LOCAL: $REGISTRY_LOCAL"
    sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-device-plugin_5.0.0/device-plugin-910-v5.0.0.yaml
    sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-volcano_5.0.0/volcano-v1.4.0/volcano-v1.4.0.yaml
    sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-noded_5.0.0/noded-v5.0.0.yaml
    sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-hccl-controller_5.0.0/hccl-controller-v5.0.0.yaml
    sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-npu-exporter_5.0.0/npu-exporter-v5.0.0.yaml
    $BASE/bin/kubectl apply -f $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-device-plugin_5.0.0/device-plugin-910-v5.0.0.yaml
    $BASE/bin/kubectl apply -f $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-volcano_5.0.0/volcano-v1.4.0/volcano-v1.4.0.yaml
    $BASE/bin/kubectl apply -f $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-noded_5.0.0/noded-v5.0.0.yaml
    $BASE/bin/kubectl apply -f $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-hccl-controller_5.0.0/hccl-controller-v5.0.0.yaml
    $BASE/bin/kubectl apply -f $BASE/ai/ai-manifests/mindx-dl/Ascend-mindxdl-npu-exporter_5.0.0/npu-exporter-v5.0.0.yaml
  fi
}

function 06_mysql() {
  logger debug "06_安装ai-mysql和minio"
  mysql_node_ip=$(grep -E 'MYSQL_NODE_IP' $BASE/ai.config | awk -F "=" '{print $2}')
  mysql_root_password=$(grep -E 'MYSQL_NODE_PASSWORD' $BASE/ai.config | awk -F "=" '{print $2}')
  minio_user=$(grep -E 'MINIO_ACCESS_KEY' $BASE/ai.config | awk -F "=" '{print $2}')
  minio_password=$(grep -E 'MINIO_SECRET_KEY' $BASE/ai.config | awk -F "=" '{print $2}')
  ai_mysql_network=$(grep -E 'AI_MYSQL_NETWORK' $BASE/ai.config | awk -F "=" '{print $2}')
  #修改ai.cfg配置文件
  /bin/cp -f $BASE/ai/ai-manifests/ai-app-arm64/ai.cfg.tmp $BASE/ai/ai-manifests/ai-app-arm64/ai.cfg
  sed -i "s/MYSQL_NODE_IP/$mysql_node_ip/g" $BASE/ai/ai-manifests/ai-app-arm64/ai.cfg
  sed -i "s/MYSQL_PASSWORD/$mysql_root_password/g" $BASE/ai/ai-manifests/ai-app-arm64/ai.cfg
  sed -i "s/MINIO_USER/$minio_user/g" $BASE/ai/ai-manifests/ai-app-arm64/ai.cfg
  sed -i "s/MINIO_PASSWORD/$minio_password/g" $BASE/ai/ai-manifests/ai-app-arm64/ai.cfg
  /bin/cp -f $BASE/ai/ai-manifests/ai-app-arm64/ai.cfg $BASE/ai/ai-manifests/ai-app-arm64/common/config/minio/env
  /bin/cp -f $BASE/ai/ai-manifests/ai-app-arm64/ai.cfg $BASE/ai/ai-manifests/ai-app-arm64/common/config/mysql/env
  #修改docker-compose.yaml
  /bin/cp -f $BASE/ai/ai-manifests/ai-app-arm64/docker-compose.yml.tmp $BASE/ai/ai-manifests/ai-app-arm64/docker-compose.yml
  sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/ai-app-arm64/docker-compose.yml
  sed -i "s#AI_MYSQL_NETWORK#$ai_mysql_network#g" $BASE/ai/ai-manifests/ai-app-arm64/docker-compose.yml
  #执行安装
  cd $BASE/ai/ai-manifests/ai-app-arm64/ && $BASE/bin/docker-compose up -d
  logger debug "初始化minio"
  if [[ -d /opt/minio/ai-pic ]];then
    logger info "minio已初始化"
  else
    #生成Minio bucket
    sleep 15
    docker exec -it ai-minio  /bin/sh /init.sh
  fi

  #创建运维平台数据库
  if [[ ! -d /opt/db/ai_kop ]];then
    logger debug "正在创建运维平台数据库"
    sleep 15
    $BASE/bin/docker exec -it ai-mysql bash -c "mysql -uroot -p$mysql_root_password -e 'create DATABASE ai_kop CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci'"
  fi
}

function 07_jwt() {
  edition=$(grep -r 'EDITION' $BASE/ai.config |cut -d'=' -f2)
  if [[ $edition == professional ]];then
    logger debug "07_安装Kubeflow JWT"
    mysql_node_ip=$(grep -E 'MYSQL_NODE_IP' $BASE/ai.config | awk -F "=" '{print $2}')
    mysql_root_password=$(grep -E 'MYSQL_NODE_PASSWORD' $BASE/ai.config | awk -F "=" '{print $2}')

    /bin/cp -f $BASE/ai/ai-manifests/yaml/arm64/jwt.yaml.tmp $BASE/ai/ai-manifests/yaml/arm64/jwt.yaml
    sed -i "s/DB_URL_VALUE/$mysql_node_ip/g"  $BASE/ai/ai-manifests/yaml/arm64/jwt.yaml
    sed -i "s/DB_PASSWORD_VALUE/$mysql_root_password/g"  $BASE/ai/ai-manifests/yaml/arm64/jwt.yaml
    sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/yaml/arm64/jwt.yaml
    $BASE/bin/kubectl apply -f $BASE/ai/ai-manifests/yaml/arm64/jwt.yaml
  fi
}

function 08_event() {
  #8、安装审计组件
  logger debug "08_安装审计组件"
  mysql_node_ip=$(grep -E 'MYSQL_NODE_IP' $BASE/ai.config | awk -F "=" '{print $2}')
  mysql_root_password=$(grep -E 'MYSQL_NODE_PASSWORD' $BASE/ai.config | awk -F "=" '{print $2}')
  /bin/cp -f $BASE/ai/ai-manifests/yaml/arm64/k8s-event.yaml.tmp $BASE/ai/ai-manifests/yaml/arm64/k8s-event.yaml
  sed -i "s/DB_URL/$mysql_node_ip/g" $BASE/ai/ai-manifests/yaml/arm64/k8s-event.yaml
  sed -i "s/DB_PASSWORD/$mysql_root_password/g" $BASE/ai/ai-manifests/yaml/arm64/k8s-event.yaml
  sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/yaml/arm64/k8s-event.yaml
  $BASE/bin/kubectl apply -f $BASE/ai/ai-manifests/yaml/arm64/k8s-event.yaml
}

function 10_kubeflow() {

  #10、安装kubeflow 1.7
  edition=$(grep -r 'EDITION' $BASE/ai.config |cut -d'=' -f2)
  if [[ $edition == professional ]];then
    logger info "安装kebeflow 1.7"
    sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g"  `grep -rl 'REGISTRY_LOCAL' /opt/ai-deployer/ai/ai-manifests/yaml/arm64/kubeflow-1.7-arm64/`
    #循环执行下面的语句5次，间隔10s 
    for time in {1..5}
    do
        $BASE/bin/kubectl apply -R -f $BASE/ai/ai-manifests/yaml/arm64/kubeflow-1.7-arm64/
        sleep 5
    done
  fi
}

function 11_backup() {
  isBackup=$(cat /etc/crontab)
  if [[ $isBackup =~ backup.sh ]];then
    logger info "已开启备份"
  else
    echo "0 0 * * * root /opt/ai-deployer/backup.sh" >> /etc/crontab
    service cron restart
  fi

}

function 12_delete_evicted_pod() {
  /bin/cp -f $BASE/ai/ai-manifests/yaml/delete-evicted-pods.yaml.tmp $BASE/ai/ai-manifests/yaml/delete-evicted-pods.yaml
  sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/yaml/delete-evicted-pods.yaml
  $BASE/bin/kubectl apply -f $BASE/ai/ai-manifests/yaml/delete-evicted-pods.yaml
}

function 13_nvidia_smi_process() {
  /bin/cp -f $BASE/ai/ai-manifests/yaml/nvidia-smi-process.yml.tmp $BASE/ai/ai-manifests/yaml/nvidia-smi-process.yml
  sed -i "s/REGISTRY_LOCAL/$REGISTRY_LOCAL/g" $BASE/ai/ai-manifests/yaml/nvidia-smi-process.yml
  $BASE/bin/kubectl apply -f $BASE/ai/ai-manifests/yaml/nvidia-smi-process.yml
}
  
function usage() {
    echo -e "\033[33mUsage:\033[0m ai-install.sh COMMAND"
    cat <<EOF
-------------------------------------------------------------------------------------
  all       to install all services for ai
  label     to label k8s nodes for ai
  gpushare  to install  gpushare device for ai
  mysql     to install mysql for ai
  jwt       to install jwt for ai
  event     to install event for ai
  volcano   to install volcano for ai
  kubeflow  to install kubeflow for ai
EOF
}

#命令执行成功@?为0；函数执行成功return 1
function main() {
  BASE=/opt/ai-deployer
  HARBOR_IP=$(grep -r 'HARBOR_DOMAIN' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  HARBOR_PORT=$(grep -r 'HARBOR_NODE_PORT' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  REGISTRY_LOCAL="$HARBOR_IP:$HARBOR_PORT"

  [ "$#" -gt 0 ] || { usage >&2; exit 2; }

  case "$1" in
      (all)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          01_label
	        02_mindl-dl_true
          06_mysql
          07_jwt
          08_event
        # 09_volcano
          10_kubeflow
        # 05_GPUshare_vGPU
          11_backup
         # 12_delete_evicted_pod
      	 # 13_nvidia_smi_process
          ;;
      (label)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          01_label
          ;;
      (gpushare)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          03_GPUshare_true
          04_GPUshare_false
          05_GPUshare_vGPU
          ;;
      (mysql)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          06_mysql
          ;;
      (jwt)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          07_jwt
          ;;
      (event)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          08_event
          ;;
      (volcano)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          09_volcano
          ;;
      (kubeflow)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          10_kubeflow
          ;;
      (backup)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          11_backup
          ;;
      (*)
          usage
          exit 0
          ;;
  esac
}

main "$@"