#!/bin/bash
set -e

function usage() {
    echo -e "\033[33mUsage:\033[0m install.sh COMMAND"
    cat <<EOF
-------------------------------------------------------------------------------------
  all             install all services
  k8s             install k8s clusters
  gpudriver       install gpudriver
  harbor          install harbor
  storage         install storage
  nvidia-runtime  install nvidia-container-runtime
  ai              install ai components
  ai-web          install ai-web
EOF
}

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 check_os() {
  RELEASE_OS=$(grep -E 'Release_OS' $BASE/Release |cut -d'=' -f2|awk '{print $1}')
  if [[ `cat /etc/os-release` =~ Ubuntu ]];then
    OS_VERSION=$(grep -E 'VERSION_ID' /etc/os-release | sed "s/'//g;s/\"//g" | cut -d'=' -f2)
    if [[ $RELEASE_OS =~ $OS_VERSION ]];then
      logger info "当前安装包for $RELEASE_OS 支持当前操作系统 ubuntu$OS_VERSION 部署"
    else
      { logger error "当前安装包for $RELEASE_OS 不支持当前操作系统 ubuntu$OS_VERSION 部署"; return 1; }
    fi
  fi
}

function download_gpudriver() {
  GENFORCE_URL=$MINIO_URL/ai-deployer/rtxseries/NVIDIA-Linux-x86_64-550.90.07.run
  TESLA_URL=$MINIO_URL/ai-deployer/tesla/NVIDIA-Linux-x86_64-550.90.07.run
  nvidia_driver_down(){
  case $1 in
    (RTX-4090|RTX-3090|RTX-3080|RTX-3070|RTX-3060|RTX-2080TI|RTX-2080|RTX-2070|RTX-2060|GTX-1080TI|GTX-1080|GTX-1070|GTX-1060|RTX-4000|RTX-5000|RTX-6000)
    if [ -f "$BASE/down/NVIDIA-Linux-x86_64-550.90.07.run" ];then
      logger info "该版本驱动已经下载"
    else
      curl -C- -O --retry 3 "$GENFORCE_URL" || { logger error "downloading gpu driver failed"; exit 1; }
      /bin/mv -f "./NVIDIA-Linux-x86_64-550.90.07.run" "$BASE/down"
    fi
  ;;
   #修改成其他的类型全部用下面的
   # (A100|A800|T4|V100|P40|K80|P100|A40|A30|A10|M60|TITIAN-XP)
    (*)
    if [ -f "$BASE/down/NVIDIA-Linux-x86_64-550.90.07.run" ];then
       logger info "该版本驱动已经下载"
    else
       curl -C- -O --retry 3 "$TESLA_URL" || { logger error "downloading gpu driver failed"; exit 1; }
       /bin/mv -f "./NVIDIA-Linux-x86_64-550.90.07.run" "$BASE/down"
    fi
  ;;
  esac
  }
  GPUTYPES=$(grep -r GPU_TYPE  $BASE/ai.config  | awk -F"="  '!($4 in GPU){GPU[$4];print $4}')
  logger debug "开始下载GPU驱动"
  for gputype in $GPUTYPES
  do
    nvidia_driver_down $gputype
  done
  /bin/cp -f $BASE/down/NVIDIA-Linux*.run $BASE/ai/prepareinstall/roles/gpudriverbefore/files/
}
function download_npudriver() {
  NPU_310_firmware_URL=http://123.157.234.230:9000/ai-deployer/arm/Ascend-hdk-310p-npu-firmware_7.5.0.1.129.run
  NPU_310_DRIVER_URL=http://123.157.234.230:9000/ai-deployer/arm/Ascend-hdk-310p-npu-driver_24.1.rc3_linux-aarch64.run
  NPU_910_firmware_URL=http://123.157.234.230:9000/ai-deployer/arm/Ascend-hdk-910b-npu-firmware_7.1.0.6.220.run
  NPU_910_DRIVER_URL=http://123.157.234.230:9000/ai-deployer/arm/Ascend-hdk-910b-npu-driver_24.1.rc1_linux-aarch64.run
  npu_driver_down(){
    case $1 in
      (310|310P|310P3|310p)
      if [ -f "$BASE/down/Ascend-hdk-310p-npu-firmware_7.5.0.1.129.run" ];then
        logger info "该版本固件已经下载"
      else
        curl -C- -O --retry 3 "$NPU_310_firmware_URL" || { logger error "固件下载失败" ; exit 1; }
        /bin/mv -f "./Ascend-hdk-310p-npu-firmware_7.5.0.1.129.run" "$BASE/down"
      fi
      if [ -f "$BASE/down/Ascend-hdk-310p-npu-driver_24.1.rc3_linux-aarch64.run" ];then
        logger info "该版本驱动已经下载"
      else
        curl -C- -O --retry 3 "$NPU_310_DRIVER_URL" || { logger error "downloading npu driver failed"; exit 1; }
        /bin/mv -f "./Ascend-hdk-310p-npu-driver_24.1.rc3_linux-aarch64.run" "$BASE/down"
      fi
      ;;
      (910|910A|910B|910PRO)
      if [ -f "$BASE/down/Ascend-hdk-910b-npu-firmware_7.1.0.6.220.run" ];then
        logger info "该版本固件已经下载"
      else
        curl -C- -O --retry 3 "$NPU_910_firmware_URL" || { logger error "固件下载失败" ; exit 1; }
        /bin/mv -f "./Ascend-hdk-910b-npu-firmware_7.1.0.6.220.run" "$BASE/down"
      fi
      if [ -f "$BASE/down/Ascend-hdk-910b-npu-driver_24.1.rc1_linux-aarch64.run" ];then
        logger info "该版本驱动已经下载"
      else
        curl -C- -O --retry 3 "$NPU_910_DRIVER_URL" || { logger error "downloading npu driver failed"; exit 1; }
        /bin/mv -f "./Ascend-hdk-910b-npu-driver_24.1.rc1_linux-aarch64.run" "$BASE/down"
      fi
      ;;
  esac
  }
  NPUTYPES=$(grep -r NPU_TYPE  $BASE/ai.config  | awk -F"="  '!($4 in NPU){NPU[$4];print $4}')
  logger debug "开始下载NPU驱动"
  for nputype in $NPUTYPES
  do
    npu_driver_down $nputype
  done
  /bin/cp -f $BASE/down/Ascend-*  $BASE/ai/prepareinstall/roles/npudriverbefore/files/
}

function download_dcudriver() {
  DCU_URL=http://123.157.234.230:9000/ai-deployer/dcu/rock-5.7.1-6.2.30-V1.6.2.aio.run
  dcu_driver_down(){
    if [ -f "$BASE/down/rock-5.7.1-6.2.30-V1.6.2.aio.run" ];then
      logger info "DCU驱动已经下载"
    else
      curl -C- -O --retry 3 "$DCU_URL" || { logger error "downloading dcu driver failed"; exit 1; }
      /bin/mv -f "./rock-5.7.1-6.2.30-V1.6.2.aio.run" "$BASE/down"
    fi
  }
  
  logger debug "开始下载DCU驱动"
  dcu_driver_down
  /bin/cp -f $BASE/down/rock-*.run $BASE/ai/prepareinstall/roles/dcudriverbefore/files/
}

function install_dcudriver() {
  ansible-playbook -i $BASE/ai/prepareinstall/hosts $BASE/ai/prepareinstall/dcudriverbefore.yml
  ansible-playbook -i $BASE/ai/prepareinstall/hosts $BASE/ai/prepareinstall/dcudriver.yml
}


function install_gpudriver() {
  ansible-playbook -i $BASE/ai/prepareinstall/hosts $BASE/ai/prepareinstall/gpudriverbefore.yml
  ansible-playbook -i $BASE/ai/prepareinstall/hosts $BASE/ai/prepareinstall/gpudriver.yml

}
function install_npudriver() {
  ansible-playbook -i $BASE/ai/prepareinstall/hosts $BASE/ai/prepareinstall/npudriverbefore.yml
  ansible-playbook -i $BASE/ai/prepareinstall/hosts $BASE/ai/prepareinstall/npudriver.yml

}

function get_sys_pkg() {
  if [[ `cat /etc/os-release` =~ CentOS || `cat /etc/os-release` =~ openEuler || `cat /etc/os-release` =~ Rocky  ]];then
    SYS_PKG_URL="$MINIO_URL/ai-deployer/centos-rpm-packages.tar.gz"
    if [[ -f "$BASE/down/centos-rpm-packages.tar.gz" ]];then
      logger warn "system packages already existed!"
    else
      logger debug "正在下载centos离线包!"
      curl -C- -O --retry 3 "$SYS_PKG_URL" || { logger error "downloading system packages failed"; exit 1; }
      /bin/mv -f "./centos-rpm-packages.tar.gz" "$BASE/down"
      tar zxf "$BASE/down/centos-rpm-packages.tar.gz" -C "$BASE/down"
    fi
  elif [[ `cat /etc/os-release` =~ Ubuntu ]];then
    UBUNTU_VERSION=$(grep -E 'VERSION_ID' /etc/os-release | sed "s/'//g;s/\"//g" | cut -d'=' -f2)
    SYS_PKG_URL="$MINIO_URL/ai-deployer/ubuntu-$UBUNTU_VERSION-deb-packages.tar.gz"
    if [[ -f "$BASE/down/ubuntu-$UBUNTU_VERSION-deb-packages.tar.gz" ]];then
      logger warn "system packages already existed!"
    else
      logger debug "正在下载ubuntu离线包!"
      curl -C- -O --retry 3 "$SYS_PKG_URL" || { logger error "downloading system packages failed"; exit 1; }
      /bin/mv -f "./ubuntu-$UBUNTU_VERSION-deb-packages.tar.gz" "$BASE/down"
      tar zxf "$BASE/down/ubuntu-$UBUNTU_VERSION-deb-packages.tar.gz" -C "$BASE/down"
    fi
  fi
}

function install_ansible() {
  if [[ -f /usr/bin/ansible ]];then
    logger warn "ansible is already installed!"
  else
    #安装ansible
    logger debug "正在安装ansible!"
    if [[ `cat /etc/os-release` =~ CentOS ]];then
      tar zxf "$BASE/down/centos-rpm-packages/ansible-rpm-packages.tar.gz" -C "$BASE/down/centos-rpm-packages"
      rpm -ivh --force --nodeps $BASE/down/centos-rpm-packages/ansible-rpm-packages/*
    else
      #tar zxf "$BASE/down/ubuntu-deb-packages/ansible-deb-packages.tar.gz" -C "$BASE/down/ubuntu-deb-packages"
      #dpkg -i $BASE/down/ubuntu-deb-packages/ansible-deb-packages/*
      apt install ansible -y
      apt install python3-pymysql -y
      ansible-galaxy collection install community-mysql-3.7.2.tar.gz --ignore-errors
    fi
  fi
}

function configure_ssh() {
  if [[ -f /usr/bin/expect ]];then
    logger warn "expect is already installed!"
  else
    #安装expect
    if [[ `cat /etc/os-release` =~ CentOS ]];then
      tar zxf "$BASE/down/centos-rpm-packages/expect-rpm-packages.tar.gz" -C "$BASE/down/centos-rpm-packages"
      rpm -ivh --force --nodeps $BASE/down/centos-rpm-packages/expect-rpm-packages/*
    else
      #tar zxf "$BASE/down/ubuntu-deb-packages/expect-deb-packages.tar.gz" -C "$BASE/down/ubuntu-deb-packages"
      #dpkg -i $BASE/down/ubuntu-deb-packages/expect-deb-packages/*
      apt install expect -y
    fi
  fi
  
  logger debug "配置免密码登录!"
  #从配置文件获取服务器IP和密码
  serverIPs=$(grep -E 'MGR_NODE_IP|COMPUTE_NODE_IP|HARBOR_NODE_IP' $BASE/ai.config |cut -d'=' -f2|awk '!($1 in IP){IP[$1];print $1}')
  serverPassword=$(grep -r 'SERVER_PASSWORD' $BASE/ai.config)
  serverPassword=${serverPassword#*=}
  sshPort=$(grep -r 'SSH_PORT' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')

  #生成密钥
  if [[ ! -f /root/.ssh/id_rsa.pub ]];then
    ssh-keygen -t rsa -b 2048 -N '' -f ~/.ssh/id_rsa
  fi

  #将生成的密钥发送至目标服务器,自动输入yes和密码
  auto_ssh_copy_id() {
    expect -c "set timeout -1
    spawn ssh-copy-id -p $sshPort $1
    expect {
      *yes/no* {send -- yes\r;exp_continue;}
      *password:* {send -- $2\r;exp_continue;}
      eof        {exit 0;}
    }"
  }
  #遍历需要发送密钥的服务器
  ssh_copy_id_to_all(){
    for serverIP in $serverIPs
    do
      auto_ssh_copy_id $serverIP $serverPassword
    done
  }
  ssh_copy_id_to_all
}

function download_k8s_packages() {
  #从配置文件获取harbor IP和Port,来配置insecure-registries
  harborExternalDomain=$(grep -r 'HARBOR_EXTERNAL_DOMAIN' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  harborDomain=$(grep -r 'HARBOR_DOMAIN' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  harborPort=$(grep -r 'HARBOR_NODE_PORT' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  # sed -i "s/__registry__external__domain__/$harborExternalDomain:$harborPort/g" "$BASE/ezdown"
  # sed -i "s/__registry__local__domain__/$harborDomain:$harborPort/g" "$BASE/ezdown"
  # sed -i "s/__registry__domain__/$REGISTRY_DOMAIN/g" "$BASE/ezdown"
  #   # 获取系统架构
  ARCH=$(uname -m)
  
  # 选择合适的ezdown脚本并检查文件是否存在
  if [ "$ARCH" = "aarch64" ]; then
    EZDOWN_CMD="./ezdown-npu"
  else
    EZDOWN_CMD="./ezdown"
  fi
   chmod +x "$EZDOWN_CMD"
  
  # 替换配置
  sed -i "s/__registry__external__domain__/$harborExternalDomain:$harborPort/g" "$EZDOWN_CMD"
  sed -i "s/__registry__local__domain__/$harborDomain:$harborPort/g" "$EZDOWN_CMD"
  sed -i "s/__registry__domain__/$REGISTRY_DOMAIN/g" "$EZDOWN_CMD"
  
  # 执行命令
  $EZDOWN_CMD -D || { logger error "$EZDOWN_CMD -D 执行失败"; return 1; }
  $EZDOWN_CMD -R || { logger error "$EZDOWN_CMD -R 执行失败"; return 1; }
  $EZDOWN_CMD -P || { logger error "$EZDOWN_CMD -P 执行失败"; return 1; }
}

function new_cluster() {
  edition=$(grep -r 'EDITION' $BASE/ai.config |cut -d'=' -f2)
   ARCH=$(grep -v '^#' $BASE/ai.config | grep ARCH | cut -d"=" -f2)
  #如果集群clusters/ai不存在则创建并且是ARCH=x86
  if [[ ! -d clusters/ai && $ARCH == "x86" ]];then
    ./ezctl new ai
  fi
  #如果集群clusters/ai不存在则创建并且是ARCH=aarch64
  if [[ ! -d clusters/ai && $ARCH == "aarch64" ]];then
     chmod +x ./ezctl-npu
    ./ezctl-npu new ai 
  fi
  logger debug "正在配置集群信息"
  #从配置文件获取etcd节点IP
  etcdIPs=$(grep -r 'MGR_NODE_IP' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  for etcdIP in $etcdIPs
  do
    # 检查etcdIP正则表达式
    [[ $etcdIP =~ ^(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})(\.(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})){3}$ ]] || { logger error "Invalid ip etcd:$etcdIP"; return 1; }
    # 检查etcdIP是否已存在
    if [[ `sed -n '/^\[etcd/,/^\[kube_master/p' "$BASE/clusters/ai/hosts"|grep $etcdIP` ]];then
      logger warn "etcd $etcdIP already existed in $BASE/clusters/ai/hosts"
    else
      sed -i "/\[etcd/a $etcdIP" "$BASE/clusters/ai/hosts"
    fi
  done
  #从配置文件获取master节点IP
  masterIPs=$(grep -r 'MGR_NODE_IP' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  for masterIP in $masterIPs
  do
    # 检查masterIP正则表达式
    [[ $masterIP =~ ^(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})(\.(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})){3}$ ]] || { logger error "Invalid ip master:$masterIP"; return 1; }
    # 检查masterIP是否已存在
    if [[ `sed -n '/^\[kube_master/,/^\[kube_node/p' "$BASE/clusters/ai/hosts"|grep $masterIP` ]];then
      logger warn "master $masterIP already existed in $BASE/clusters/ai/hosts"
    else
      sed -i "/\[kube_master/a $masterIP" "$BASE/clusters/ai/hosts"
    fi
  done

  #从配置文件获取node节点IP
  nodeIPs=$(grep -r 'COMPUTE_NODE_IP' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  for nodeIP in $nodeIPs
  do
    # 检查nodeIP正则表达式
    [[ $nodeIP =~ ^(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})(\.(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})){3}$ ]] || { logger error "Invalid ip node:$nodeIP"; return 1; }
    # 检查nodeIP是否已存在
    if [[ `sed -n '/^\[kube_node/,/^\[harbor/p' "$BASE/clusters/ai/hosts"|grep $nodeIP` ]];then
      logger warn "node $nodeIP already existed in $BASE/clusters/ai/hosts"
    else
      sed -i "/\[kube_node/a $nodeIP" "$BASE/clusters/ai/hosts"
    fi
  done

  #从配置文件获取chrony节点IP
  chronyIP=$(grep -r 'MGR_NODE_IP' $BASE/ai.config |cut -d'=' -f2|awk 'NR==1{print $1}')
  # 检查chronyIP正则表达式
  [[ $chronyIP =~ ^(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})(\.(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})){3}$ ]] || { logger error "Invalid ip chrony:$chronyIP"; return 1; }
  # 检查chronyIP是否已存在
  if [[ `sed -n '/^\[chrony/,/^\[all:vars/p' "$BASE/clusters/ai/hosts"|grep $chronyIP` ]];then
    logger warn "chrony $chronyIP already existed in $BASE/clusters/ai/hosts"
  else
    sed -i "/\[chrony/a $chronyIP" "$BASE/clusters/ai/hosts"
  fi

  #从配置文件获取harbor IP 端口和admin用户密码
  harborNodeIP=$(grep -r 'HARBOR_NODE_IP' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  harborDomain=$(grep -r 'HARBOR_DOMAIN' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  harborExternalDomain=$(grep -r 'HARBOR_EXTERNAL_DOMAIN' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  harborPort=$(grep -r 'HARBOR_NODE_PORT' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  harborAdminPassword=$(grep -r 'HARBOR_ADMIN_PASSWORD' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  # 检查harborIP正则表达式
#  [[ $harborIP =~ ^(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})(\.(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})){3}$ ]] || { logger error "Invalid ip harbor:$harborIP"; return 1; }
  # 检查harborIP是否已存在
  if [[ `sed -n '/^\[harbor/,/^\[ex_lb/p' "$BASE/clusters/ai/hosts"|grep $harborNodeIP` ]];then
    logger warn "harbor $harborNodeIP already existed in $BASE/clusters/ai/hosts"
  else
    sed -i "s/__harbor_ip__/$harborNodeIP/g" "$BASE/clusters/ai/hosts"
    sed -i "s/__harbor_domain__/$harborDomain/g" "$BASE/clusters/ai/config.yml"
    sed -i "s/__harbor_external_domain__/$harborExternalDomain/g" "$BASE/clusters/ai/config.yml"
    sed -i "s/__harbor_port__/$harborPort/g" "$BASE/clusters/ai/config.yml"
    sed -i "s/__harbor_admin_password__/$harborAdminPassword/g" "$BASE/clusters/ai/config.yml"
  fi

  #从配置文件获取glusterfs/longhorn IP
  STORAGE_TYPE=$(grep -v '^#' $BASE/ai.config | grep 'STORAGE_TYPE' | cut -d'=' -f2)
  if [[ $STORAGE_TYPE == glusterfs || $STORAGE_TYPE == kadalu ]];then
    glusterfsIPs=$(grep -r 'STORAGE_NODE_IP' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
    for glusterfsIP in $glusterfsIPs
    do
      # 检查glusterfsIP正则表达式
      [[ $glusterfsIP =~ ^(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})(\.(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})){3}$ ]] || { logger error "Invalid ip glusterfs:$glusterfsIP"; return 1; }
      # 检查glusterfsIP是否已存在
      if [[ `sed -n '/^\[glusterfs/,/^\[longhorn/p' "$BASE/clusters/ai/hosts"|grep $glusterfsIP` ]];then
        logger warn "glusterfs $glusterfsIP already existed in $BASE/clusters/ai/hosts"
      else
#      if grep -q "RAW_DISK" "$BASE/ai.config"; then 修改脚本先去掉注释行,然后再看配置文件中是否有RAW_DISK
       if grep -v '^#' $BASE/ai.config | grep -q "RAW_DISK" ; then
          rawDisk=$(grep -r 'RAW_DISK' $BASE/ai.config |grep $glusterfsIP |cut -d'=' -f4)
          sed -i "/\[glusterfs/a $glusterfsIP RAW_DISK=$rawDisk" "$BASE/clusters/ai/hosts"
        else
          localPath=$(grep -r 'LOCAL_STORAGE_PATH' $BASE/ai.config| grep -v '^#' |grep $glusterfsIP |cut -d'=' -f4)
          sed -i "/\[glusterfs/a $glusterfsIP LOCAL_STORAGE_PATH=$localPath" "$BASE/clusters/ai/hosts"
        fi
      fi
    done
  else
    longhornIPs=$(grep -r 'STORAGE_NODE_IP' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
    for longhornIP in $longhornIPs
    do
      # 检查longhornIP正则表达式
      [[ $longhornIP =~ ^(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})(\.(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})){3}$ ]] || { logger error "Invalid ip longhorn:$longhornIP"; return 1; }
      # 检查longhornIP是否已存在
      if [[ `sed -n '/^\[longhorn/,/^\[chrony/p' "$BASE/clusters/ai/hosts"|grep $longhornIP` ]];then
        logger warn "longhorn $longhornIP already existed in $BASE/clusters/ai/hosts"
      else
        rawDisk=$(grep -r 'STORAGE_NODE_IP' $BASE/ai.config |grep $longhornIP |cut -d'=' -f4)
        sed -i "/\[longhorn/a $longhornIP RAW_DISK=$rawDisk" "$BASE/clusters/ai/hosts"
      fi
    done
  fi
     
  
  #从配置文件获取EFK IP
  if [[ $edition == professional ]];then
    efkIP=$(grep -r 'EFK_NODE_IP' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
    # 检查efkIP正则表达式
    [[ $efkIP =~ ^(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})(\.(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})){3}$ ]] || { logger error "Invalid ip efk:$efkIP"; return 1; }
    sed -i "s/__efk_ip__/$efkIP/g" "$BASE/clusters/ai/config.yml"
    sed -i 's/efk_install: "no"/efk_install: "yes"/g' "$BASE/clusters/ai/config.yml"
  fi

  #从配置文件获取PROMETHEUS IP
  prometheusIP=$(grep -r 'PROMETHEUS_IP' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  # 检查prometheusIP正则表达式
  [[ $prometheusIP =~ ^(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})(\.(2(5[0-5]{1}|[0-4][0-9]{1})|[0-1]?[0-9]{1,2})){3}$ ]] || { logger error "Invalid ip prometheus:$prometheusIP"; return 1; }
  sed -i "s/__prometheus_ip__/$prometheusIP/g" "$BASE/clusters/ai/config.yml"
  
  #从配置文件获取服务器ssh端口
  sshPort=$(grep -r 'SSH_PORT' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  sed -i "s/__ansible_ssh_port__/$sshPort/g" "$BASE/clusters/ai/hosts"

  #从配置文件获取nodeport range
  nodeportRange=$(grep -r 'NODE_PORT_RANGE' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  sed -i "s/__node_port_range__/$nodeportRange/g" "$BASE/clusters/ai/hosts"

  #从配置文件获取cluster cidr
  clusterCidr=$(grep -r 'CLUSTER_CIDR' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  sed -i "s#__cluster_cidr__#$clusterCidr#g" "$BASE/clusters/ai/hosts"

  #从配置文件获取harbor网络
  harborNetwork=$(grep -r 'HARBOR_NETWORK' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  harborClairNetwork=$(grep -r 'HARBOR_CLAIR_NETWORK' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  harborChartmuseumNetwork=$(grep -r 'HARBOR_CHARTMUSEUM_NETWORK' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  sed -i "s#__harbor_network__#$harborNetwork#g" "$BASE/clusters/ai/config.yml"
  sed -i "s#__harbor_clair_network__#$harborClairNetwork#g" "$BASE/clusters/ai/config.yml"
  sed -i "s#__harbor_chartmuseum_network__#$harborChartmuseumNetwork#g" "$BASE/clusters/ai/config.yml"
}

function setup_cluster() {
  ./ezctl setup ai all
}

function install_nvidia_runtime() {
  logger debug "安装nvidia-container-runtime"
  #在所有gpu主机上安装和配置nvidia-container-runtime
  cat $BASE/ai.config | grep -E 'GPU_TYPE'|awk -F " " '{print $1}'|awk -F "=" '!($2 in IP){IP[$2];print $2}' > /tmp/gpu_nodes.tmp
  sshPort=$(grep -r 'SSH_PORT' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  echo "[all:vars]" >> /tmp/gpu_nodes.tmp
  sed -i "/\[all:vars/a ansible_ssh_port=$sshPort" /tmp/gpu_nodes.tmp
  /bin/cp -f $BASE/ai/ai-manifests/02-nvidia-container-runtime.sh.tmp $BASE/ai/ai-manifests/02-nvidia-container-runtime.sh
  sed -i "s#MINIO_URL#$MINIO_URL#g" $BASE/ai/ai-manifests/02-nvidia-container-runtime.sh

  ansible all -i /tmp/gpu_nodes.tmp -m copy -a "src=$BASE/ai/ai-manifests/02-nvidia-container-runtime.sh dest=/opt/02-nvidia-container-runtime.sh"
  ansible all -i /tmp/gpu_nodes.tmp -m shell -a "cd /opt && /bin/bash /opt/02-nvidia-container-runtime.sh"
  #ansible all -i /tmp/gpu_nodes.tmp -m shell -a "/usr/bin/nvidia-container-runtime -v"
}

function install_ascend_docker_runtime() {
  logger debug "安装ascend-docker-runtime"
  #在所有NPU主机上安装和配置ascend-docker-runtime
  cat $BASE/ai.config | grep -v '^#' |grep -E 'NPU_TYPE'|awk -F " " '{print $1}'|awk -F "=" '!($2 in IP){IP[$2];print $2}' > /tmp/npu_nodes.tmp
  sshPort=$(grep -r 'SSH_PORT' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
  echo "[all:vars]" >> /tmp/npu_nodes.tmp
  sed -i "/\[all:vars/a ansible_ssh_port=$sshPort" /tmp/npu_nodes.tmp
  /bin/cp -f $BASE/ai/ai-manifests/ascend-docker-runtime.sh.tmp $BASE/ai/ai-manifests/ascend-docker-runtime.sh

  ansible all -i /tmp/npu_nodes.tmp -m copy -a "src=$BASE/ai/ai-manifests/ascend-docker-runtime.sh dest=/opt/ascend-docker-runtime.sh"
  ansible all -i /tmp/npu_nodes.tmp -m shell -a "cd /opt && /bin/bash /opt/ascend-docker-runtime.sh"
}

function setup_harbor() {
  #判断docker是否已经安装
  dockerState=$(ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "systemctl status docker|grep Active || echo 'NoFound'" |awk 'NR>1')
  if [[ $dockerState =~ NoFound ]];then
    logger info "正在安装docker"
    ansible-playbook -i "clusters/ai/hosts" -e "@clusters/ai/config.yml" "playbooks/08.harbor-docker.yml"
  fi
  #判断harbor是否已经安装
  harborState=$(ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "ls /data/harbor" |awk 'NR>1')
  if [[ $harborState =~ registry ]];then
    logger info "harbor已经安装"
    isActive=$(ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "docker ps -a|grep goharbor/nginx-photon"|awk 'NR>1')
    if [[ $isActive =~ healthy ]];then
      logger info "harbor正在运行中"
    else
      ARCH=$(grep -r "ARCH" $BASE/ai.config|cut -d"=" -f2)
      if [[ "$ARCH" == "aarch64" ]];then
              ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "/usr/bin/docker load < /data/harbor/harbor.v1.9.3.tar.gz"
              ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "/usr/bin/docker load < /data/harbor/harbor-db.tar"
              ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "cd /data/harbor && /usr/bin/docker-compose up -d"
      else
             ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "/usr/bin/docker load < /data/harbor/harbor.v1.10.16.tar.gz"
             ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "/usr/bin/docker load < /data/harbor/harbor-core.tar"
             ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "/usr/bin/docker load < /data/harbor/harbor-portal.tar"
             ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "cd /data/harbor && /usr/bin/docker-compose up -d"
      fi
      until [[ $isActive =~ healthy ]]
      do
        isActive=$(ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "docker ps -a|grep goharbor/nginx-photon"|awk 'NR>1')
        if [[ $isActive =~ healthy ]];then
          logger info "harbor正在运行中"
        fi
      done
    fi
  else
    ARCH=$(grep -v '^#' $BASE/ai.config | grep ARCH | cut -d"=" -f2)
    if [[ "$ARCH" == "aarch64" ]];then
    chmod +x  ezctl-npu
    ./ezctl-npu setup ai harbor
    else
    ./ezctl setup ai harbor
    fi
  fi
}

function get_images() {
  #从线上镜像仓库拉取部署所需的镜像
  if [[ -f /mnt/got-images ]];then
    logger info "镜像已经拉取到本地仓库"
  else
    logger debug "正在拉取镜像至本地仓库"
    > $BASE/images-tag
    nohup ./get-images.sh & >/dev/null 2>&1
    ./push-images.sh
 
    touch /mnt/got-images
  fi
  
}

function get_frame_images() {
  #安装完nvidia_runtime需要重启docker,重启docker后等待harbor运行
  isActive=$(ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "docker ps -a|grep goharbor/nginx-photon"|awk 'NR>1')
  if [[ $isActive =~ healthy ]];then
    logger info "harbor正在运行中"
  else
    logger debug "正在重启harbor"
    ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "cd /data/harbor && /usr/bin/docker-compose up -d"
    until [[ $isActive =~ healthy ]]
    do
      isActive=$(ansible -i $BASE/clusters/ai/hosts harbor -m shell -a "docker ps -a|grep goharbor/nginx-photon"|awk 'NR>1')
      if [[ $isActive =~ healthy ]];then
        logger info "harbor正在运行中"
      fi
    done
  fi

  if [[ -f /mnt/got-frame-images ]];then
    logger info "框架镜像已经拉取到本地仓库"
  else
    logger debug "正在拉取框架镜像至本地仓库"
    cd $BASE/ai/ai-manifests
    ARCH=$(grep -v '^#' $BASE/ai.config | grep ARCH | cut -d"=" -f2)
    if [[ $ARCH == "aarch64" ]];then
        chmod +x  get-npu-frame-images.sh
        nohup ./get-npu-frame-images.sh & >/dev/null 2>&1
    else
        nohup ./get-frame-images.sh & >/dev/null 2>&1
    fi
    cd $BASE
  fi
}

function setup_storage() {
  if [[ ! -f /usr/bin/python ]];then
    ln -s /usr/bin/python3 /usr/bin/python
  fi

  STORAGE_TYPE=$(grep -v '^#' $BASE/ai.config | grep 'STORAGE_TYPE' | cut -d'=' -f2)
  if [[ $STORAGE_TYPE == glusterfs ]];then
    namespaces=$($BASE/bin/kubectl get ns|awk 'NR>1{print $1}')
    if [[ $namespaces =~ glusterfs ]];then
      logger info "glusterfs服务已经安装"
    else
      ./ezctl setup ai glusterfs
    fi
  else
    namespaces=$($BASE/bin/kubectl get ns|awk 'NR>1{print $1}')
    if [[ $namespaces =~ longhorn-system ]];then
      logger info "longhorn服务已经安装"
    else
      ./ezctl setup ai longhorn
    fi
  fi
}

function setup_storage_new() {
  STORAGE_TYPE=$(grep -v '^#' $BASE/ai.config | grep 'STORAGE_TYPE' | cut -d'=' -f2)
  if [[ $STORAGE_TYPE == glusterfs ]];then
    if [[ `cat /etc/os-release` =~ CentOS ]];then
      GLUSTERFS_PKG_URL="$MINIO_URL/ai-deployer/glusterfs8-6/gfs86-server-centos.tar"
      if [[ -f "$BASE/down/gfs86-server-centos.tar" ]];then
        logger info "glusterfs8.6 离线包已下载"
      else
        logger debug "正在下载glusterfs8.6 离线包"
        curl -C- -O --retry 3 "$GLUSTERFS_PKG_URL" || { logger error "下载glusterfs 离线包失败"; exit 1; }
        /bin/mv -f "./gfs86-server-centos.tar" "$BASE/down"
      fi
    else
      UBUNTU_VERSION=$(grep -E 'VERSION_ID' /etc/os-release | sed "s/'//g;s/\"//g" | cut -d'=' -f2)
      GLUSTERFS_PKG_URL="$MINIO_URL/ai-deployer/glusterfs8-6/glusterfs-86-ubuntu-$UBUNTU_VERSION.tar"
      if [[ -f "$BASE/down/glusterfs-86-ubuntu.tar" ]];then
        logger info "glusterfs8.6 离线包已下载"
      else
        logger debug "正在下载glusterfs8.6 离线包"
        curl -C- -O --retry 3 "$GLUSTERFS_PKG_URL" || { logger error "下载glusterfs 离线包失败"; exit 1; }
        /bin/mv -f "./glusterfs-86-ubuntu-$UBUNTU_VERSION.tar" "$BASE/down/glusterfs-86-ubuntu.tar"
      fi
    fi
  
    #下载heketi离线包
    if [[ -f "$BASE/down/heketi-v9.0.0.linux.amd64.tar.gz" ]];then
      logger info "heketi 离线包已下载"
    else
      HEKETI_PGK_URL="$MINIO_URL/ai-deployer/glusterfs8-6/heketi-v9.0.0.linux.amd64.tar.gz"
      logger debug "正在下载heketi离线包"
      curl -C- -O --retry 3 "$HEKETI_PGK_URL" || { logger error "下载heketi 离线包失败"; exit 1; }
      /bin/mv -f "./heketi-v9.0.0.linux.amd64.tar.gz" "$BASE/down"
      tar zxf "$BASE/down/heketi-v9.0.0.linux.amd64.tar.gz" -C "$BASE/down"
    fi
  fi

  if [[ $STORAGE_TYPE == kadalu ]];then
    if [[ `cat /etc/os-release` =~ CentOS ]];then
      GLUSTERFS_PKG_URL="$MINIO_URL/ai-deployer/glusterfs8-6/gfs86-server-centos.tar"
      if [[ -f "$BASE/down/gfs86-server-centos.tar" ]];then
        logger info "glusterfs8.6 离线包已下载"
      else
        logger debug "正在下载glusterfs8.6 离线包"
        curl -C- -O --retry 3 "$GLUSTERFS_PKG_URL" || { logger error "下载glusterfs 离线包失败"; exit 1; }
        /bin/mv -f "./gfs86-server-centos.tar" "$BASE/down"
      fi
    else
      UBUNTU_VERSION=$(grep -E 'VERSION_ID' /etc/os-release | sed "s/'//g;s/\"//g" | cut -d'=' -f2)
      GLUSTERFS_PKG_URL="$MINIO_URL/ai-deployer/glusterfs8-6/glusterfs-86-ubuntu-$UBUNTU_VERSION.tar"
      if [[ -f "$BASE/down/glusterfs-86-ubuntu.tar" ]];then
        logger info "glusterfs8.6 离线包已下载"
      else
        logger debug "正在下载glusterfs8.6 离线包"
        curl -C- -O --retry 3 "$GLUSTERFS_PKG_URL" || { logger error "下载glusterfs 离线包失败"; exit 1; }
        /bin/mv -f "./glusterfs-86-ubuntu-$UBUNTU_VERSION.tar" "$BASE/down/glusterfs-86-ubuntu.tar"
      fi
    fi
  fi

  if [[ $STORAGE_TYPE == glusterfs ]];then
    storageclasses=$($BASE/bin/kubectl get storageclass|awk 'NR>1{print $1}')
    if [[ $storageclasses =~ sc-glusterfs ]];then
      logger info "glusterfs服务已安装"
    else
      logger debug "正在安装glusterfs"
      ansible-playbook -i "clusters/ai/hosts" -e "@clusters/ai/config.yml" "playbooks/14.glusterfs-8.6.yml"
    fi


  elif [[ $STORAGE_TYPE == kadalu ]];then
    storageclasses=$($BASE/bin/kubectl get storageclass|awk 'NR>1{print $1}')
    if [[ $storageclasses =~ kadalu.ext ]];then
      logger info "kadalu服务已安装" 
    else
      logger debug "正在安装kadalu"
      ansible-playbook -i "clusters/ai/hosts" -e "@clusters/ai/config.yml" "playbooks/43.kadalu.yml"
    fi

  elif [[ $STORAGE_TYPE == localpath ]];then
    storageclasses=$($BASE/bin/kubectl get storageclass|awk 'NR>1{print $1}')
    if [[ $storageclasses =~ local-path ]];then
      logger info "localpath storageclass服务已安装" 
    else
      logger debug "正在安装localpath storageclass"
      LOCALPATH=$(grep -r 'LOCALPATH' $BASE/ai.config |grep -v '^#'|cut -d'=' -f2)
      sed -i "s#LOCALPATH#$LOCALPATH#g" /opt/ai-deployer/ai/ai-manifests/yaml/local-path-storage.yaml
      sed -i "s#registry.mopaas.com#registry.local:38080#" /opt/ai-deployer/ai/ai-manifests/yaml/local-path-storage.yaml
      $BASE/bin/kubectl apply -f /opt/ai-deployer/ai/ai-manifests/yaml/local-path-storage.yaml
    fi


  else
    storageclasses=$($BASE/bin/kubectl get storageclass|awk 'NR>1{print $1}')
    if [[ $storageclasses =~ longhorn ]];then
      logger info "longhorn服务已经安装"
    else
      logger debug "正在安装longhorn"
      ./ezctl setup ai longhorn
    fi
  fi
}
function setup_storage_npu() {
  if [[ ! -f /usr/bin/python ]];then
    ln -s /usr/bin/python3 /usr/bin/python
  fi

  STORAGE_TYPE=$(grep -v '^#' $BASE/ai.config | grep 'STORAGE_TYPE' | cut -d'=' -f2)
  if [[ $STORAGE_TYPE == glusterfs ]];then
    namespaces=$($BASE/bin/kubectl get ns|awk 'NR>1{print $1}')
    if [[ $namespaces =~ glusterfs ]];then
      logger info "glusterfs服务已经安装"
    else
      ./ezctl-npu i-glusterfs ai
    fi
  else
    namespaces=$($BASE/bin/kubectl get ns|awk 'NR>1{print $1}')
    if [[ $namespaces =~ longhorn-system ]];then
      logger info "longhorn服务已经安装"
    else
      ./ezctl-npu setup ai longhorn
    fi
  fi
}

function setup_ai() {
   chmod +x ./ai/ai-manifests/ai-install.sh
  ./ai/ai-manifests/ai-install.sh all
}
function setup_ai_npu() {
  chmod +x ./ai/ai-manifests/ai-npu-install.sh
  ./ai/ai-manifests/ai-npu-install.sh all
}
function setup_ai_web() {
  cd $BASE/ai/ai-manifests/ai-auto-deploy/
  ./install.sh
  logger info "web应用部署成功"
  cd $BASE
}
function setup_ai_web_npu() {
  cd $BASE/ai/ai-manifests/ai-auto-deploy-arm64/
  logger info "进入arm64"
  chmod +x install.sh
  ./install.sh
  logger info "web应用部署成功"
  cd $BASE
}
### Main Lines ##################################################
function main() {
  BASE=$(grep -E 'BASE' ./ai.config|cut -d'=' -f2)
  mkdir -p "$BASE/bin" "$BASE/down"
  DEFAULT_BASE=/opt/ai-deployer
  if [[ $BASE != $DEFAULT_BASE ]];then
    logger debug "正在配置工作目录"
    sed -i "s#$DEFAULT_BASE#$BASE#g" `grep -r $DEFAULT_BASE . -rl`
  fi
  #线上MINIO
  MINIO_URL=$(grep -E 'MINIO_URL' $BASE/ai.config|cut -d'=' -f2)
  
  #线上镜像仓库
  REGISTRY_URL=$(grep -E 'REGISTRY_URL' $BASE/ai.config|cut -d'=' -f2)
  REGISTRY_DOMAIN=${REGISTRY_URL#*//}
  #本地镜像仓库
  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"

  DEFAULT_REGISTRY_DOMAIN=registry.local:38080
  if [[ $REGISTRY_LOCAL != $DEFAULT_REGISTRY_DOMAIN ]];then
    logger debug "正在配置线上镜像仓库"
    sed -i "s#$DEFAULT_REGISTRY_DOMAIN#$REGISTRY_LOCAL#g" $BASE/install.sh
    sed -i "s#$DEFAULT_REGISTRY_DOMAIN#$REGISTRY_LOCAL#g" $BASE/operate.sh
    sed -i "s#$DEFAULT_REGISTRY_DOMAIN#$REGISTRY_LOCAL#g" $BASE/clean.sh
    sed -i "s#$DEFAULT_REGISTRY_DOMAIN#$REGISTRY_LOCAL#g" `grep -r $DEFAULT_REGISTRY_DOMAIN $BASE/ai/ai-manifests/kubeflow-manifests -rl`
    sed -i "s#$DEFAULT_REGISTRY_DOMAIN#$REGISTRY_LOCAL#g" `grep -r $DEFAULT_REGISTRY_DOMAIN $BASE/ai/ai-manifests/yaml/kadalu -rl`
  fi
  [ "$#" -gt 0 ] || { usage >&2; exit 2; }

  case "$1" in
      (all)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          # 获取ARCH值,排除注释行
          ARCH=$(grep -v '^#' $BASE/ai.config | grep ARCH | cut -d"=" -f2)
          
          # 检查配置文件中的有效配置(排除注释)
          HAS_GPU=$(grep -v '^#' $BASE/ai.config | grep 'GPU_TYPE' || echo "")
          HAS_DCU=$(grep -v '^#' $BASE/ai.config | grep 'DCU_TYPE' || echo "")
          HAS_NPU=$(grep -v '^#' $BASE/ai.config | grep 'NPU_TYPE' || echo "")
          if [[ "$ARCH" == x86 ]];then
              if [[ ! -z "$HAS_GPU" ]]; then
                  download_gpudriver
                  install_gpudriver
              elif [[ ! -z "$HAS_DCU" ]]; then
                  download_dcudriver  
                  install_dcudriver
              fi
              get_sys_pkg
              install_ansible
              configure_ssh
              download_k8s_packages
              new_cluster
              setup_harbor
              get_images
              setup_cluster
              if [[ ! -z "$HAS_GPU" ]]; then    
                install_nvidia_runtime
              fi
              get_frame_images
              #setup_storage
              setup_storage_new
              setup_ai
              setup_ai_web
          elif [[ "$ARCH" == aarch64 ]];then
              if [[ ! -z "$HAS_NPU" ]]; then
                 download_npudriver
                 install_npudriver
              fi
              get_sys_pkg
              install_ansible
              configure_ssh
              download_k8s_packages
              new_cluster
              setup_harbor
              get_images
              setup_cluster
              install_ascend_docker_runtime
              get_frame_images
              setup_storage_npu
              setup_ai_npu
              setup_ai_web_npu
          fi
          ;;
      (k8s)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          get_sys_pkg
          install_ansible
          configure_ssh
          download_k8s_packages
          new_cluster
          setup_harbor
          get_images
          setup_cluster
          ;;
      (gpudriver)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          download_gpudriver
          install_gpudriver
          ;;
      (harbor)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          setup_harbor
          ;;
      (storage)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          #setup_storage 如何arch的值是x86 则执行
          ARCH=$(grep -v '^#' $BASE/ai.config | grep ARCH | cut -d"=" -f2)
          if [[ "$ARCH" == x86 ]];then
            setup_storage_new
          else
            setup_storage_npu
          fi 
          ;;
      (nvidia-runtime)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          install_nvidia_runtime
          ;;
      (ai)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          ARCH=$(grep -v '^#' $BASE/ai.config | grep ARCH | cut -d"=" -f2)
          if [[ "$ARCH" == x86 ]];then
             setup_ai
          else
             setup_ai_npu
          fi
          ;;
      (ai-web)
          [ "$#" -eq 1 ] || { usage >&2; exit 2; }
          ARCH=$(grep -v '^#' $BASE/ai.config | grep ARCH | cut -d"=" -f2)
  
          if [[ "$ARCH" == x86 ]];then
              setup_ai_web
          else
              setup_ai_web_npu
          fi
          ;;
      (*)
          usage
          exit 0
          ;;
  esac
}
main "$@"