#!/bin/bash

# 设定一些常量
HARBOR_REPO="harbor-dev.platform.baai-inner.ac.cn/flagdiagnose"
HARBOR_USERNAME="diagnose"
HARBOR_PASSWORD="RnMdBqgj3J"
CURRENT_TIME=$(date +"%Y%m%d%H%M")  # 当前时间
COMMIT_ID=$(git rev-parse --short HEAD)  # 获取当前代码的 commit id
SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)   # 获取脚本当前目录
echo $SCRIPT_DIR
ROOT_DIR="$SCRIPT_DIR/../.."
DOCKERFILE="$SCRIPT_DIR/../../env/Dockerfile"

#获取 version
version=$(PYTHONPATH=../ python3 -c "from version import __version__; print(__version__)")

BRANCH_NAME=$1
if [ -z "$BRANCH_NAME" ]; then
    BRANCH_NAME="main"
    echo "没有提供分支名，默认使用 main 分支"
fi

# 切换到指定分支
cd "$SCRIPT_DIR/../../" || { echo "无法进入代码目录"; exit 1; }
echo "切换到分支 $BRANCH_NAME"
git checkout "$BRANCH_NAME" || { echo "切换到分支 $BRANCH_NAME 失败"; exit 1; }


# 登录到 Harbor 仓库
docker login $HARBOR_REPO -u $HARBOR_USERNAME -p $HARBOR_PASSWORD
if [ $? -ne 0 ]; then
  echo "Docker login to Harbor failed! Exiting..."
  exit 1
fi

# 镜像标签格式：<repo>/<image_name>:<commit_id>-<timestamp>
build_and_push() {
  BASE_IMAGE=$1  # "nv" 或 "cambricon"

  # 选择基础镜像
  if [[ "$BASE_IMAGE" == "nv" ]]; then
    FROM_IMAGE="nvcr.io/nvidia/pytorch:24.09-py3"  # NVIDIA 镜像
  elif [[ "$BASE_IMAGE" == "cambricon" ]]; then
    FROM_IMAGE="nvcr.io/nvidia/pytorch:24.09-py3"  # Cambricon 镜像
  else
    echo "Unknown base image type: $BASE_IMAGE"
    exit 1
  fi

  # 修改 Dockerfile 中的 FROM 指令
  sed -i "s|^FROM .*|FROM $FROM_IMAGE|" $DOCKERFILE
  BASE_IMAGE="toolkits_${BASE_IMAGE}"
  TAG="$HARBOR_REPO/$BASE_IMAGE:$BRANCH_NAME-$COMMIT_ID-$CURRENT_TIME-$version"  # 镜像标签

  # 构建 Docker 镜像
  echo "Building image $TAG..."
  docker build -t $TAG -f $DOCKERFILE $ROOT_DIR

  if [ $? -eq 0 ]; then
    echo "Build successful, pushing image to Harbor..."
    # 推送到 Harbor
    docker push $TAG
    if [ $? -eq 0 ]; then
      echo "Successfully pushed $TAG to Harbor!"
    else
      echo "Failed to push $TAG to Harbor!"
      exit 1
    fi
  else
    echo "Docker build failed!"
    exit 1
  fi
# 返回镜像的完整标签，供后续拉取使用
  echo "$TAG"
}

# 生成两个镜像，一个基于 nv，另一个基于 cambricon
NV_TAG=$(build_and_push "nv")
CAMBRICON_TAG=$(build_and_push "cambricon")
DOCKER_PULL_NV="docker pull $NV_TAG"
DOCKER_PULL_CAMBRICON="docker pull $CAMBRICON_TAG"
echo "$DOCKER_PULL_NV"
echo "$DOCKER_PULL_CAMBRICON"


curl -X POST -H "Content-Type: application/json" \
    -d "{\"msg_type\":\"text\",\"content\":{\"text\":\"最新版本诊断工具下载地址: \nnv: $DOCKER_PULL_NV\ncambricon: $DOCKER_PULL_CAMBRICON\"}}" \
    https://open.feishu.cn/open-apis/bot/v2/hook/6fc85b34-fa64-4377-98f4-c56df546a39c


echo "All operations completed successfully."
