#!/bin/bash
# ============================================================
# 预加载 Terminal-Bench 2.1 数据集和 Docker 镜像
#
# 用途：
#   terminal_bench_v2_1 运行时 Harbor 会按需拉取 Docker 镜像，
#   导致首次运行很慢。本脚本先把数据集和 89 个镜像拉到本地，
#   后续评测就只需要跑 agent，不用等镜像下载。
#
# 用法：
#   bash bash/images_load/preload_terminal_bench_images.sh
#
# 注意：
#   - 需要 Docker 环境
#   - 89 个镜像大约需要几十 GB 磁盘空间
#   - 执行时间取决于网速，可能需要 30 分钟 ~ 数小时
# ============================================================

set -euo pipefail

DATASET_NAME="terminal-bench/terminal-bench-2-1"
CACHE_DIR="${HOME}/.cache/harbor/tasks/packages/terminal-bench"

echo "============================================================"
echo "1. 预下载 Terminal-Bench 2.1 数据集"
echo "============================================================"
harbor download "${DATASET_NAME}" --cache

echo ""
echo "============================================================"
echo "2. 提取所有任务镜像列表"
echo "============================================================"
IMAGE_FILE="/tmp/terminal_bench_images.txt"
if [[ ! -d "${CACHE_DIR}" ]]; then
    echo "ERROR: 数据集缓存目录不存在: ${CACHE_DIR}"
    exit 1
fi

grep -hR "^docker_image" "${CACHE_DIR}" 2>/dev/null \
    | sed 's/docker_image = //; s/"//g' \
    | sort -u > "${IMAGE_FILE}"

IMAGE_COUNT=$(wc -l < "${IMAGE_FILE}" | tr -d ' ')
echo "找到 ${IMAGE_COUNT} 个唯一镜像，列表保存在 ${IMAGE_FILE}"

echo ""
echo "============================================================"
echo "3. 并行拉取 Docker 镜像"
echo "============================================================"

# 脚本放在 bash/images_load/，项目根目录是上一级再上一级
PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
LOG_DIR="${PROJECT_ROOT}/logs"
LOG_FILE="${LOG_DIR}/pull_terminal_images.log"
mkdir -p "${LOG_DIR}"
> "${LOG_FILE}"

echo "开始拉取镜像（后台 4 并发），日志保存在 ${LOG_FILE}"

PULL_FAILED=0
while IFS= read -r image; do
    [[ -z "$image" ]] && continue
    (
        echo "[$(date '+%Y-%m-%d %H:%M:%S')] Pulling ${image} ..." >> "${LOG_FILE}"
        if docker pull "${image}" >> "${LOG_FILE}" 2>&1; then
            echo "[$(date '+%Y-%m-%d %H:%M:%S')] OK ${image}" >> "${LOG_FILE}"
        else
            echo "[$(date '+%Y-%m-%d %H:%M:%S')] FAILED ${image}" >> "${LOG_FILE}"
            exit 1
        fi
    ) &
    # 控制并发数 4
    if (( $(jobs -r -p | wc -l) >= 4 )); then
        wait -n || PULL_FAILED=1
    fi
done < "${IMAGE_FILE}"

wait || PULL_FAILED=1

echo ""
echo "============================================================"
if [[ "$PULL_FAILED" -eq 0 ]]; then
    echo "所有镜像拉取完成"
else
    echo "部分镜像拉取失败，请查看日志: ${LOG_FILE}"
    exit 1
fi
echo "============================================================"