#!/bin/bash
# 将本地已拉取的 swebench 实例镜像打包保存，方便通过移动硬盘/内网传送到其他机器。
# 注意：swe_bench_verified 全量约 500 个镜像、~2TB，无法上传到 ModelScope 这类代码/文件仓库。

set -e

OUTPUT_DIR="${1:-/data1/sora/evalscope/docker/swe_images}"
BATCH_SIZE="${2:-50}"

mkdir -p "$OUTPUT_DIR"

echo "正在列出所有 swebench 镜像..."
mapfile -t IMAGES < <(docker images --format '{{.Repository}}:{{.Tag}}' | grep '^swebench/')
TOTAL=${#IMAGES[@]}

echo "共找到 $TOTAL 个 swebench 镜像，将按每批 $BATCH_SIZE 个保存到 $OUTPUT_DIR"

BATCH=()
BATCH_IDX=0
COUNT=0

for IMG in "${IMAGES[@]}"; do
    BATCH+=("$IMG")
    COUNT=$((COUNT + 1))

    if ((${#BATCH[@]} >= BATCH_SIZE)) || ((COUNT == TOTAL)); then
        BATCH_IDX=$((BATCH_IDX + 1))
        TAR_FILE="$OUTPUT_DIR/swebench_batch_$(printf '%03d' $BATCH_IDX).tar"
        echo "[$COUNT/$TOTAL] 保存批次 $BATCH_IDX -> $TAR_FILE"
        docker save "${BATCH[@]}" -o "$TAR_FILE"
        gzip -f "$TAR_FILE"
        BATCH=()
    fi
done

echo "完成。镜像列表："
ls -lh "$OUTPUT_DIR"