#!/bin/bash

# 定义颜色输出
GREEN='\033[0;32m'
RED='\033[0;31m'
NC='\033[0m' # No Color

# 定义源和目标registry
SOURCE_REGISTRY="123.157.234.230:58080"
TARGET_REGISTRY="registry.local:38080"

# 登录源registry
login_registry() {
    echo "正在登录 ${SOURCE_REGISTRY}..."
    if ! docker login ${SOURCE_REGISTRY} -u admin -p Harbor@llm123; then
        echo -e "${RED}登录失败！${NC}"
        exit 1
    fi
    echo -e "${GREEN}登录成功！${NC}"
}

# 定义镜像列表
IMAGES=(
    "frame_system/llama2:v5"
    "frame_system/llama2-serving:0.13.0-cuda12.1-py38"
    "frame_system/stablediffusion-serving:v1.4-torch2.3-cuda12.1-py310"
    "frame_system/stablediffusion:v1.4-lora"
    "frame_system/llama2:7b-lora"
    "frame_system/chatglm2:6b-ptuning"
    "ai/gitlab-ce:15.11.13-ce.0"
    "frame_system/chatglm2-serving:0.13.0-cuda12.1-py310-6b"
    "frame_system/chatglm2:torch2.1-cuda12.1-py310-6b"
    "frame_system/ollama:notebook"
    "frame_system/sdwebui:notebook"
)

# 拉取镜像
pull_images() {
    echo "开始拉取镜像..."
    for image in "${IMAGES[@]}"; do
        echo "正在拉取 ${SOURCE_REGISTRY}/${image}"
        if ! docker pull ${SOURCE_REGISTRY}/${image}; then
            echo -e "${RED}拉取 ${image} 失败！${NC}"
            return 1
        fi
    done
    echo -e "${GREEN}所有镜像拉取完成！${NC}"
}

# 标记镜像
tag_images() {
    echo "开始标记镜像..."
    for image in "${IMAGES[@]}"; do
        echo "正在标记 ${image}"
        if ! docker tag ${SOURCE_REGISTRY}/${image} ${TARGET_REGISTRY}/${image}; then
            echo -e "${RED}标记 ${image} 失败！${NC}"
            return 1
        fi
    done
    echo -e "${GREEN}所有镜像标记完成！${NC}"
}

# 推送镜像
push_images() {
    echo "开始推送镜像..."
    for image in "${IMAGES[@]}"; do
        echo "正在推送 ${TARGET_REGISTRY}/${image}"
        if ! docker push ${TARGET_REGISTRY}/${image}; then
            echo -e "${RED}推送 ${image} 失败！${NC}"
            return 1
        fi
    done
    echo -e "${GREEN}所有镜像推送完成！${NC}"
}

# 清理镜像
cleanup_images() {
    echo "开始清理镜像..."
    for image in "${IMAGES[@]}"; do
        echo "正在删除 ${SOURCE_REGISTRY}/${image}"
        if ! docker rmi ${SOURCE_REGISTRY}/${image}; then
            echo -e "${RED}删除源镜像 ${image} 失败！${NC}"
        fi
        
        echo "正在删除 ${TARGET_REGISTRY}/${image}"
        if ! docker rmi ${TARGET_REGISTRY}/${image}; then
            echo -e "${RED}删除目标镜像 ${image} 失败！${NC}"
        fi
    done
    echo -e "${GREEN}镜像清理完成！${NC}"
}

# 主函数
main() {
    login_registry || exit 1
    pull_images || exit 1
    tag_images || exit 1
    push_images || exit 1
    cleanup_images
    echo -e "${GREEN}镜像迁移完成！${NC}"
}

# 执行主函数
main