#!/bin/bash

# 一键启动脚本
# 用于初始化目录、构建镜像并启动所有服务

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TEST_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
cd "$TEST_DIR"

# 解析参数
FORCE_REBUILD=false
if [[ "$1" == "--rebuild" ]]; then
    FORCE_REBUILD=true
fi

echo "=========================================="
echo "  Argus Metrics 一键启动脚本"
echo "=========================================="
echo ""
echo "当前工作目录: $TEST_DIR"
echo ""

# 检查 Docker 和 Docker Compose
if ! command -v docker &> /dev/null; then
    echo "错误: 未找到 docker 命令，请先安装 Docker"
    exit 1
fi

# 检查 docker compose 命令
if ! docker compose version &> /dev/null 2>&1; then
    echo "错误: 未找到 docker compose 命令，请确保 Docker Compose V2 已安装"
    exit 1
fi
echo "使用: docker compose"
echo "Compose 文件: $TEST_DIR/docker-compose.yml"
echo ""

# 检查必要的构建目录
echo "检查构建目录..."
BUILD_DIRS=(
    "../ftp/build"
    "../prometheus/build"
    "../grafana/build"
    "client-test-node/build"
    "client-test-gpu-node/build"
)

for dir in "${BUILD_DIRS[@]}"; do
    if [ ! -d "$dir" ]; then
        echo "错误: 构建目录不存在: $dir"
        echo "完整路径: $(cd "$(dirname "$dir")" 2>/dev/null && pwd)/$(basename "$dir")"
        exit 1
    else
        echo "  ✓ 找到: $dir"
    fi
done
echo ""

# 检查并创建 .env 文件
if [ ! -f .env ]; then
    echo "未找到 .env 文件，从 env.example 创建..."
    cp env.example .env
    echo "已创建 .env 文件，请根据需要修改配置"
fi

# 加载环境变量
source .env

# 检查并创建 Docker 网络
echo "检查 Docker 网络..."
NETWORK_NAME="argus-debug-net"
if docker network inspect "$NETWORK_NAME" >/dev/null 2>&1; then
    echo "网络 $NETWORK_NAME 已存在"
else
    echo "创建网络 $NETWORK_NAME..."
    docker network create --driver bridge --subnet 172.30.0.0/16 "$NETWORK_NAME"
    echo "网络创建成功"
fi
echo ""

echo "1. 初始化目录结构..."
bash "$SCRIPT_DIR/init-directories.sh"

echo ""
echo "2. 准备 Docker 镜像..."

# 检查镜像是否存在
IMAGE_CACHE_DIR="$TEST_DIR/images-cache"
IMAGES=("argus-metric-ftp:latest" "argus-metric-prometheus:latest" "argus-metric-grafana:latest" "argus-metric-test-node:latest" "argus-metric-test-gpu-node:latest")
all_images_exist=true

for image in "${IMAGES[@]}"; do
    if ! docker images --format "{{.Repository}}:{{.Tag}}" | grep -q "^${image}$"; then
        all_images_exist=false
        break
    fi
done

if $FORCE_REBUILD; then
    echo "强制重新构建镜像（--rebuild 模式）..."
    cd "$TEST_DIR"
    docker compose build --no-cache
    echo "镜像重新构建完成"
elif $all_images_exist; then
    echo "所有镜像已存在，跳过构建"
else
    echo "检测到缺失镜像，尝试从缓存加载..."
    
    # 尝试从缓存加载
    loaded_from_cache=false
    if [ -d "$IMAGE_CACHE_DIR" ]; then
        for image in "${IMAGES[@]}"; do
            if ! docker images --format "{{.Repository}}:{{.Tag}}" | grep -q "^${image}$"; then
                # 镜像不存在，尝试加载
                case "$image" in
                    "argus-metric-ftp:latest")
                        cache_file="${IMAGE_CACHE_DIR}/argus-ftp.tar"
                        ;;
                    "argus-metric-prometheus:latest")
                        cache_file="${IMAGE_CACHE_DIR}/argus-prometheus.tar"
                        ;;
                    "argus-metric-grafana:latest")
                        cache_file="${IMAGE_CACHE_DIR}/argus-grafana.tar"
                        ;;
                    "argus-metric-test-node:latest")
                        cache_file="${IMAGE_CACHE_DIR}/argus-test-node.tar"
                        ;;
                    "argus-metric-test-gpu-node:latest")
                        cache_file="${IMAGE_CACHE_DIR}/argus-test-gpu-node.tar"
                        ;;
                esac
                
                if [ -f "$cache_file" ]; then
                    echo "  从缓存加载: $image"
                    docker load -i "$cache_file"
                    loaded_from_cache=true
                fi
            fi
        done
    fi
    
    # 检查加载后是否还有缺失的镜像
    need_build=false
    for image in "${IMAGES[@]}"; do
        if ! docker images --format "{{.Repository}}:{{.Tag}}" | grep -q "^${image}$"; then
            need_build=true
            break
        fi
    done
    
    if $need_build; then
        echo ""
        echo "部分镜像缺失，开始构建..."
        echo "工作目录: $(pwd)"
        cd "$TEST_DIR"
        docker compose build --no-cache
        
        # 询问是否保存镜像
        echo ""
        read -p "是否保存镜像到缓存以便下次快速启动? (Y/n): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Nn]$ ]]; then
            mkdir -p "$IMAGE_CACHE_DIR"
            echo "保存镜像到缓存..."
            for image in "${IMAGES[@]}"; do
                case "$image" in
                    "argus-metric-ftp:latest")
                        docker save -o "${IMAGE_CACHE_DIR}/argus-ftp.tar" "$image" && echo "  已保存: argus-ftp.tar"
                        ;;
                    "argus-metric-prometheus:latest")
                        docker save -o "${IMAGE_CACHE_DIR}/argus-prometheus.tar" "$image" && echo "  已保存: argus-prometheus.tar"
                        ;;
                    "argus-metric-grafana:latest")
                        docker save -o "${IMAGE_CACHE_DIR}/argus-grafana.tar" "$image" && echo "  已保存: argus-grafana.tar"
                        ;;
                    "argus-metric-test-node:latest")
                        docker save -o "${IMAGE_CACHE_DIR}/argus-test-node.tar" "$image" && echo "  已保存: argus-test-node.tar"
                        ;;
                    "argus-metric-test-gpu-node:latest")
                        docker save -o "${IMAGE_CACHE_DIR}/argus-test-gpu-node.tar" "$image" && echo "  已保存: argus-test-gpu-node.tar"
                        ;;
                esac
            done
            echo "镜像已保存到: $IMAGE_CACHE_DIR/"
        fi
    elif $loaded_from_cache; then
        echo ""
        echo "所有镜像已从缓存加载完成！"
    fi
fi

echo ""
echo "3. 启动基础服务..."
cd "$TEST_DIR"
# 启动除GPU节点外的所有服务
docker compose up -d ftp prometheus grafana test-node test-gpu-node

echo ""
echo "4. 等待服务启动..."
sleep 5