#!/bin/bash
# ============================================================
# EvalScope 完整部署脚本
# 用法: bash deploy_evalscope.sh <BASE_DIR> [TOKEN]
# 示例: bash deploy_evalscope.sh /data1/sora ms-3d554a39-6e07-496d-8022-0b0ee64a6389
# ============================================================

set -euo pipefail

# --------------------------------------------------
# 1. 参数解析
# --------------------------------------------------
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"

# --------------------------------------------------
# 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 镜像..."
docker load -i "$DOCKER_DIR/evalscope-complete-py312.tar.gz"

# --------------------------------------------------
# 8. 克隆代码仓库
# --------------------------------------------------
echo "==> 克隆 evalscope 代码..."
cd "$EVALSCOPE_DIR"
if [ ! -d "evalstone" ]; then
    git clone https://git.meta-stone.net/sora/evalstone.git
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 [ -f "$SANDBOX_IMAGES_DIR/$file" ]; then
        echo "==> 加载 $file ..."
        docker load -i "$SANDBOX_IMAGES_DIR/$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 镜像包（可选）
# --------------------------------------------------
if [ -n "$MODELSCOPE_TOKEN" ]; 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 镜像下载失败"
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
    for TAR in "${TARS[@]}"; do
        IDX=$((IDX + 1))
        echo "[$IDX/$TOTAL] 加载 $(basename "$TAR")..."
        docker load -i "$TAR"
    done

    echo "SWE-Bench 镜像加载完成，已加载 $(docker images | grep '^swebench/' | wc -l) 个"
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 "========================================"