# cd /data1/p800_deploy_scripts/DeepSeek-V4-Flash-xsglang-20260511-deploy
# docker rm -f /deepseek-v4-flash-xsglang-0511
# bash /data1/p800_deploy_scripts/DeepSeek-V4-Flash-xsglang-20260511-deploy/scripts/docker_run.sh
# bash scripts/run_server.sh

import os
import yaml
from copy import deepcopy

from evalscope import run_task, TaskConfig
from evalscope.api.agent import NativeAgentConfig

# ============================================================
# 1. 用户可配置项
# ============================================================
CONFIG_PATH = '/data1/benchmark/evalscope/bash/generation_config_nothinking.yaml'

# 显式开关：True 表示开启模型的 thinking 能力，False 表示关闭
ENABLE_THINKING = False

# 选择本次实际要跑的数据集
DATASETS = [
    # 'aime26',
    # 'terminal_bench_v2',
    # 'gpqa_diamond',
    # 'hle',
    # 'aime24',
    # 'aime25',
    # 'mmlu_pro',
    # 'simple_qa',
    # 'arc',
    # 'bbh',
    # 'browsecomp',
    # 'live_code_bench',
    # 'swe_bench_pro',
    # 'hmmt26',
    # 'imo_answerbench',
    # 'super_gpqa',
    # 'drop',
    # 'hellaswag',
    # 'mmlu',
    # 'openai_mrcr',
    # 'mcp_atlas',
    # 'swe_bench_multilingual_agentic',
    # 'swe_bench_verified',
    'bigcodebench',
    # 'humaneval',
    # 'gsm8k',
    # 'competition_math',
    # 'cmmlu',
    # 'trivia_qa',
    # 'winogrande',
    # 'longbench_v2',
    # 'tau2_bench',
]

# 公共配置（所有数据集共享）
COMMON_CONFIG = {
    'seed': 42,
    'collect_perf': True,
    'no_timestamp': True,
    'model': 'DeepSeek-V4-Flash-Int8',
    'api_url': 'http://localhost:30000/v1',
    # 'api_key': '',
    'eval_type': 'openai_api',
    'dataset_dir': '/data1/benchmark/bash/datasets',
    'judge_model_args': {
        'model_id': 'deepseek-v4-pro',
        'api_url': 'https://api.deepseek.com/v1',
        'api_key': 'sk-9ed86ef546ca47e3afa7c3b014dea268',
        'eval_type': 'openai_api',
        'generation_config': {
            'temperature': 0.0,
            'max_tokens': 1024 * 10,
        },
    },
}

# 数学类任务集合
MATH_DATASETS = {
    'aime24', 'aime25', 'aime26', 'hmmt26',
    'gsm8k', 'competition_math', 'imo_answerbench',
}

# 数学类任务 prompt 模板
MATH_PROMPT_TEMPLATE = (
    "{question}\n"
    "Please reason step by step, and put your final answer within \\boxed{{}}."
)

batch_size_list = [16]


# ============================================================
# 2. 加载 YAML 配置
# ============================================================
with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
    DATASET_CONFIGS = yaml.safe_load(f)


# ============================================================
# 3. 辅助函数
# ============================================================
def configure_thinking(generation_config: dict, enable: bool) -> dict:
    """显式开启或关闭 generation_config 中的 thinking 能力。"""
    if enable:
        generation_config['thinking'] = {'type': 'enabled'}
    else:
        generation_config.pop('thinking', None)
    return generation_config


def build_agent_config(agent_cfg: dict) -> NativeAgentConfig:
    """从 YAML 的 agent_config 构造 NativeAgentConfig。

    YAML 里可能包含 strategy 未声明的字段（如 per_step_tokens），
    这些字段会被放到 kwargs 里，透传给 Agent Strategy。
    """
    agent_cfg = deepcopy(agent_cfg or {})

    # NativeAgentConfig 已知的顶层字段
    known_fields = {'mode', 'strategy', 'tools', 'max_steps', 'mcp_servers', 'environment', 'environment_extra'}
    kwargs = agent_cfg.pop('kwargs', {})

    # 把未知字段收敛进 kwargs
    for key in list(agent_cfg.keys()):
        if key not in known_fields:
            kwargs[key] = agent_cfg.pop(key)

    if kwargs:
        agent_cfg['kwargs'] = kwargs

    return NativeAgentConfig(**agent_cfg)


def build_task_config(dataset_name: str, ds_cfg: dict, batch_size: int, enable_thinking: bool) -> TaskConfig:
    """根据数据集配置构造单个 TaskConfig。"""
    work_dir = f'/data1/benchmark/temp/output_{batch_size}/{dataset_name}'

    # 构造 generation_config，并显式处理 thinking 开关
    generation_config = configure_thinking(deepcopy(ds_cfg['generation_config']), enable_thinking)

    # 构造 dataset_args：仅对应当前数据集
    dataset_args = deepcopy(ds_cfg.get('dataset_args', {}))

    # 数学类任务覆盖 prompt_template
    if dataset_name in MATH_DATASETS:
        dataset_args['prompt_template'] = MATH_PROMPT_TEMPLATE

    dataset_args_dict = {dataset_name: dataset_args}

    # 构造 agent_config（如果 YAML 里配了的话）
    agent_config = None
    if 'agent_config' in ds_cfg:
        agent_config = build_agent_config(ds_cfg['agent_config'])

    return TaskConfig(
        **COMMON_CONFIG,
        work_dir=work_dir,
        use_cache=work_dir,
        datasets=[dataset_name],
        generation_config=generation_config,
        dataset_args=dataset_args_dict,
        agent_config=agent_config,
        eval_batch_size=batch_size,
    )


# ============================================================
# 4. 主循环
# ============================================================
if __name__ == '__main__':
    for batch_size in batch_size_list:
        task_cfg_list = []

        for dataset_name in DATASETS:
            if dataset_name not in DATASET_CONFIGS:
                raise ValueError(
                    f"Dataset '{dataset_name}' 不在 {CONFIG_PATH} 中，"
                    f"请先在 YAML 里补充它的 generation_config。"
                )

            ds_cfg = DATASET_CONFIGS[dataset_name]
            task_cfg = build_task_config(dataset_name, ds_cfg, batch_size, ENABLE_THINKING)
            task_cfg_list.append(task_cfg)

        # run_task 支持 list，会串行执行每个数据集的评测
        run_task(task_cfg_list)