#!/bin/bash
# ============================================================
# 清理评测后遗留的 Docker 镜像
#
# 用途：
#   SWE-bench、Terminal-Bench 等 benchmark 会拉取大量 Docker 镜像，
#   评测完后很占磁盘。本脚本按命名空间批量删除这些镜像。
#
# 用法：
#   # 默认只预览，不会真的删除
#   bash bash/cleanup_docker_images.sh --preset swe
#
#   # 确认删除
#   bash bash/cleanup_docker_images.sh --preset swe --yes
#
#   # 删除 terminal-bench 镜像
#   bash bash/cleanup_docker_images.sh --preset terminal --yes
#
#   # 删除所有常见 benchmark 镜像
#   bash bash/cleanup_docker_images.sh --preset all-bench --yes
#
#   # 自定义匹配规则（支持多个，逗号分隔）
#   bash bash/cleanup_docker_images.sh --pattern 'swebench/*,alexgshaw/*' --yes
#
#   # 同时清理 dangling 镜像和未使用卷
#   bash bash/cleanup_docker_images.sh --preset all-bench --yes --prune
# ============================================================

set -euo pipefail

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

# 预设的 benchmark 镜像命名空间
PRESET_SWE=('swebench/*')
PRESET_TERMINAL=('alexgshaw/*')
PRESET_ALL_BENCH=()
PRESET_ALL_BENCH+=("${PRESET_SWE[@]}")
PRESET_ALL_BENCH+=("${PRESET_TERMINAL[@]}")

DRY_RUN=1
PRESET=''
PATTERNS=''
PRUNE=0

usage() {
    sed -n '7,34p' "$0"
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --preset)
            PRESET="${2:-}"
            shift 2
            ;;
        --pattern)
            PATTERNS="${2:-}"
            shift 2
            ;;
        --yes)
            DRY_RUN=0
            shift
            ;;
        --prune)
            PRUNE=1
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# 解析出要删除的 pattern 列表
PATTERN_LIST=()
if [[ -n "$PRESET" ]]; then
    case "$PRESET" in
        swe)
            PATTERN_LIST+=("${PRESET_SWE[@]}")
            ;;
        terminal)
            PATTERN_LIST+=("${PRESET_TERMINAL[@]}")
            ;;
        all-bench)
            PATTERN_LIST+=("${PRESET_ALL_BENCH[@]}")
            ;;
        *)
            echo "ERROR: unknown preset '$PRESET'. Available: swe, terminal, all-bench"
            exit 1
            ;;
    esac
fi

if [[ -n "$PATTERNS" ]]; then
    IFS=',' read -ra CUSTOM_PATTERNS <<< "$PATTERNS"
    PATTERN_LIST+=("${CUSTOM_PATTERNS[@]}")
fi

if [[ ${#PATTERN_LIST[@]} -eq 0 ]]; then
    echo "ERROR: 请指定 --preset 或 --pattern"
    usage
    exit 1
fi

echo "============================================================"
echo "清理目标 pattern:"
for p in "${PATTERN_LIST[@]}"; do
    echo "  - $p"
done
echo "============================================================"

# 收集所有匹配且没有被运行中容器使用的镜像 ID
MATCHED_IMAGES=()
RUNNING_IMAGES=$(docker ps --format '{{.Image}}' | sort -u)

for pattern in "${PATTERN_LIST[@]}"; do
    while IFS= read -r img; do
        [[ -z "$img" ]] && continue
        repo="${img%:*}"
        # 检查是否有运行中容器在使用该镜像
        if echo "$RUNNING_IMAGES" | grep -qx "$img"; then
            echo "SKIP (running container uses): $img"
            continue
        fi
        # 避免重复
        if [[ ! " ${MATCHED_IMAGES[*]} " =~ " ${img} " ]]; then
            MATCHED_IMAGES+=("$img")
        fi
    done < <(docker images --format '{{.Repository}}:{{.Tag}}' --filter=reference="$pattern")
done

if [[ ${#MATCHED_IMAGES[@]} -eq 0 ]]; then
    echo "没有匹配到可删除的镜像"
    exit 0
fi

echo ""
echo "匹配到 ${#MATCHED_IMAGES[@]} 个镜像："
for img in "${MATCHED_IMAGES[@]}"; do
    size=$(docker images --format '{{.Size}}' "$img" | head -1)
    echo "  $img ($size)"
done

# 计算标记总大小（注意共享层会被多次计算，仅作参考）
TOTAL_SIZE_GB=$(docker images --format '{{.Size}} {{.Repository}}:{{.Tag}}' \
    | awk -v img_list="${MATCHED_IMAGES[*]}" '
    BEGIN {
        n = split(img_list, imgs, " ");
        for (i=1; i<=n; i++) target[imgs[i]] = 1;
    }
    {
        size = $1;
        img = $2;
        if (!(img in target)) next;
        unit = substr(size, length(size)-1);
        val = substr(size, 1, length(size)-2);
        if (unit == "GB") sum += val * 1024 * 1024 * 1024;
        else if (unit == "MB") sum += val * 1024 * 1024;
        else if (unit == "kB") sum += val * 1024;
    }
    END {
        printf "%.2f", sum / 1024 / 1024 / 1024;
    }')

echo ""
echo "标记总大小约: ${TOTAL_SIZE_GB} GB（含共享层，实际释放会小于该值）"

if [[ "$DRY_RUN" -eq 1 ]]; then
    echo ""
    echo "当前是预览模式，不会删除。如需删除请加上 --yes"
    exit 0
fi

echo ""
echo "开始删除..."
FAILED=0
for img in "${MATCHED_IMAGES[@]}"; do
    if docker rmi -f "$img" >/dev/null 2>&1; then
        echo "  DELETED $img"
    else
        echo "  FAILED  $img"
        FAILED=1
    fi
done

if [[ "$PRUNE" -eq 1 ]]; then
    echo ""
    echo "清理 dangling 镜像和未使用卷..."
    docker image prune -f
    docker volume prune -f
fi

echo ""
echo "============================================================"
if [[ "$FAILED" -eq 0 ]]; then
    echo "清理完成"
else
    echo "部分镜像删除失败"
fi
echo "============================================================"