#!/bin/bash

# 批量构建所有更新的 Agent 并推送到 ACR (ARM64)
# 使用方法: ./build_all_agents.sh

set -e

# 默认配置
ACR_NAME="${ACR_NAME:-agnettaiji.azurecr.io}"
TAG="${1:-latest}"
PLATFORM="linux/arm64"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}  批量构建 Agent Templates (ARM64)${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""

# 登录 ACR
echo -e "${GREEN}==> 登录到 ACR...${NC}"
az acr login --name $(echo ${ACR_NAME} | cut -d'.' -f1)

# 检查 Docker buildx
if ! docker buildx version > /dev/null 2>&1; then
    echo -e "${RED}错误: Docker buildx 未安装${NC}"
    exit 1
fi

# 创建 builder
if ! docker buildx ls | grep -q multiarch-builder; then
    echo -e "${YELLOW}创建 multiarch-builder...${NC}"
    docker buildx create --name multiarch-builder --use
fi

docker buildx use multiarch-builder

# Agent 列表（更新路径）
declare -A AGENTS=(
    ["search-agent"]="agents/search_agent/search_agent.Dockerfile"
    ["search-agent-a2a"]="agents/search_agent/search_agent_A2A/search_agent_A2A.Dockerfile"
    ["search-agent-mcp"]="agents/search_agent/search_agent_MCP/search_agent_MCP.Dockerfile"
    ["azure-blob-agent"]="agents/azure_blob_agent/azure_blob_agent.Dockerfile"
    ["azure-blob-agent-a2a"]="agents/azure_blob_agent_a2a/azure_blob_agent_a2a.Dockerfile"
    ["azure-blob-agent-mcp"]="agents/azure_blob_agent_mcp/azure_blob_agent_mcp.Dockerfile"
    ["a2a-litellm-agent"]="agents/a2a_litellm_agent/a2a_litellm_agent.Dockerfile"
    ["mysql-agent"]="agents/mysql_agent/mysql_agent.Dockerfile"
    ["postgresql-agent"]="agents/postgresql_agent/postgresql_agent.Dockerfile"
    ["jina-search-agent"]="agents/jina_search_agent/jina_search_agent.Dockerfile"
    ["echo-agent"]="agents/echo_agent/echo_agent.Dockerfile"
)

# 构建函数
build_agent() {
    local name=$1
    local dockerfile=$2
    local image="${ACR_NAME}/ai-agents/${name}:${TAG}"
    
    echo ""
    echo -e "${BLUE}========================================${NC}"
    echo -e "${BLUE}构建: ${name}${NC}"
    echo -e "${BLUE}========================================${NC}"
    echo "Dockerfile: ${dockerfile}"
    echo "镜像: ${image}"
    echo "平台: ${PLATFORM}"
    echo ""
    
    # 构建并推送
    docker buildx build \
        --platform ${PLATFORM} \
        -f ${dockerfile} \
        -t ${image} \
        --push \
        .
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✅ ${name} 构建成功${NC}"
    else
        echo -e "${RED}❌ ${name} 构建失败${NC}"
        return 1
    fi
}

# 构建所有 Agent
SUCCESS_COUNT=0
FAIL_COUNT=0
FAILED_AGENTS=()

for name in "${!AGENTS[@]}"; do
    dockerfile="${AGENTS[$name]}"
    
    if [ -f "${dockerfile}" ]; then
        if build_agent "${name}" "${dockerfile}"; then
            ((SUCCESS_COUNT++))
        else
            ((FAIL_COUNT++))
            FAILED_AGENTS+=("${name}")
        fi
    else
        echo -e "${YELLOW}⚠️  跳过 ${name}: Dockerfile ${dockerfile} 不存在${NC}"
    fi
done

# 总结
echo ""
echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}  构建总结${NC}"
echo -e "${BLUE}========================================${NC}"
echo -e "${GREEN}成功: ${SUCCESS_COUNT}${NC}"
echo -e "${RED}失败: ${FAIL_COUNT}${NC}"

if [ ${FAIL_COUNT} -gt 0 ]; then
    echo -e "${RED}失败的 Agents:${NC}"
    for agent in "${FAILED_AGENTS[@]}"; do
        echo -e "  - ${agent}"
    done
    exit 1
fi

echo ""
echo -e "${GREEN}✅ 所有 Agent 构建成功！${NC}"
echo ""
echo "已推送的镜像:"
for name in "${!AGENTS[@]}"; do
    echo "  - ${ACR_NAME}/ai-agents/${name}:${TAG}"
done