#!/bin/bash
set -e
#定义成数组，方便以后扩展 
declare -a images=(
    "ai-arm64/adapt-job:v3.5"
    "ai-arm64/ai-adminer:4.7.5"
    "ai-arm64/ai-explorer:v1.10"
    "ai-arm64/ai-mgr:v3.5"
    "ai-arm64/ai-mgr-web:v3.5"
    "ai-arm64/ai-user:v3.5"
    "ai-arm64/ai-user-web:v3.5"
    "ai-arm64/l4t-ml:r32.4.4-py3"
    "ai-arm64/labelme:latest" 
    "ai-arm64/label-studio:1.7.3-1"
    "ai-arm64/mariadb:10.7"
    "ai-arm64/minio/minio:RELEASE.2020-12-29T23-29-29Z"
    "ai-arm64/page-render:v3.5"
    "ai-arm64/pytorch-serving:0.10.0"
    "ai-arm64/tensorflow-serving:0.10.0"
    "ai-arm64/testarmbuild:latest"
    "frame_system/ascend-pytorch:24.0.0-A2-2.1.0-ubuntu20.04"
    "frame_system/tensorflow-modelzoo:22.0.0-notebook"
    "frame_system_npu/mis-tei:6.0.0-300I-Duo-aarch64"
)
#工作目录
BASE=/opt/ai-deployer
#记录已下载镜像的文件
DOWNLOAD_RECORD="/mnt/downloaded_images.txt"

#线上镜像仓库地址
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}')
HARBOR_ADMIN_PASSWORD=$(grep -r 'HARBOR_ADMIN_PASSWORD' $BASE/ai.config |cut -d'=' -f2|awk '{print $1}')
REGISTRY_LOCAL="$HARBOR_IP:$HARBOR_PORT"

#创建记录文件（如果不存在）
touch "$DOWNLOAD_RECORD"

#下载和推送单个镜像的函数
download_and_push_image() {
    local image=$1
    # ai-arm64的是ai，frame_system_npu的是frame_system

    local new_image=${image/ai-arm64/ai}
    new_image=${new_image/frame_system_npu/frame_system}

    echo "######开始处理镜像：$image######"
    
    #检查是否已经下载过
    if grep -q "^$image$" "$DOWNLOAD_RECORD"; then
        echo "镜像 $image 已经下载过，跳过"
        return 0
    fi
    
    #尝试下载和推送，如果失败则返回非0值
    if docker pull "$REGISTRY_DOMAIN/$image" && \
       docker tag "$REGISTRY_DOMAIN/$image" "$REGISTRY_LOCAL/$new_image" && \
       docker push "$REGISTRY_LOCAL/$new_image"; then
        #记录成功下载的镜像
        echo "$image" >> "$DOWNLOAD_RECORD"
        echo "######成功处理镜像：$image######"
        return 0
    else
        echo "######处理镜像失败：$image######"
        return 1
    fi
}

#主下载循环
for image in "${images[@]}"; do
    #如果下载失败，最多重试3次
    max_retries=3
    retry_count=0
    
    while [ $retry_count -lt $max_retries ]; do
        if download_and_push_image "$image"; then
            break
        else
            retry_count=$((retry_count + 1))
            if [ $retry_count -lt $max_retries ]; then
                echo "重试第 $retry_count 次下载 $image"
                sleep 5
            else
                echo "镜像 $image 下载失败，已达到最大重试次数"
            fi
        fi
    done
done

#检查是否所有镜像都已下载
total_images=${#images[@]}
downloaded_images=$(wc -l < "$DOWNLOAD_RECORD")

if [ "$total_images" -eq "$downloaded_images" ]; then
    echo "######所有镜像已成功下载并推送至本地仓库######"
    touch /mnt/got-frame-images
    #清理下载记录文件
    rm -f "$DOWNLOAD_RECORD"
else
    echo "######警告：只完成了 $downloaded_images/$total_images 个镜像的下载######"
    echo "请重新运行脚本以继续未完成的下载"
    exit 1
fi

#杀死后台进程
ps -ef | grep "get-frame-images" | grep -v grep | awk '{print $2}' | xargs kill -s 9