#!/bin/bash
set -e
VERSION="$1"
WORK_DIR=$(mktemp -d /tmp/rdma_agent_XXXXXX)
TARGET_PARENT="/etc/sysop/rdma_service/"
LOG_FILE="/var/log/rdma-agent-install.log"
TARGET_NAME="agent_tobhpc"
DOWNLOAD_RC=1
INSTALL_RC=1
START_SCRIPT="$TARGET_PARENT$TARGET_NAME/start-rdma-agent-service.sh"
PKG_TYPE=""
PKG_DEB="deb"
PKG_RPM="rpm"
PKG_UNKNOWN="unknown"

# 记录日志函数
function log() {
  log_text="$(date '+%Y-%m-%d %H:%M:%S') - $1"
  echo "$log_text"
  echo "$log_text" >> "$LOG_FILE"
}

get_pkg_type() {
  # 方案一：检查 /etc/os-release
  if [ -f /etc/os-release ]; then
    # 在当前 shell 环境中加载变量
    . /etc/os-release
    case "$ID $ID_LIKE" in
      *debian*|*ubuntu*|*mint*)
        printf '%s' "$PKG_DEB"
        return
        ;;
      *rhel*|*centos*|*fedora*|*suse*|*rocky*|*almalinux*)
        printf '%s' "$PKG_RPM"
        return
        ;;
    esac
  fi

  # 方案二：检查特定发行版的文件
  if [ -f /etc/debian_version ]; then
    printf '%s' "$PKG_DEB"
    return
  elif [ -f /etc/redhat-release ]; then
    printf '%s' "$PKG_RPM"
    return
  fi

  # 方案三：检查高层级包管理命令
  if command -v apt >/dev/null 2>&1; then
    printf '%s' "$PKG_DEB"
    return
  elif command -v yum >/dev/null 2>&1 || command -v dnf >/dev/null 2>&1; then
    printf '%s' "$PKG_RPM"
    return
  fi
  
  # 如果所有方法都失败
  printf '%s' "$PKG_UNKNOWN"
}

function download_package() {
    # 定义最大重试次数
    MAX_RETRIES=2
    RETRY_COUNT=0
    # 获取 REGION 的重试逻辑
    while [ $RETRY_COUNT -le $MAX_RETRIES ]; do
      REGION=$(wget -O - -a /var/log/wget.log -t 1 --timeout=10 "http://100.96.0.96/latest/region_id" 2>/dev/null)
      # 检查是否成功获取 REGION
      if [ -n "$REGION" ]; then
        log "成功获取 REGION: $REGION"
        break
      else
        RETRY_COUNT=$((RETRY_COUNT + 1))
        log "获取 REGION 失败,重试次数: $RETRY_COUNT"
        sleep 2  # 重试前等待 2 秒
      fi
    done
    # 如果最终未获取到 REGION,退出脚本
    if [ -z "$REGION" ]; then
      log "错误：无法获取 REGION,请检查网络连接或 URL 是否正确。"
      return
    fi

    rdma_bucket_url="rdma-toc-config-${REGION}.tos-${REGION}.ivolces.com"
    if [ "$PKG_TYPE" = "$PKG_RPM" ] ; then
        PKG_FILE="rdma-agent-${VERSION}-0.noarch.rpm"
        PKG_URL="http://${rdma_bucket_url}/$PKG_FILE"
        log "检测到 CentOS 系统"
        log "下载 $PKG_URL ..."
        wget "$PKG_URL"
        if [ $? -ne 0 ]; then
            log "下载$PKG_URL失败"
            return
        fi
    elif [ "$PKG_TYPE" = "$PKG_DEB" ] ; then
        PKG_FILE="rdma-agent_${VERSION}-0_all.deb"
        PKG_URL="http://${rdma_bucket_url}/$PKG_FILE"
        log "检测到 Debian/Ubuntu 系统"
        log "下载 $PKG_URL ..."
        wget "$PKG_URL"
        if [ $? -ne 0 ]; then
            log "下载$PKG_URL失败"
            return
        fi
    elif [ "$PKG_TYPE" = "$PKG_UNKNOWN" ] ; then
        log "未知的包类型,无法安装"
        exit 1
    fi
    DOWNLOAD_RC=0
}

function pkg_install() {
    if [ "$PKG_TYPE" = "$PKG_RPM" ] ; then
        if rpm -qa | grep -qw rdma-agent-infer; then
            log "检测到已安装 rdma-agent-infer,先卸载..."
            rpm -e rdma-agent-infer
        fi
        if rpm -qa | grep -qw rdma-agent; then
            log "检测到已安装 rdma-agent,先卸载..."
            rpm -e rdma-agent
        fi
        log "安装 $PKG_FILE ..."
        if rpm -i "$PKG_FILE"; then
            log "rpm 安装成功！"
        fi
    elif [ "$PKG_TYPE" = "$PKG_DEB" ] ; then
        if dpkg -l | grep -qw rdma-agent-infer; then
            log "检测到已安装 rdma-agent-infer,先卸载..."
            dpkg --purge rdma-agent-infer
        fi
        if dpkg -l | grep -qw rdma-agent; then
            log "检测到已安装 rdma-agent,先卸载..."
            dpkg --purge rdma-agent
        fi
        log "安装 $PKG_FILE ..."
        OUTPUT=$(dpkg -i "$PKG_FILE" 2>&1)
        RET=$?
        if [[ $RET -eq 0 ]]; then
            log "dpkg 安装成功！"
        fi
    elif [ "$PKG_TYPE" = "$PKG_UNKNOWN" ] ; then
        log "未知的包类型,无法安装"
        exit 1
    fi
    INSTALL_RC=0
}

move_agent_tobhpc() {
    SRC="${WORK_DIR}/etc/sysop/rdma_service/agent_tobhpc"
    DEST="$TARGET_PARENT"
    if [[ -d "$SRC" ]]; then
        mkdir -p "$DEST"
        rm -rf "$DEST$TARGET_NAME"
        mv "$SRC" "$DEST"
        log "agent_tobhpc 已移动到 $DEST"
    else
        log "未找到 $SRC,移动失败"
        exit 10
    fi
}

run_start_script() {
    if [[ -f "$START_SCRIPT" ]]; then
        chmod +x "$START_SCRIPT"
        log "执行 $START_SCRIPT ..."
        bash "$START_SCRIPT"
        log "start-rdma-agent-service.sh 已执行"
    else
        log "未找到 $START_SCRIPT, 无法执行"
        exit 11
    fi
}

function extract_install() {
    if [ "$PKG_TYPE" = "$PKG_RPM" ] ; then
        log "rpm 安装失败,自动切换为解包模式。"
        log "用 rpm2cpio + cpio 解包 ..."
        rpm2cpio "$PKG_FILE" | cpio -idmv
        log "解包完成"
        move_agent_tobhpc
        run_start_script
        log "安装完成"
    elif [ "$PKG_TYPE" = "$PKG_DEB" ] ; then
        log "dpkg 安装失败,自动切换为解包模式。"
        echo "用 ar 解包到当前目录..."
        ar x "$PKG_FILE"
        DATA_TAR=$(ls ./data.tar.* 2>/dev/null | head -n 1)
        if [[ -n "$DATA_TAR" ]]; then
            case "$DATA_TAR" in
                *.gz)  tar_opt="z";;
                *.xz)  tar_opt="J";;
                *.bz2) tar_opt="j";;
                *)     tar_opt="";;
            esac
            echo "解压$DATA_TAR到当前目录..."
            tar -x${tar_opt}vf "./$DATA_TAR" -C "./"
            log "解包完成"
            move_agent_tobhpc
            run_start_script
            log "安装完成"
        fi
    elif [ "$PKG_TYPE" = "$PKG_UNKNOWN" ] ; then
        log "未知的包类型,无法安装"
        exit 1
    fi
    INSTALL_RC=0
}



if [[ -z "$VERSION" ]]; then
    log "用法: $0 <版本号>"
    log "例如: $0 1.1.8"
    exit 1
fi
PKG_TYPE=$(get_pkg_type)
cd $WORK_DIR
download_package
if [ $DOWNLOAD_RC -ne 0 ]; then
    log "下载软件包失败，退出"
    exit 2
fi

pkg_install
if [ $INSTALL_RC -ne 0 ]; then
    extract_install
else
    log "安装成功"
fi
cd /
rm -rf $WORK_DIR