#!/bin/sh
# 已经测试过安装ok的系统
# ubuntu: 16.04/18.04/20.04/22.04
# centos: 7
# debian: 11
# mac: amd架构/arm架构
# windows: 10/11
#
# 已经测试过支持的SHELL类型
# bash/zsh/oh-my-zsh(同zsh)

# 下载压缩包的地址
version=${LEPTON_CLI_VERSION:-latest}
# 默认使用cn环境的
repo="https://scc.aoss.cn-sh-01.sensecoreapi-oss.cn/$version"
# 根据SCC_ENV环境变量确定
ENV_VALUE=${SCC_ENV}
if [ -z "$ENV_VALUE" ]; then
  # 如果SCC_ENV环境变量不存在，则使用AILAB_ENV环境变量
  ENV_VALUE=${AILAB_ENV}
fi
if [ -n "$ENV_VALUE" ]; then
  if [ "$ENV_VALUE" = "dev" ]; then
    repo="https://scc.aoss.cn-sh-01.sensecoreapi-oss.dev/$version"
  elif [ "$ENV_VALUE" = "tech" ]; then
    repo="https://scc.aoss.cn-sh-01.sensecoreapi-oss.tech/$version"
  elif [ "$ENV_VALUE" = "st-sh-01" ]; then
    repo="https://scc.aoss.st-sh-01.sensecoreapi-oss.cn/$version"
  elif [ "$ENV_VALUE" = "cn-fj-01" ]; then
    repo="https://scc.aoss.cn-fj-01.thinkheadbdapi-oss.com/$version"
  elif [ "$ENV_VALUE" = "ms-sc-01" ]; then
    repo="https://scc.aoss.ms-sc-01.maoshanwangtechapi-oss.com/$version"
  elif [ "$ENV_VALUE" = "cn-fz-01" ]; then
    repo="https://scc.aoss.cn-fz-01.fjscmsapi-oss.com/$version"
  elif [ "$ENV_VALUE" = "cn-tj-01" ]; then
    repo="https://scc.aoss.cn-tj-01.sensecoreapi-oss.cn/$version"
  elif [ "$ENV_VALUE" = "cn-sz-02" ]; then
    repo="https://scc.aoss.cn-sz-02.sensecoreapi-oss.cn/$version"
  else
    repo="https://scc.aoss.cn-sh-01.sensecoreapi-oss.cn/$version"
  fi
fi
# 根据执行脚本时传入的参数变量确定
if [ $# -gt 0 ]; then
  if [ "$1" = "st-sh-01" ]; then
      repo="https://scc.aoss.st-sh-01.sensecoreapi-oss.cn/$version"
  elif [ "$1" = "cn-fj-01" ]; then
      repo="https://scc.aoss.cn-fj-01.thinkheadbdapi-oss.com/$version"
  elif [ "$1" = "ms-sc-01" ]; then
      repo="https://scc.aoss.ms-sc-01.maoshanwangtechapi-oss.com/$version"
  elif [ "$1" = "cn-fz-01" ]; then
      repo="https://scc.aoss.cn-fz-01.fjscmsapi-oss.com/$version"
  else
      repo="https://scc.aoss.cn-sh-01.sensecoreapi-oss.$1/$version"
  fi
fi

echo "Install slurm compatiable CLI begins"

# 获取当前系统
case $(uname -s) in
    Linux) os=linux ;;
    Darwin) os=darwin ;;
    MINGW64*) os=windows ;; # windows还有其他类型,这里为常见的win7/10等
    *) os= ;;
esac
# windows下安装cli的几种方式
# 1。通过exe/msi后缀的安装文件。
# 2。借助包管理工具，例如：winget/scoop/choco。
# 3。借助PowerShell的脚本文件（sp1后缀的文件）。
# 4。通过Bash运行安装脚本，例如：使用git bash跑shell脚本。wsl和cygwin这种通过在windows安装linux虚拟机的方式不可行，就算安装上了在虚拟机外也无法使用。
# 5。手动下载压缩包解压后配置PATH。
# 6。通过编程语言提供的包管理工具进行安装，例如：Python的pip，Golang的install。（本方式需要代码开源）

# 获取当前操作系统的架构
case $(uname -m) in
    amd64|x86_64) arch=amd64 ;;
    arm64|aarch64) arch=arm64 ;;
    *) arch= ;;
esac

# 检查系统
if [ -z "$os" ]; then
    echo "OS $(uname -s) not supported. \n\tonly support windows/Linux/Mac" >&2
    exit 1
fi

if [ -z "$arch" ]; then
    echo "Architecture  $(uname -m) not supported. \n\tonly support amd64" >&2
    exit 1
fi

# 获取shell类型 注意：如果是docker的ubuntu容器运行sh install.sh 这个时候$SHELL为空，但是source运行就正常
shell=$(echo $SHELL | awk 'BEGIN {FS="/";} { print $NF }')
if [ -z "$shell" ]; then
    echo "try another way to run the script. \ne.g.\n\tsource install.sh" >&2
    exit 1
fi

if [ -z "$SLURM_COMPATIABLE_CLI_HOME" ]; then
    SLURM_COMPATIABLE_CLI_HOME=$HOME/.scc
fi

# 文本加粗
bold=$(tput bold 2>/dev/null)
# 恢复默认
sgr0=$(tput sgr0 2>/dev/null)
# 执行文件路径
bin_dir=$SLURM_COMPATIABLE_CLI_HOME/bin
# 这里应该直接下对应版本的压缩文件
download_url="$repo/lepton_cli_${os}_${arch}.tar.gz"
download_check_upgrade="$repo/check_upgrade.sh"
tmp_file="/tmp/lepton_cli_${os}_${arch}.tar.gz"
check_upgrade_file="$SLURM_COMPATIABLE_CLI_HOME/check_upgrade.sh"

# 下载编译好的压缩文件 TODO：这里的文件名待确认，并且解压的方式也需要考虑到压缩包的目录组织形式
download_binary() {
    echo "start download $download_url"
    curl -# $download_url -o $tmp_file || return 1
    echo "start download $download_check_upgrade"
    curl -# $download_check_upgrade -o $check_upgrade_file || return 1
    return 0
}

# 解压缩
extract_binary(){
    echo "start extract $tmp_file"
    tar -zxf $tmp_file -C "$SLURM_COMPATIABLE_CLI_HOME" --strip-components 1 --overwrite || return 1
    rm $tmp_file
    return 0
}

check_depends() {
    pass=0
    command -v curl >/dev/null || {
        echo "Dependency check failed: please install 'curl' before proceeding."
        pass=1
    }
    command -v tar >/dev/null || {
        echo "Dependency check failed: please install 'tar' before proceeding."
        pass=1
    }
    command -v gzip >/dev/null || {
        echo "Dependency check failed: please install 'gzip' before proceeding."
        pass=1
    }

    return $pass
}

# install_manual() {
#     echo "start install manual"
#     mkdir -p /usr/local/man/man1
#     for file in "$SLURM_COMPATIABLE_CLI_HOME"/man/*
#     do
#         gzip "$file" > /dev/null
#     done
#     cp -rf "$SLURM_COMPATIABLE_CLI_HOME"/man/* /usr/local/man/man1
#     mandb > /dev/null
#     rm -rf "$SLURM_COMPATIABLE_CLI_HOME"/man
# }

if ! check_depends; then
    exit 1
fi

mkdir -p "$SLURM_COMPATIABLE_CLI_HOME"

if ! download_binary; then
    echo "Failed to download lepton-cli archive."
    exit 1
fi

if ! extract_binary; then
    echo "Failed to extract lepton-cli archive."
    exit 1
fi

# if ! install_manual; then
#     echo "Failed to install CLI manual."
#     exit 1
# fi

echo "Detected shell: ${bold}$shell${sgr0}"

# linux和mac相同的环境配置
if [ "$os" = linux ] || [ "$os" = darwin ]; then
    for file in "$bin_dir"
    do 
        chmod 755 $file
    done
   
    # if [ -f "${HOME}/.${shell}_profile" ]; then
    #     PROFILE=${HOME}/.${shell}_profile
    # el
    if [ -f "${HOME}/.${shell}_login" ]; then
        PROFILE=${HOME}/.${shell}_login
    elif [ -f "${HOME}/.${shell}rc" ]; then
        PROFILE=${HOME}/.${shell}rc
    else
        PROFILE=${HOME}/.profile
    fi

    echo "Shell profile:  ${bold}$PROFILE${sgr0}"

    sed -i '/scc_upgrade/d' ${PROFILE} 2>/dev/null

    case :$PATH: in
        *:$bin_dir:*)
            echo "PATH already contains $bin_dir"
            echo "open a new terminal or ${bold}source ${PROFILE}${sgr0} to use it";;
        *)
            printf '\nexport PATH=%s:$PATH\n' "$bin_dir" >> "$PROFILE"
            echo "$PROFILE has been modified to add lepton-cli to PATH"
            echo "open a new terminal or ${bold}source ${PROFILE}${sgr0} to use it"
            ;;
    esac
    
    # 确保scc_upgrade在export 环境变量后
    echo "scc_upgrade 2>/dev/null" >> "$PROFILE"
    #已在开发机部署时加入
    #grep -q -c "bash $check_upgrade_file" $PROFILE || printf "bash "$check_upgrade_file"" >> "$PROFILE"
fi

# 这里cli相关配置的初始化
if [ -f ~/.scc/config.yaml ]
then
    echo "Config file already exists"
    rm $SLURM_COMPATIABLE_CLI_HOME/config.yaml.default
else
    mv $SLURM_COMPATIABLE_CLI_HOME/config.yaml.default $SLURM_COMPATIABLE_CLI_HOME/config.yaml
fi
echo "Installed successfully, use srun/sinfo/scancel/squeue/scontrol command to show more info."
echo "Lepton CLI $version"
echo "SCC v1.2.x supports ACP (Classic), while v2.x.x supports ACP"
echo "Please note that SCC between these two versions does not support automatic upgrades"
echo "Kindly ensure that you use the corresponding SCC version for the respective ACP"