#!/bin/bash
# ============================================================
# 通用 API 模型评测脚本（原 GLM52_API_TEST1.sh 升级版）
#
# 用法:
#   # 1. 环境变量方式（推荐，避免命令行泄露 key）
#   export EVAL_API_KEY="sk-xxxx"
#   export EVAL_API_URL="https://api.example.com/v1"
#   export EVAL_MODEL="glm-5.2"
#   export EVAL_DATASETS="gsm8k,aime24,arc"
#   bash bash/case/GLM52_API_TEST1.sh
#
#   # 2. 命令行方式
#   bash bash/case/GLM52_API_TEST1.sh \
#       --api-key sk-xxxx \
#       --api-url https://api.example.com/v1 \
#       --model glm-5.2 \
#       --datasets gsm8k,aime24,arc
#
#   # 3. 使用内置 mode（quick / lite / mid / full / official / custom）
#   bash bash/case/GLM52_API_TEST1.sh --mode quick
#
# 修改 datasets：改 EVAL_DATASETS 环境变量或 --datasets 参数即可。
# ============================================================

set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "$ROOT_DIR"

# --------------------------------------------------
# 默认配置（可通过环境变量或命令行覆盖）
# --------------------------------------------------
API_KEY="${EVAL_API_KEY:-}"
API_URL="${EVAL_API_URL:-https://api.vectron.meta-stone.com/v1}"
MODEL="${EVAL_MODEL:-DeepSeek/DeepSeek-V4-Flash}"
DATASETS="${EVAL_DATASETS:-gsm8k,aime24,arc}"
FOLDER_NAME="${EVAL_FOLDER_NAME:-API-Test}"
CONFIG="${EVAL_CONFIG:-config/dpv4-int8_nothinking.yaml}"
BATCH_SIZE="${EVAL_BATCH_SIZE:-4}"
LIMIT="${EVAL_LIMIT:-none}"
SEED="${EVAL_SEED:-42}"
THINKING="${EVAL_THINKING:-false}"
THINKING_BUDGET_TOKENS="${EVAL_THINKING_BUDGET_TOKENS:-}"
MODE="${EVAL_MODE:-custom}"

# --------------------------------------------------
# 解析命令行参数
# --------------------------------------------------
while [[ $# -gt 0 ]]; do
    case "$1" in
        --api-key)       API_KEY="$2"; shift 2 ;;
        --api-url)       API_URL="$2"; shift 2 ;;
        --model)         MODEL="$2"; shift 2 ;;
        --datasets)      DATASETS="$2"; shift 2 ;;
        --folder-name)   FOLDER_NAME="$2"; shift 2 ;;
        --config)        CONFIG="$2"; shift 2 ;;
        --batch-size)    BATCH_SIZE="$2"; shift 2 ;;
        --limit)         LIMIT="$2"; shift 2 ;;
        --seed)          SEED="$2"; shift 2 ;;
        --thinking)      THINKING="true"; shift ;;
        --no-thinking)   THINKING="false"; shift ;;
        --thinking-budget-tokens) THINKING_BUDGET_TOKENS="$2"; shift 2 ;;
        --mode)          MODE="$2"; shift 2 ;;
        -h|--help)
            grep '^# ' "$0" | sed 's/^# //'
            exit 0
            ;;
        *) echo "未知参数: $1"; exit 1 ;;
    esac
done

# --------------------------------------------------
# 模式预设：按需改这里即可扩展常用组合
# --------------------------------------------------
case "$MODE" in
    quick)
        DATASETS="gsm8k,aime24,arc"
        LIMIT="20"
        ;;
    lite)
        DATASETS="gsm8k,aime24,humaneval,arc"
        LIMIT="none"
        ;;
    mid)
        DATASETS="gsm8k,aime24,aime25,bbh,humaneval,arc,simple_qa"
        LIMIT="none"
        ;;
    full)
        DATASETS="bigcodebench,humaneval,live_code_bench,aime24,aime25,aime26,hmmt26,imo_answerbench,gsm8k,competition_math,bbh,drop,gpqa_diamond,mmlu_pro,simple_qa,mmlu,cmmlu,arc,hellaswag,trivia_qa,winogrande,longbench_v2,openai_mrcr,general_fc,bfcl_v3"
        LIMIT="none"
        ;;
    official)
        # 与 DP4-Flash 官方套件对齐
        DATASETS="bigcodebench,humaneval,live_code_bench,aime24,aime25,aime26,hmmt26,imo_answerbench,gsm8k,competition_math,bbh,drop,gpqa_diamond,mmlu_pro,simple_qa,mmlu,cmmlu,arc,hellaswag,trivia_qa,winogrande,longbench_v2,openai_mrcr,tau2_bench,general_fc,bfcl_v3"
        LIMIT="none"
        ;;
    custom)
        # 使用 DATASETS 环境变量或命令行传入的值
        ;;
    *)
        echo "未知模式: $MODE"
        echo "可用模式: quick | lite | mid | full | official | custom"
        exit 1
        ;;
esac

# --------------------------------------------------
# API key 校验
# --------------------------------------------------
if [[ -z "$API_KEY" ]]; then
    echo "ERROR: 请设置 EVAL_API_KEY 环境变量或传入 --api-key"
    exit 1
fi

export EVALSCOPE_API_KEY="$API_KEY"
export OPENAI_API_KEY="$API_KEY"

HEALTH=$(curl -s -o /dev/null -w "%{http_code}" \
    -H "Authorization: Bearer ${API_KEY}" \
    "${API_URL}/models")
if [[ "$HEALTH" != "200" ]]; then
    echo "ERROR: API key 校验失败，${API_URL}/models 返回 HTTP $HEALTH"
    exit 1
fi
echo "API key 校验通过 (${API_URL})"

# --------------------------------------------------
# 组装 run.py 参数
# --------------------------------------------------
ARGS=(
    --model "$MODEL"
    --api-url "$API_URL"
    --dataset-dir "$ROOT_DIR"
    --output-dir "$ROOT_DIR/output"
    --folder-name "$FOLDER_NAME"
    --config "$CONFIG"
    --batch-size "$BATCH_SIZE"
    --seed "$SEED"
    --limit "$LIMIT"
    --datasets "$DATASETS"
)

if [[ "$THINKING" == "true" ]]; then
    ARGS+=(--thinking)
fi

if [[ -n "$THINKING_BUDGET_TOKENS" ]]; then
    ARGS+=(--thinking-budget-tokens "$THINKING_BUDGET_TOKENS")
fi

echo "============================================================"
echo "API 评测启动"
echo "Mode:         $MODE"
echo "Model:        $MODEL"
echo "API URL:      $API_URL"
echo "Datasets:     $DATASETS"
echo "Folder:       $FOLDER_NAME"
echo "Config:       $CONFIG"
echo "Batch size:   $BATCH_SIZE"
echo "Limit:        $LIMIT"
echo "Thinking:     $THINKING"
if [[ -n "$THINKING_BUDGET_TOKENS" ]]; then
    echo "Thinking budget_tokens: $THINKING_BUDGET_TOKENS"
fi
echo "============================================================"

python bash/run.py "${ARGS[@]}"