#!/bin/bash
# ============================================================
# 预下载 Pier mini-swe-agent 安装依赖（uv 二进制 + wheel）
#
# DeepSWE 每个任务镜像都会再执行：
#   curl https://astral.sh/uv/0.7.13/install.sh | sh
#   uv tool install mini-swe-agent
# FROM 镜像不同，Docker 层缓存不共享，全样本会反复打 GitHub/PyPI。
#
# 本脚本把 uv 和 mini-swe-agent 的 wheel 放到
#   bash/images_load/offline_pier_agent/
# 评测时 evalscope 会 COPY 进构建上下文，构建不再访问 GitHub。
#
# 用法：
#   bash bash/images_load/preload_deep_swe_agent_offline.sh
# ============================================================

set -euo pipefail

PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
VENDOR_DIR="${PIER_OFFLINE_AGENT_DIR:-${PROJECT_ROOT}/bash/images_load/offline_pier_agent}"
WHEEL_DIR="${VENDOR_DIR}/wheels"
UV_BIN="${VENDOR_DIR}/uv"
HOST_UV="${UV:-}"
PIP_INDEXES=(
    "https://pypi.tuna.tsinghua.edu.cn/simple"
    "https://mirrors.aliyun.com/pypi/simple/"
)
UV_MIRRORS=(
    "https://github.com/astral-sh/uv/releases/download/0.7.13/uv-x86_64-unknown-linux-gnu.tar.gz"
    "https://ghproxy.net/https://github.com/astral-sh/uv/releases/download/0.7.13/uv-x86_64-unknown-linux-gnu.tar.gz"
    "https://mirror.ghproxy.com/https://github.com/astral-sh/uv/releases/download/0.7.13/uv-x86_64-unknown-linux-gnu.tar.gz"
)
PYTHON_VERSIONS=(3.12 3.11)

mkdir -p "${WHEEL_DIR}"

if [[ -z "${HOST_UV}" ]]; then
    if [[ -x /data1/env/uv/bin/uv ]]; then
        HOST_UV=/data1/env/uv/bin/uv
    elif command -v uv >/dev/null 2>&1; then
        HOST_UV="$(command -v uv)"
    fi
fi

echo "============================================================"
echo "1. 准备 uv 二进制 -> ${UV_BIN}"
echo "============================================================"
if [[ -x "${UV_BIN}" ]]; then
    echo "已有 ${UV_BIN} ($("${UV_BIN}" --version 2>/dev/null || true))"
else
    extracted=0
    # 优先从已经 pull 下来的 DeepSWE 镜像里拷，避免 GitHub
    while IFS= read -r image; do
        [[ -z "${image}" ]] && continue
        if docker image inspect "${image}" >/dev/null 2>&1; then
            echo "尝试从 ${image} 提取 uv"
            cid="$(docker create --entrypoint /bin/true "${image}" 2>/dev/null || true)"
            if [[ -n "${cid}" ]]; then
                for src in /root/.local/bin/uv /usr/local/bin/uv /usr/bin/uv; do
                    rm -f "${UV_BIN}"
                    if docker cp "${cid}:${src}" "${UV_BIN}" 2>/dev/null && [[ -f "${UV_BIN}" && ! -L "${UV_BIN}" ]]; then
                        chmod +x "${UV_BIN}"
                        docker rm "${cid}" >/dev/null
                        echo "已从镜像 ${src} 提取 uv: $("${UV_BIN}" --version)"
                        extracted=1
                        break 2
                    fi
                done
                docker rm "${cid}" >/dev/null 2>/dev/null || true
            fi
        fi
    done < <(docker images --format '{{.Repository}}:{{.Tag}}' | grep -E 'swe-bench|mars-base' || true)

    if [[ "${extracted}" -eq 0 && -n "${HOST_UV}" ]]; then
        echo "复制宿主机 uv: ${HOST_UV}"
        cp "${HOST_UV}" "${UV_BIN}"
        chmod +x "${UV_BIN}"
        extracted=1
    fi

    if [[ "${extracted}" -eq 0 ]]; then
        tmp_tar="$(mktemp --suffix=.tar.gz)"
        for url in "${UV_MIRRORS[@]}"; do
            echo "下载 ${url}"
            if curl -L --fail --retry 3 --retry-delay 2 -o "${tmp_tar}" "${url}"; then
                tmp_dir="$(mktemp -d)"
                tar -C "${tmp_dir}" -xzf "${tmp_tar}"
                found="$(find "${tmp_dir}" -type f -name uv -print -quit)"
                if [[ -n "${found}" ]]; then
                    cp "${found}" "${UV_BIN}"
                    chmod +x "${UV_BIN}"
                    extracted=1
                    echo "已下载 uv: $("${UV_BIN}" --version)"
                    rm -rf "${tmp_dir}" "${tmp_tar}"
                    break
                fi
                rm -rf "${tmp_dir}"
            fi
        done
        rm -f "${tmp_tar}"
    fi

    if [[ "${extracted}" -eq 0 ]]; then
        echo "ERROR: 无法获得 uv 二进制"
        exit 1
    fi
fi

echo ""
echo "============================================================"
echo "2. 下载 mini-swe-agent wheel（清华/阿里云）"
echo "============================================================"
HOST_PYTHON="${HOST_PYTHON:-}"
if [[ -z "${HOST_PYTHON}" ]]; then
    if [[ -x /data1/env/conda/envs/syy/bin/python ]]; then
        HOST_PYTHON=/data1/env/conda/envs/syy/bin/python
    else
        HOST_PYTHON="$(command -v python3)"
    fi
fi

download_ok=0
for pyver in "${PYTHON_VERSIONS[@]}"; do
    abi="cp${pyver//./}"
    for index in "${PIP_INDEXES[@]}"; do
        echo "pip download mini-swe-agent  py${pyver}  ${index}"
        if "${HOST_PYTHON}" -m pip download mini-swe-agent \
            -d "${WHEEL_DIR}" \
            --python-version "${pyver}" \
            --implementation cp \
            --abi "${abi}" \
            --platform manylinux2014_x86_64 \
            --only-binary=:all: \
            -i "${index}" \
            --default-timeout=120; then
            download_ok=1
            break
        fi
        echo "该源失败，尝试下一个..."
    done
done

if [[ "${download_ok}" -eq 0 ]]; then
    echo "ERROR: mini-swe-agent wheel 下载失败"
    exit 1
fi

wheel_count="$(find "${WHEEL_DIR}" -name '*.whl' | wc -l | tr -d ' ')"
echo "wheel 数量: ${wheel_count}"
if [[ "${wheel_count}" -lt 5 ]]; then
    echo "ERROR: wheel 过少，拒绝写 .offline_strict"
    exit 1
fi
touch "${VENDOR_DIR}/.offline_strict"

echo ""
echo "============================================================"
echo "3. 可选：LiteLLM 价格表（失败可忽略，构建时只是 warning）"
echo "============================================================"
COST_JSON="${VENDOR_DIR}/model_prices_and_context_window.json"
if [[ ! -s "${COST_JSON}" ]]; then
    curl -L --fail --retry 2 -o "${COST_JSON}" \
        "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json" \
        || rm -f "${COST_JSON}"
fi

echo ""
echo "完成。离线目录: ${VENDOR_DIR}"
ls -lh "${UV_BIN}"
ls "${WHEEL_DIR}" | head
echo ""
echo "评测容器会挂载 evalstone；DeepSWEAdapter 会自动 COPY 该目录进 Pier 构建上下文。"
echo "下一步仍建议先: bash bash/images_load/preload_deep_swe_images.sh"