#!/bin/bash
# ============================================================
# EvalScope 完整部署脚本
# 用法: bash deploy_evalscope.sh [OPTIONS] <BASE_DIR> [TOKEN]
# 示例:
#   bash deploy_evalscope.sh /data1/sora
#   bash deploy_evalscope.sh --download-swe /data1/sora
#
# 断点续传：已加载的镜像/tar 包会记录在 <BASE_DIR>/evalscope/docker/.deploy_checkpoint
#           重复执行时会自动跳过，避免二次加载。
# ============================================================

set -euo pipefail

# --------------------------------------------------
# 1. 参数解析
# --------------------------------------------------
DOWNLOAD_SWE=0
POSITIONAL_ARGS=()

while [[ $# -gt 0 ]]; do
    case "$1" in
        --download-swe)
            DOWNLOAD_SWE=1
            shift
            ;;
        --help|-h)
            echo "用法: bash deploy_evalscope.sh [OPTIONS] <BASE_DIR> [TOKEN]"
            echo "选项:"
            echo "  --download-swe    同时下载 SWE-Bench 镜像包（默认不下载，体积很大）"
            echo "  -h, --help        显示帮助"
            exit 0
            ;;
        -*)
            echo "未知选项: $1"
            echo "使用 --help 查看用法"
            exit 1
            ;;
        *)
            POSITIONAL_ARGS+=("$1")
            shift
            ;;
    esac
done

set -- "${POSITIONAL_ARGS[@]}"

BASE_DIR="${1:-/data1/sora}"           # 基础目录，默认 /data1/sora
DEFAULT_TOKEN="ms-3d554a39-6e07-496d-8022-0b0ee64a6389"
MODELSCOPE_TOKEN="${2:-$DEFAULT_TOKEN}"

# 派生路径
EVALSCOPE_DIR="$BASE_DIR/evalscope"
DOCKER_DIR="$EVALSCOPE_DIR/docker"
IMAGES_DIR="$DOCKER_DIR/images"
CONTAINERD_DIR="$DOCKER_DIR/containerd"
SWE_IMAGES_DIR="$DOCKER_DIR/swe_images"
DATASETS_DIR="$EVALSCOPE_DIR/datasets"
CHECKPOINT_FILE="$DOCKER_DIR/.deploy_checkpoint"

# 断点：记录/检查已加载的镜像或 tar 包
mark_loaded() {
    local key="$1"
    mkdir -p "$DOCKER_DIR"
    touch "$CHECKPOINT_FILE"
    if ! grep -qx "$key" "$CHECKPOINT_FILE" 2>/dev/null; then
        echo "$key" >> "$CHECKPOINT_FILE"
    fi
}

is_loaded() {
    local key="$1"
    if [ -f "$CHECKPOINT_FILE" ] && grep -qx "$key" "$CHECKPOINT_FILE" 2>/dev/null; then
        return 0
    fi
    return 1
}

# --------------------------------------------------
# 2. 安装依赖
# --------------------------------------------------
echo "==> 安装 modelscope..."
pip install modelscope

# --------------------------------------------------
# 3. 登录 ModelScope
# --------------------------------------------------
if [ -n "$MODELSCOPE_TOKEN" ]; then
    echo "==> 登录 ModelScope..."
    modelscope login --token "$MODELSCOPE_TOKEN"
fi

# --------------------------------------------------
# 4. 下载数据集
# --------------------------------------------------
echo "==> 下载 evalscope 数据集..."
mkdir -p "$DATASETS_DIR"
python3 -c "
from modelscope.hub.snapshot_download import snapshot_download
snapshot_download(
    'SoraAmami/evalscope-datasets',
    repo_type='dataset',
    cache_dir='$EVALSCOPE_DIR',
    local_dir='$DATASETS_DIR'
)
"

# --------------------------------------------------
# 5. 下载 Docker 镜像包
# --------------------------------------------------
echo "==> 下载 evalscope Docker 镜像..."
mkdir -p "$DOCKER_DIR"
python3 -c "
from modelscope.hub.file_download import model_file_download
model_file_download(
    model_id='SoraAmami/evalscope-docker',
    file_path='evalscope-complete-py312.tar.gz',
    local_dir='$DOCKER_DIR'
)
"

# --------------------------------------------------
# 6. 配置 Docker 和 containerd 数据目录
# --------------------------------------------------
echo "==> 配置 Docker 和 containerd..."

# 停止服务
sudo systemctl stop docker.socket 2>/dev/null || true
sudo systemctl stop docker 2>/dev/null || true
sudo systemctl stop containerd 2>/dev/null || true

# 创建数据目录
mkdir -p "$IMAGES_DIR"
mkdir -p "$CONTAINERD_DIR"

# 配置 Docker
echo "==> 写入 Docker 配置..."
sudo tee /etc/docker/daemon.json <<EOF
{
  "data-root": "$IMAGES_DIR",
  "features": {
    "containerd-snapshotter": true
  },
  "registry-mirrors": [
    "https://docker.m.daocloud.io",
    "https://docker.1ms.run",
    "https://hub.rat.dev",
    "https://docker.1panel.live",
    "https://dockerproxy.com",
    "https://hub-mirror.c.163.com",
    "https://mirror.baidubce.com",
    "https://docker.mirrors.ustc.edu.cn",
    "https://docker.mirrors.sjtug.sjtu.edu.cn",
    "https://docker.nju.edu.cn",
    "https://docker.mirrors.tuna.tsinghua.edu.cn"
  ]
}
EOF

# 配置 containerd
echo "==> 写入 containerd 配置..."
sudo tee /etc/containerd/config.toml <<EOF
root = "$CONTAINERD_DIR"
state = "/run/containerd"
EOF

# 重启服务
sudo systemctl reset-failed docker.service 2>/dev/null || true
sudo systemctl start containerd
sudo systemctl start docker

# --------------------------------------------------
# 7. 加载 evalscope 环境镜像
# --------------------------------------------------
echo "==> 加载 evalscope Docker 镜像..."
if docker images --format '{{.Repository}}:{{.Tag}}' | grep -qx "evalscope-complete-py312:latest"; then
    echo "evalscope-complete-py312:latest 已存在，跳过加载"
elif is_loaded "evalscope-complete-py312.tar.gz"; then
    echo "evalscope-complete-py312.tar.gz 已加载过（checkpoint），跳过加载"
else
    docker load -i "$DOCKER_DIR/evalscope-complete-py312.tar.gz"
    mark_loaded "evalscope-complete-py312.tar.gz"
fi

# --------------------------------------------------
# 8. 克隆代码仓库到根目录
# --------------------------------------------------
echo "==> 克隆 evalscope 代码..."

if [ -f "$EVALSCOPE_DIR/bash/run.py" ]; then
    echo "代码已存在于 $EVALSCOPE_DIR，跳过克隆"
else
    echo "正在克隆代码到 $EVALSCOPE_DIR ..."
    tmp_backup="$BASE_DIR/.evalscope_deploy_backup_$(date +%s)"
    mkdir -p "$tmp_backup"

    # 保留已有的 datasets/ 和 docker/，其余清空后克隆
    for item in "$EVALSCOPE_DIR"/* "$EVALSCOPE_DIR"/.[!.]* "$EVALSCOPE_DIR"/..?*; do
        [ -e "$item" ] || continue
        basename_item=$(basename "$item")
        case "$basename_item" in
            .|..|.git) continue ;;
            datasets|docker)
                mv "$item" "$tmp_backup/$basename_item" ;;
            *)
                rm -rf "$item" ;;
        esac
    done

    git clone "https://git.meta-stone.net/sora/evalstone.git" "$EVALSCOPE_DIR"

    # 恢复 datasets/ docker/
    for item in "$tmp_backup"/* "$tmp_backup"/.[!.]* "$tmp_backup"/..?*; do
        [ -e "$item" ] || continue
        basename_item=$(basename "$item")
        [ "$basename_item" = "." ] || [ "$basename_item" = ".." ] && continue
        if [ -e "$EVALSCOPE_DIR/$basename_item" ]; then
            continue
        fi
        mv "$item" "$EVALSCOPE_DIR/$basename_item"
    done
    rm -rf "$tmp_backup" 2>/dev/null || true
fi

# --------------------------------------------------
# 9. 准备 sandbox 镜像
# --------------------------------------------------
echo "==> 准备 sandbox 镜像..."
SANDBOX_IMAGES_DIR="$DOCKER_DIR/sandbox_images"
mkdir -p "$SANDBOX_IMAGES_DIR"

# 从 ModelScope 下载预打包的 sandbox 镜像（不再使用 docker pull，避免 DockerHub 限速）
download_sandbox_image_from_modelscope() {
    local file="$1"
    if [ -f "$SANDBOX_IMAGES_DIR/$file" ]; then
        echo "sandbox 镜像包 $file 已存在，跳过下载"
        return 0
    fi
    echo "==> 从 ModelScope 下载 $file ..."
    # 优先使用新版 modelscope_hub API（支持 dataset repo）
    python3 -c "
from modelscope_hub import HubApi
api = HubApi()
api.download_file(
    repo_id='SoraAmami/evalscope-sandbox-images',
    repo_type='dataset',
    path='$file',
    local_dir='$SANDBOX_IMAGES_DIR'
)
" && return 0
    echo "警告：从 ModelScope 下载 $file 失败"
    return 1
}

load_sandbox_image() {
    local file="$1"
    local expected_image="$2"
    if docker images --format '{{.Repository}}:{{.Tag}}' | grep -qx "$expected_image"; then
        echo "镜像 $expected_image 已存在，跳过加载"
        return 0
    fi
    if is_loaded "$file"; then
        echo "镜像包 $file 已加载过（checkpoint），跳过加载"
        return 0
    fi
    if [ -f "$SANDBOX_IMAGES_DIR/$file" ]; then
        echo "==> 加载 $file ..."
        docker load -i "$SANDBOX_IMAGES_DIR/$file"
        mark_loaded "$file"
        return 0
    fi
    return 1
}

# 9.1 bigcodebench-sandbox:latest
if docker images --format '{{.Repository}}:{{.Tag}}' | grep -qx "bigcodebench-sandbox:latest"; then
    echo "bigcodebench-sandbox:latest 已存在，跳过"
else
    if ! load_sandbox_image "bigcodebench-sandbox.tar.gz" "bigcodebench-sandbox:latest"; then
        download_sandbox_image_from_modelscope "bigcodebench-sandbox.tar.gz" && \
            load_sandbox_image "bigcodebench-sandbox.tar.gz" "bigcodebench-sandbox:latest"
    fi
fi

# 9.2 python:3.11-slim
if docker images --format '{{.Repository}}:{{.Tag}}' | grep -qx "python:3.11-slim"; then
    echo "python:3.11-slim 已存在，跳过"
else
    if ! load_sandbox_image "python-3.11-slim.tar.gz" "python:3.11-slim"; then
        download_sandbox_image_from_modelscope "python-3.11-slim.tar.gz" && \
            load_sandbox_image "python-3.11-slim.tar.gz" "python:3.11-slim"
    fi
fi

# --------------------------------------------------
# 10. 下载 SWE-Bench 镜像包（默认不下载，需加 --download-swe）
# --------------------------------------------------
if [ "$DOWNLOAD_SWE" -eq 1 ]; then
    echo "==> 下载 SWE-Bench 镜像包..."
    mkdir -p "$SWE_IMAGES_DIR"
    modelscope download \
        --repo-type dataset \
        --local_dir "$SWE_IMAGES_DIR" \
        SoraAmami/swe-bench-verified-images || echo "警告：SWE-Bench 镜像下载失败"
else
    echo "==> 跳过 SWE-Bench 镜像包下载（如需下载请加 --download-swe）"
fi

# --------------------------------------------------
# 11. 加载 SWE-Bench 镜像（如果存在）
# --------------------------------------------------
if [ -d "$SWE_IMAGES_DIR" ] && ls "$SWE_IMAGES_DIR"/swebench_batch_*.tar.gz 1>/dev/null 2>&1; then
    echo "==> 加载 SWE-Bench 镜像..."
    TARS=("$SWE_IMAGES_DIR"/swebench_batch_*.tar.gz)
    TOTAL=${#TARS[@]}
    echo "共找到 $TOTAL 个镜像包"

    IDX=0
    LOADED_COUNT=0
    SKIPPED_COUNT=0
    for TAR in "${TARS[@]}"; do
        IDX=$((IDX + 1))
        BASENAME=$(basename "$TAR")
        if is_loaded "$BASENAME"; then
            echo "[$IDX/$TOTAL] 跳过（已加载过）: $BASENAME"
            SKIPPED_COUNT=$((SKIPPED_COUNT + 1))
            continue
        fi
        echo "[$IDX/$TOTAL] 加载 $BASENAME ..."
        docker load -i "$TAR"
        mark_loaded "$BASENAME"
        LOADED_COUNT=$((LOADED_COUNT + 1))
    done

    echo "SWE-Bench 镜像加载完成：本次加载 $LOADED_COUNT 个，跳过 $SKIPPED_COUNT 个，已加载 $(docker images | grep -c '^swebench/' || true) 个"
else
    echo "跳过 SWE-Bench 镜像加载（未找到镜像包）"
fi

# --------------------------------------------------
# 12. 验证部署
# --------------------------------------------------
echo ""
echo "========================================"
echo "部署完成！"
echo "========================================"
echo "基础目录:     $BASE_DIR"
echo "EvalScope:    $EVALSCOPE_DIR"
echo "数据集:       $DATASETS_DIR"
echo "Docker 数据:  $IMAGES_DIR"
echo "Containerd:   $CONTAINERD_DIR"
echo ""
echo "Docker Root Dir:"
docker info 2>/dev/null | grep "Docker Root Dir" || echo "Docker 未运行"
echo ""
echo "已加载镜像:"
docker images | grep -E "evalscope|swebench|bigcodebench|python" || true
echo ""
echo "运行 EvalScope:"
echo "  cd $EVALSCOPE_DIR"
echo "  docker run -it --rm \\"
echo "    --network host \\"
echo "    -v $EVALSCOPE_DIR:/opt/evalscope \\"
echo "    -v /var/run/docker.sock:/var/run/docker.sock \\"
echo "    evalscope-complete-py312:latest \\"
echo "    bash"
echo "========================================"