#!/usr/bin/env python3
"""
Unified benchmark runner for EvalScope.

A single entry point for lite / mid / full / group1 / group2 / group3 evaluations.
All tunable parameters can be controlled via command-line arguments.

Examples:
    # Full evaluation (all benchmarks, multi-run for stability)
    python bash/run.py \
        --model DeepSeek-V4-Flash-Int8 \
        --api-url http://localhost:30000/v1 \
        --dataset-dir /data1/sora/evalscope \
        --output-dir /data1/sora/evalscope/output \
        --suite full \
        --limit none

    # Lite smoke test (~5h with full samples)
    python bash/run.py --suite lite --limit none

    # Run only selected benchmarks
    python bash/run.py --datasets aime24,gsm8k,arc --limit 20

    # Custom judge model
    python bash/run.py \
        --judge-model deepseek-v4-pro \
        --judge-api-url https://api.deepseek.com/v1 \
        --judge-api-key sk-xxx
"""

import argparse
import json
import os
import subprocess
import sys
import threading
import time
from copy import deepcopy
from pathlib import Path

# LiteLLM 默认会联网拉取 GitHub 上的 model cost map，内网/无网环境会超时。
# 强制使用本地备份，避免运行时被卡住。
os.environ.setdefault('LITELLM_LOCAL_MODEL_COST_MAP', 'True')

import yaml
sys.path.insert(0, str(Path(__file__).parent.parent / "evalscope"))

from evalscope import run_task, TaskConfig
from evalscope.api.agent import NativeAgentConfig
from evalscope.config import SandboxTaskConfig
from evalscope.utils.logger import get_logger

logger = get_logger()

SCRIPT_DIR = Path(__file__).parent.resolve()
PROJECT_ROOT = SCRIPT_DIR.parent

# Make collect_results importable
sys.path.insert(0, str(SCRIPT_DIR))
import collect_results as collect_results_module
import perf_backup as perf_backup_module

# Make fingerprint helpers importable（失败报告写入用）
sys.path.insert(0, str(SCRIPT_DIR / "fingerprint"))
import common as fingerprint_common

# ============================================================
# Default configuration (override via CLI)
# ============================================================

DEFAULT_MODEL = 'DeepSeek-V4-Flash-Int8'
DEFAULT_API_URL = 'http://localhost:30000/v1'
DEFAULT_API_KEY = os.environ.get('EVALSCOPE_API_KEY', 'EMPTY')
# 数据集缓存根：evalscope 会在其下找 datasets/<名字>-<hash>。
# 镜像内通过 EVALSTONE_DATASET_DIR 指到挂载卷，宿主机目录直接命中已有缓存。
DEFAULT_DATASET_DIR = os.environ.get('EVALSTONE_DATASET_DIR', str(PROJECT_ROOT))
DEFAULT_OUTPUT_DIR = str(PROJECT_ROOT / 'output')
DEFAULT_CONFIG = str(PROJECT_ROOT / 'config' / 'dpv4-int8_nothinking.yaml')
DEFAULT_TOKENIZER_PATH = '/data1/models/DeepSeek-V4-Flash-INT8'
DEFAULT_LIMIT = None
DEFAULT_SEED = 42
DEFAULT_BATCH_SIZE = 4
DEFAULT_ENABLE_THINKING = False

# 新 benchmark 没在 YAML 里配时的默认生成参数
DEFAULT_GENERATION_CONFIG = {
    'temperature': 0.0,
    'top_p': 1.0,
    'stream': True,
    'max_tokens': 32768,
}

DEFAULT_JUDGE_MODEL = 'DeepSeek/DeepSeek-V4-Pro'
DEFAULT_JUDGE_API_URL = 'https://api.vectron.meta-stone.com/v1'
DEFAULT_JUDGE_API_KEY = 'sk-dbd8a665f7634081b87ec409c7636500'
DEFAULT_JUDGE_MAX_TOKENS = 10240

# 长文本 middle-truncation 上限（token 数）。当前默认 128k。
DEFAULT_TRUNCATION_TOKENS = 32768 * 4

# ============================================================
# Benchmark suites
# ============================================================

# 多次采样配置：总样本数控制在 ~400-500
MULTI_RUN_CONFIG = {
    'aime24': 12,
    'aime25': 12,
    'aime26': 12,
    'hmmt26': 12,
    'live_code_bench': 5,
    'imo_answerbench': 4,
    'humaneval': 3,
    'gpqa_diamond': 2,
}

# 能力域完整列表
ALL_MULTI_RUN = [
    'humaneval', 'live_code_bench',
    'aime24', 'aime25', 'aime26', 'hmmt26',
    'imo_answerbench', 'gpqa_diamond',
]
ALL_SINGLE_RUN = [
    'bigcodebench', 'bfcl_v3', 'competition_math', 'gsm8k', 'hle', 'super_gpqa',
    'arc', 'bbh', 'cmmlu', 'drop', 'hellaswag', 'mmlu', 'mmlu_pro',
    'simple_qa', 'trivia_qa', 'winogrande',
    'openai_mrcr', 'longbench_v2',
]
ALL_AGENT = ['tau2_bench', 'general_fc']
# Kimi-K3 可直接跑且备注为空的集 + HLE-Full（dataset id 为 hle）。
# tau3 子集 banking_knowledge 在 config/kimi-k3.yaml 里配置。
K3_SINGLE = [
    # Coding
    'deep_swe',
    'terminal_bench_v2_1',
    'scicode',
    # Agentic
    'browsecomp',
    'deepsearchqa',
    'job_bench',
    'officeqa',
    'tau3_bench',
    'researchrubrics',
    # Reasoning
    'gpqa_diamond',
    'aa_lcr',
    'hle',
    # Vision
    'mmmu_pro',
    'charxiv',
    'math_vision',
    'omni_doc_bench',
]

# ============================================================
# Fingerprint / model-identity benchmarks
# ============================================================
# 这三个 benchmark 不经过 EvalScope 数据集管线：由 bash/fingerprint/ 下的
# 执行器直接探测 --api-url 端点，并产出与 EvalScope 同构的
# output/<folder>/<benchmark>/seed_<seed>/reports/<benchmark>.json（含 score），
# collect_results.py 可像普通 benchmark 一样汇总。
ALL_FINGERPRINT = ['llmmap', 'llm_verify', 'llm_fingerprint_detector']
# fp_fusion 为融合型指纹基准(strict), 只通过 --datasets 单跑, 不进 --suite fingerprint
FUSION_BENCHMARKS = ['fp_fusion']
FINGERPRINT_SCRIPTS = {
    'llmmap': SCRIPT_DIR / 'fingerprint' / 'run_llmmap.py',
    'llm_verify': SCRIPT_DIR / 'fingerprint' / 'run_llm_verify.py',
    'llm_fingerprint_detector': SCRIPT_DIR / 'fingerprint' / 'run_llm_detector.py',
    'fp_fusion': SCRIPT_DIR / 'fingerprint' / 'fp_fusion' / 'run_fp_fusion.py',
}
# 三个指纹工具仓库（LLMmap/llm-verify/llm-fingerprint-detector）优先用
# evalstone 内置的 bash/fingerprint/tools（随仓库走、自包含），可用
BUILTIN_TOOLS_ROOT = SCRIPT_DIR / 'fingerprint' / 'tools'
DEFAULT_TOOLS_ROOT = (
    os.environ.get('FP_TOOLS_ROOT')
    or (str(BUILTIN_TOOLS_ROOT) if (BUILTIN_TOOLS_ROOT / 'LLMmap').is_dir() else None)
    or '/data1/xii'
)
# 单个指纹 benchmark 的整体子进程超时（秒）。verify 的 32 条探测较慢，给足余量。
FP_OVERALL_TIMEOUT = 7200

# 分组基于 CSV 单次时间 + multi-run 后的 wall time 平衡：
#   Group1: ~61h | Group2: ~62h | Group3: ~55h
SUITES = {
    'full': {
        'multi': ALL_MULTI_RUN,
        'single': ALL_SINGLE_RUN,
        'agent': ALL_AGENT,
    },
    'lite': {
        'multi': ['aime24', 'humaneval'],
        'single': ['gsm8k', 'arc', 'longbench_v2'],
        'agent': ['general_fc'],
    },
    "k3": {"multi": [], "single": K3_SINGLE, "agent": []},
    "tencent_hunyuan": {"multi": [], "single": [], "agent": []},
    'mid': {
        'multi': ['aime24', 'humaneval'],
        'single': [
            'live_code_bench', 'bigcodebench', 'competition_math', 'gsm8k',
            'gpqa_diamond', 'mmlu_pro', 'simple_qa', 'longbench_v2', 'openai_mrcr',
        ],
        'agent': ['general_fc', 'tau2_bench'],
    },
    # 多机组分组，基于 CSV 实测完整时间（已含 multi-run）平衡：
    #   Group1: ~22.7h | Group2: ~25.6h | Group3: ~27.0h | 合计 ~75.2h
    'group1': {
        'multi': ['live_code_bench', 'aime24', 'aime25', 'aime26', 'hmmt26', 'imo_answerbench', 'humaneval'],
        'single': ['bigcodebench', 'competition_math', 'gsm8k', 'drop', 'arc', 'hellaswag', 'winogrande'],
        'agent': [],
    },
    'group2': {
        'multi': [],
        'single': ['hle', 'mmlu_pro', 'trivia_qa'],
        'agent': [],
    },
    'group3': {
        'multi': ['gpqa_diamond'],
        'single': ['openai_mrcr', 'longbench_v2', 'bfcl_v3', 'mmlu', 'cmmlu', 'bbh', 'simple_qa'],
        'agent': ['tau2_bench', 'general_fc'],
    },
    'official': {
        'multi': ['aime25', 'aime26', 'live_code_bench'],
        'single': [

            'hle',
             'mmlu_pro',
              'gpqa_diamond',
               'longbench_v2',
            'swe_bench_verified',
        ],
        'agent': [
            'tau2_bench'
        ],
    },
    # 模型指纹/安全套件：LLMmap 身份识别 + LLM Verify 欺诈检测 + 单 token 分布验证
    'fingerprint': {
        'multi': [],
        'single': [],
        'agent': [],
        'fingerprint': ALL_FINGERPRINT,
    },
}

# ============================================================
# Fixed configuration
# ============================================================

MATH_DATASETS = {
    'aime24', 'aime25', 'aime26', 'hmmt26',
    'gsm8k', 'competition_math', 'imo_answerbench',
}

MATH_PROMPT_TEMPLATE = (
    "{question}\n"
    "Please reason step by step, and put your final answer within \\boxed{{}}."
)

SANDBOX_DATASETS = {'humaneval', 'bigcodebench', 'scicode'}
SANDBOX_CONFIGS = {
    'bigcodebench': {
        'image': 'bigcodebench-sandbox:latest',
        'working_dir': '/tmp',
        'tools_config': {
            'shell_executor': {},
            'python_executor': {}
        }
    },
    'humaneval': {
        'image': 'python:3.11-slim',
        'tools_config': {
            'shell_executor': {},
            'python_executor': {}
        }
    },
    'scicode': {
        'image': 'scicode-benchmark:latest',
        'working_dir': '/workspace',
        'tools_config': {
            'shell_executor': {},
            'python_executor': {}
        }
    },
}


# ============================================================
# CLI parser
# ============================================================

def build_parser():
    parser = argparse.ArgumentParser(
        description='Unified EvalScope benchmark runner',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='Suites: full, lite, mid, group1, group2, group3, official',
    )

    # Model / API
    parser.add_argument('--model', default=DEFAULT_MODEL,
                        help='Served model name (default: %(default)s)')
    parser.add_argument('--api-url', default=DEFAULT_API_URL,
                        help='OpenAI-compatible API URL (default: %(default)s)')
    parser.add_argument('--api-key', default=DEFAULT_API_KEY,
                        help='API key for --api-url (default: env EVALSCOPE_API_KEY or EMPTY)')

    # Paths
    parser.add_argument('--dataset-dir', default=DEFAULT_DATASET_DIR,
                        help='Parent directory containing datasets/ subdir (default: %(default)s)')
    parser.add_argument('--output-dir', default=DEFAULT_OUTPUT_DIR,
                        help='Output root directory (default: %(default)s)')
    parser.add_argument('--results-dir', default=None,
                        help='Directory for summary CSV/Excel (default: <repo>/results)')
    parser.add_argument('--folder-name', default=None,
                        help='Top-level output folder name; defaults to --model, or --model_THINKING when --thinking is enabled')
    parser.add_argument('--config', default=DEFAULT_CONFIG,
                        help='YAML config path (default: %(default)s)')
    parser.add_argument('--tokenizer-path', default=DEFAULT_TOKENIZER_PATH,
                        help='Local tokenizer path for middle-truncation (default: %(default)s)')

    # Run control
    parser.add_argument('--suite', default='full', choices=list(SUITES.keys()),
                        help='Benchmark suite to run (default: %(default)s)')
    parser.add_argument('--datasets', '--benchmarks', dest='datasets', default=None,
                        help='Override suite with comma-separated benchmark names, e.g. aime24,gsm8k')
    parser.add_argument('--exclude', default=None,
                        help='Comma-separated benchmarks to exclude from the chosen suite')
    parser.add_argument('--limit', default=None,
                        help='Max samples per benchmark; "none"/"all" for no limit (default: none)')
    parser.add_argument('--seed', type=int, default=DEFAULT_SEED,
                        help='Random seed (default: %(default)s)')
    parser.add_argument('--batch-size', type=int, default=DEFAULT_BATCH_SIZE,
                        help='Evaluation batch size (default: %(default)s)')
    parser.add_argument('--parallel-runs', type=int, default=1,
                        help='Number of multi-run seeds to execute in parallel '
                             '(default: %(default)s). Each run still uses --batch-size '
                             'concurrent requests, so total API concurrency is '
                             'parallel-runs * batch-size.')
    parser.add_argument('--parallel-benchmarks', type=int, default=1,
                        help='Number of different benchmarks to execute in parallel '
                             '(default: %(default)s). When >1, each benchmark runs in its '
                             'own thread; completed runs immediately free a slot for the '
                             'next benchmark.')

    # Decoding / thinking
    parser.add_argument('--thinking', action='store_true', default=None,
                        help='Enable thinking mode (sglang chat_template_kwargs.thinking=True)')
    parser.add_argument('--no-thinking', dest='thinking', action='store_false',
                        help='Disable thinking mode (default)')
    parser.add_argument('--thinking-max-tokens-scale', type=float, default=1.0,
                        help='Scale max_tokens by this factor when --thinking is enabled (default: %(default)s)')
    parser.add_argument('--thinking-budget-tokens', type=int, default=None,
                        help='Add extra_body.thinking={"type": "enabled", "budget_tokens": N} to control thinking budget. '
                             'Set to 0 to attempt disabling thinking via API.')
    parser.add_argument('--max-tokens-add', type=int, default=0,
                        help='Add this many tokens to every benchmark max_tokens (applied after scale)')

    # Judge model
    parser.add_argument('--judge-model', default=DEFAULT_JUDGE_MODEL,
                        help='Judge model name (default: %(default)s)')
    parser.add_argument('--judge-api-url', default=DEFAULT_JUDGE_API_URL,
                        help='Judge model API URL (default: %(default)s)')
    parser.add_argument('--judge-api-key', default=DEFAULT_JUDGE_API_KEY,
                        help='Judge model API key')
    parser.add_argument('--judge-max-tokens', type=int, default=DEFAULT_JUDGE_MAX_TOKENS,
                        help='Judge model max_tokens (default: %(default)s)')

    # Truncation
    parser.add_argument('--truncation-tokens', type=int, default=DEFAULT_TRUNCATION_TOKENS,
                        help='Middle-truncation token budget for long-context benchmarks (default: %(default)s)')

    # Fingerprint benchmarks (llmmap / llm_verify / llm_fingerprint_detector)
    parser.add_argument('--tools-root', default=DEFAULT_TOOLS_ROOT,
                        help='Root dir containing the three fingerprint tool repos '
                             '(LLMmap/, llm-verify/, llm-fingerprint-detector/) (default: %(default)s)')
    parser.add_argument('--llmmap-python',
                        default=os.environ.get('LLMMAP_PYTHON',
                                               '/root/miniconda3/envs/llmmap/bin/python'),
                        help='Python interpreter with torch/transformers for the LLMmap runner')
    parser.add_argument('--verify-python',
                        default=os.environ.get('LLMVERIFY_PYTHON',
                                               '/root/miniconda3/envs/llmverify/bin/python'),
                        help='Python interpreter with fastapi/httpx for the LLM Verify runner')
    parser.add_argument('--detector-node', default=os.environ.get('DETECTOR_NODE', 'node'),
                        help='Node executable for the llm-fingerprint-detector runner')
    parser.add_argument('--detector-reference', default=None,
                        help='Optional same-protocol reference fingerprint JSON for the detector; '
                             'omit to run in self-consistency mode')
    parser.add_argument('--detector-preset', default='standard',
                        choices=['quick', 'standard', 'strict'],
                        help='Sampling preset for the detector benchmark (default: %(default)s)')
    parser.add_argument('--expected-model', default=None,
                        help='Ground-truth model identity; when set, the llmmap score becomes a '
                             'strict Top-1 identity match flag instead of a distance confidence')
    parser.add_argument('--fingerprint-timeout', type=int, default=120,
                        help='Per-request timeout (seconds) passed to fingerprint runners (default: %(default)s)')

    # Result collection
    parser.add_argument('--no-summary', dest='write_summary', action='store_false',
                        help='Skip writing summary Excel/CSV after each benchmark')

    return parser


# ============================================================
# Middle-truncation helpers
# ============================================================

_TOKENIZER = None


def get_tokenizer(tokenizer_path: str):
    global _TOKENIZER
    if _TOKENIZER is not None:
        return _TOKENIZER

    from transformers import AutoTokenizer
    import os

    # Try local path first
    if tokenizer_path and os.path.exists(tokenizer_path):
        try:
            _TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
            logger.info(f'Loaded tokenizer from local path: {tokenizer_path}')
            return _TOKENIZER
        except Exception as e:
            logger.warning(f'Failed to load tokenizer from local path {tokenizer_path}: {e}')

    # If a local path was explicitly provided but does not exist, warn clearly
    if tokenizer_path and tokenizer_path != '/data1/models/DeepSeek-V4-Flash-INT8':
        logger.warning(f'Local tokenizer path does not exist: {tokenizer_path}. '
                       f'Will try to download from model hub.')

    # Fallback: try model hub. Respect ModelScope hub if configured.
    fallback_model = 'deepseek-ai/DeepSeek-V4-Flash'
    hub_source = 'ModelScope' if os.environ.get('USE_MODELSCOPE_HUB') == '1' else 'HuggingFace'
    logger.info(f'Trying to load tokenizer from {hub_source}: {fallback_model}')
    try:
        _TOKENIZER = AutoTokenizer.from_pretrained(fallback_model, trust_remote_code=True)
        logger.info(f'Loaded tokenizer from {hub_source}: {fallback_model}')
        return _TOKENIZER
    except Exception as e:
        raise RuntimeError(
            f'Failed to load tokenizer.\n'
            f'  Local path: {tokenizer_path} (exists: {os.path.exists(tokenizer_path) if tokenizer_path else False})\n'
            f'  {hub_source} fallback: {fallback_model}\n'
            f'  Error: {e}\n'
            f'  Hint: mount the tokenizer to the container and set --tokenizer-path correctly, '
            f'or set USE_MODELSCOPE_HUB=1 and HF_ENDPOINT=https://hf-mirror.com if downloading.'
        ) from e


def truncate_middle(text: str, max_tokens: int, tokenizer_path: str) -> str:
    if max_tokens <= 0:
        return text
    tokenizer = get_tokenizer(tokenizer_path)
    token_ids = tokenizer.encode(text, add_special_tokens=False)
    if len(token_ids) <= max_tokens:
        return text
    keep_head = max_tokens // 2
    keep_tail = max_tokens - keep_head
    truncated_ids = token_ids[:keep_head] + token_ids[-keep_tail:]
    return tokenizer.decode(truncated_ids, skip_special_tokens=True)


def _patch_adapters_for_truncation(tokenizer_path: str, truncation_tokens: int):
    from evalscope.benchmarks.aa_lcr.aa_lcr_adapter import AALCRAdapter
    from evalscope.benchmarks.longbench_v2.longbench_v2_adapter import LongBenchV2Adapter
    from evalscope.benchmarks.openai_mrcr.openai_mrcr_adapter import OpenAIMRCRAdapter

    _orig_aa_lcr_sample = AALCRAdapter.record_to_sample

    def _patched_aa_lcr_sample(self, record):
        sample = _orig_aa_lcr_sample(self, record)
        if not sample.input:
            return sample
        msg = sample.input[0]
        content = getattr(msg, 'content', None)
        if isinstance(content, str) and content:
            msg.content = truncate_middle(content, truncation_tokens, tokenizer_path)
        return sample

    AALCRAdapter.record_to_sample = _patched_aa_lcr_sample

    _orig_longbench_format = LongBenchV2Adapter.format_prompt_template

    def _patched_longbench_format(self, sample):
        if sample.metadata and 'context' in sample.metadata:
            sample.metadata['context'] = truncate_middle(sample.metadata['context'], truncation_tokens, tokenizer_path)
        return _orig_longbench_format(self, sample)

    LongBenchV2Adapter.format_prompt_template = _patched_longbench_format

    _orig_mrcr_record = OpenAIMRCRAdapter.record_to_sample

    def _patched_mrcr_record(self, record):
        per_msg_max_tok = 8192
        if 'prompt' in record:
            try:
                prompt_data = json.loads(record['prompt'])
                if not isinstance(prompt_data, list) or len(prompt_data) == 0:
                    return _orig_mrcr_record(self, record)

                tokenizer = get_tokenizer(tokenizer_path)
                total_tok = sum(
                    len(tokenizer.encode(msg.get('content', '') if isinstance(msg, dict) else '', add_special_tokens=False))
                    for msg in prompt_data
                )
                if total_tok <= truncation_tokens:
                    return _orig_mrcr_record(self, record)

                desired_idx = record.get('desired_msg_index', 0)
                if not isinstance(desired_idx, int) or desired_idx < 0 or desired_idx >= len(prompt_data):
                    desired_idx = 0

                n = len(prompt_data)
                keep = set()
                keep.update(range(min(2, n)))
                keep.update(range(max(0, n - 2), n))
                window = 2
                keep.update(range(max(0, desired_idx - window), min(n, desired_idx + window + 1)))
                keep = sorted(keep)

                new_prompt = []
                for idx in keep:
                    msg = prompt_data[idx]
                    if isinstance(msg, dict):
                        msg = dict(msg)
                        content = msg.get('content', '')
                        if len(tokenizer.encode(content, add_special_tokens=False)) > per_msg_max_tok:
                            msg['content'] = truncate_middle(content, per_msg_max_tok, tokenizer_path)
                    new_prompt.append(msg)

                record = dict(record)
                record['prompt'] = json.dumps(new_prompt)
            except (json.JSONDecodeError, TypeError):
                pass
        return _orig_mrcr_record(self, record)

    OpenAIMRCRAdapter.record_to_sample = _patched_mrcr_record


# ============================================================
# Helpers
# ============================================================

def load_dataset_configs(config_path: str):
    if not Path(config_path).exists():
        raise FileNotFoundError(f'Config file not found: {config_path}')
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


def configure_thinking(generation_config: dict, enable: bool, thinking_budget_tokens: int = None) -> dict:
    extra_body = generation_config.get('extra_body', {})
    chat_template_kwargs = extra_body.get('chat_template_kwargs', {})
    if enable:
        chat_template_kwargs['thinking'] = True
    else:
        chat_template_kwargs.pop('thinking', None)
    if chat_template_kwargs:
        extra_body['chat_template_kwargs'] = chat_template_kwargs

    # 支持 API 级别的 thinking budget 控制（如 {"type": "enabled", "budget_tokens": 0}）
    if thinking_budget_tokens is not None:
        extra_body['thinking'] = {'type': 'enabled', 'budget_tokens': thinking_budget_tokens}
    else:
        extra_body.pop('thinking', None)

    if extra_body:
        generation_config['extra_body'] = extra_body
    return generation_config


def build_agent_config(agent_cfg: dict) -> NativeAgentConfig:
    agent_cfg = deepcopy(agent_cfg or {})
    known_fields = {'mode', 'strategy', 'tools', 'max_steps', 'mcp_servers', 'environment', 'environment_extra'}
    kwargs = agent_cfg.pop('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,
    seed: int,
    limit,
    output_dir: str,
    model: str,
    api_url: str,
    api_key: str,
    dataset_dir: str,
    judge_model_args: dict,
    run_idx: int = 0,
    thinking_max_tokens_scale: float = 1.0,
    max_tokens_add: int = 0,
    thinking_budget_tokens: int = None,
) -> TaskConfig:
    if run_idx > 0:
        work_dir = Path(output_dir) / dataset_name / f'seed_{seed}_run_{run_idx}'
    else:
        work_dir = Path(output_dir) / dataset_name / f'seed_{seed}'
    work_dir.mkdir(parents=True, exist_ok=True)
    # 升级后的 EvalScope：只要设置了 use_cache，就会校验 evaluation identity。
    # 空目录（无 task_config.yaml）会被当成 previous=unknown 直接报错；
    # 仅有 task_config、无 predictions 的失败/改配残留也会因 fingerprint 变化秒挂。
    # 仅在已有完整缓存快照（config + 至少一个 prediction）时才启用 resume。
    cache_snapshot = work_dir / 'configs' / 'task_config.yaml'
    pred_dir = work_dir / 'predictions'
    has_predictions = pred_dir.is_dir() and any(pred_dir.rglob('*'))
    if cache_snapshot.is_file() and has_predictions:
        use_cache = str(work_dir)
    else:
        if cache_snapshot.is_file() and not has_predictions:
            print(
                f'   [cache] skip resume for {dataset_name}: '
                'stale/incomplete cache (task_config without predictions)'
            )
        use_cache = None
    work_dir = str(work_dir)

    generation_config = configure_thinking(
        deepcopy(ds_cfg['generation_config']), enable_thinking, thinking_budget_tokens
    )
    if thinking_budget_tokens is not None:
        print(f'   [thinking] budget_tokens: {thinking_budget_tokens}')
    if enable_thinking and thinking_max_tokens_scale != 1.0:
        original_max_tokens = generation_config.get('max_tokens', 32768)
        scaled = int(original_max_tokens * thinking_max_tokens_scale)
        generation_config['max_tokens'] = scaled
        print(f'   [thinking] max_tokens scaled: {original_max_tokens} -> {scaled}')
    if max_tokens_add > 0:
        original_max_tokens = generation_config.get('max_tokens', 32768)
        new_max_tokens = original_max_tokens + max_tokens_add
        generation_config['max_tokens'] = new_max_tokens
        print(f'   max_tokens add: {original_max_tokens} -> {new_max_tokens}')
    dataset_args = deepcopy(ds_cfg.get('dataset_args', {}))
    dataset_args.setdefault('shuffle', True)

    if dataset_name in MATH_DATASETS:
        dataset_args['prompt_template'] = MATH_PROMPT_TEMPLATE

    dataset_args_dict = {dataset_name: dataset_args}

    agent_config = None
    if 'agent_config' in ds_cfg:
        agent_config = build_agent_config(ds_cfg['agent_config'])

    return TaskConfig(
        model=model,
        api_url=api_url,
        api_key=api_key,
        eval_type='openai_api',
        dataset_dir=dataset_dir,
        judge_model_args=judge_model_args,
        seed=seed,
        limit=limit,
        collect_perf=True,
        no_timestamp=True,
        work_dir=work_dir,
        use_cache=use_cache,
        datasets=[dataset_name],
        generation_config=generation_config,
        dataset_args=dataset_args_dict,
        agent_config=agent_config,
        eval_batch_size=batch_size,
        sandbox=SandboxTaskConfig(
            enabled=True,
            engine='docker',
            default_config=SANDBOX_CONFIGS.get(dataset_name, {
                'image': 'python:3.11-slim',
                'tools_config': {
                    'shell_executor': {},
                    'python_executor': {}
                }
            })
        ) if dataset_name in SANDBOX_DATASETS else None,
    )


# ============================================================
# Main
# ============================================================

# Overridden in main() from --results-dir
RESULTS_DIR = PROJECT_ROOT / 'results'


def write_summary(output_dir: str, model_name: str, folder_name: str,
                  benchmark_names: list = None):
    """Re-aggregate results for the given benchmarks (or all on disk if None)."""
    try:
        excel_output_dir = Path(RESULTS_DIR)
        excel_output_dir.mkdir(parents=True, exist_ok=True)
        if benchmark_names is not None:
            collect_results_module.eval_benchmark(
                benchmark_names, Path(output_dir), model_name, folder_name,
                excel_output_dir=excel_output_dir,
            )
        else:
            collect_results_module.collect_all(
                Path(output_dir), model_name, folder_name,
                excel_output_dir=excel_output_dir,
            )
    except Exception as e:
        print(f'WARNING: failed to write summary: {e}')


def run_and_summarize(task_cfg, write_summary_flag: bool, output_dir: str,
                       model_name: str, folder_name: str, benchmark_names: list,
                       summary_lock: threading.Lock = None):
    """Run a benchmark task and optionally refresh the summary table.

    ``benchmark_names`` is the running list of canonical benchmark names that
    were scheduled for this invocation. Those rows are refreshed in the
    summary table; other benchmarks already written to CSV/Excel are kept.

    After a successful ``run_task()`` we also snapshot the cumulative
    ``perf_metrics.summary`` and the per-sample predictions into the durable
    backup files maintained by ``perf_backup.py`` so checkpoint restarts can
    recover them.

    If ``summary_lock`` is provided, the summary write is protected so multiple
    seeds/runs can safely refresh the table as each one finishes.
    """
    dataset_name = task_cfg.datasets[0]
    work_dir = Path(task_cfg.work_dir)
    restore_before_run(output_dir, dataset_name, model_name, work_dir)

    start_ts = time.monotonic()
    try:
        run_task(task_cfg)
    finally:
        elapsed = time.monotonic() - start_ts
        perf_backup_module.record_active_time(output_dir, dataset_name, model_name, elapsed)
        print(f'Active time for {dataset_name}: {elapsed:.1f}s (total accumulated)')

    backup_after_run(output_dir, dataset_name, model_name, work_dir)
    if write_summary_flag:
        if summary_lock is not None:
            with summary_lock:
                write_summary(output_dir, model_name, folder_name, benchmark_names=benchmark_names)
        else:
            write_summary(output_dir, model_name, folder_name, benchmark_names=benchmark_names)


def backup_after_run(output_dir: str, benchmark: str, model_name: str,
                     work_dir: Path):
    """Snapshot the just-finished run's perf summary and predictions."""
    try:
        report_json = work_dir / 'reports' / f'{benchmark}.json'
        perf_backup_module.backup_perf_stats(
            Path(output_dir), benchmark, model_name, report_json,
        )
        predictions_dir = work_dir / 'predictions'
        perf_backup_module.archive_predictions(
            Path(output_dir), benchmark, model_name, predictions_dir,
        )
    except Exception as e:
        print(f'WARNING: perf backup for {benchmark} failed: {e}')


def restore_before_run(output_dir: str, benchmark: str, model_name: str,
                        work_dir: Path) -> bool:
    """If durable backups exist for a benchmark/model, materialise them into
    the new ``work_dir`` before evalscope starts so the next run resumes from
    the larger historical state instead of overwriting it.
    """
    try:
        predictions_dir = work_dir / 'predictions'
        report_json = work_dir / 'reports' / f'{benchmark}.json'
        restored = perf_backup_module.restore_from_backup(
            Path(output_dir), benchmark, model_name,
            predictions_dir, report_json,
        )
        if restored:
            print(f'Restored {benchmark}/{model_name} from backup before run')
        return restored
    except Exception as e:
        print(f'WARNING: perf restore for {benchmark} failed: {e}')
        return False


def main():
    global RESULTS_DIR
    parser = build_parser()
    args = parser.parse_args()
    RESULTS_DIR = Path(args.results_dir) if args.results_dir else (PROJECT_ROOT / 'results')

    # Resolve limit
    limit = args.limit
    if limit is not None:
        if str(limit).lower() in ('none', 'all'):
            limit = None
        else:
            limit = int(limit)

    enable_thinking = DEFAULT_ENABLE_THINKING if args.thinking is None else args.thinking

    # Resolve suite or custom datasets
    if args.datasets:
        custom = [d.strip() for d in args.datasets.split(',') if d.strip()]
        multi_run = [d for d in custom if d in MULTI_RUN_CONFIG]
        single_run = [d for d in custom if d not in MULTI_RUN_CONFIG]
        agent = [d for d in custom if d in ALL_AGENT]
        single_run = [d for d in single_run if d not in ALL_AGENT]
        fingerprint = [d for d in custom if d in ALL_FINGERPRINT or d in FUSION_BENCHMARKS]
        single_run = [d for d in single_run if d not in ALL_FINGERPRINT
                      and d not in FUSION_BENCHMARKS]
    else:
        suite = SUITES[args.suite]
        multi_run = list(suite['multi'])
        single_run = list(suite['single'])
        agent = list(suite['agent'])
        fingerprint = list(suite.get('fingerprint', []))

    # Apply --exclude
    if args.exclude:
        exclude = {d.strip() for d in args.exclude.split(',') if d.strip()}
        multi_run = [d for d in multi_run if d not in exclude]
        single_run = [d for d in single_run if d not in exclude]
        agent = [d for d in agent if d not in exclude]
        fingerprint = [d for d in fingerprint if d not in exclude]

    judge_model_args = {
        'model_id': args.judge_model,
        'api_url': args.judge_api_url,
        'api_key': args.judge_api_key,
        'eval_type': 'openai_api',
        'generation_config': {
            'temperature': 0.0,
            'max_tokens': args.judge_max_tokens,
        },
    }

    truncation_tokens = args.truncation_tokens
    _patch_adapters_for_truncation(args.tokenizer_path, truncation_tokens)

    # Auto-select thinking config if --thinking is enabled and user did not
    # explicitly provide a different --config.
    config_path = args.config
    if enable_thinking and args.config == DEFAULT_CONFIG:
        thinking_config = str(PROJECT_ROOT / 'config' / 'dpv4-int8_thinking.yaml')
        if Path(thinking_config).exists():
            config_path = thinking_config
            print(f'Auto-selected thinking config: {config_path}')
        else:
            print(f'WARNING: thinking config not found at {thinking_config}, falling back to {args.config}')

    dataset_configs = load_dataset_configs(config_path)

    # Resolve top-level output folder name
    if args.folder_name:
        folder_name = args.folder_name
    elif enable_thinking:
        safe_model = args.model.replace('/', '_').replace('\\', '_').replace(' ', '_')
        folder_name = f'{safe_model}_THINKING'
    else:
        folder_name = args.model.replace('/', '_').replace('\\', '_').replace(' ', '_')
    model_output_dir = Path(args.output_dir) / folder_name
    model_output_dir.mkdir(parents=True, exist_ok=True)
    results_dir = RESULTS_DIR
    results_dir.mkdir(parents=True, exist_ok=True)

    print('=' * 60)
    print(f'Config: {args.config}')
    print(f'Model: {args.model}')
    print(f'API URL: {args.api_url}')
    print(f'Dataset Dir: {args.dataset_dir}')
    print(f'Output Root: {args.output_dir}')
    print(f'Summary CSV/Excel: {results_dir}')
    print(f'Output Folder: {folder_name}')
    print(f'Suite: {args.suite}')
    print(f'Limit: {limit if limit is not None else "ALL"}')
    print(f'Thinking: {enable_thinking}')
    if enable_thinking:
        print(f'Thinking max_tokens scale: {args.thinking_max_tokens_scale}')
    if args.thinking_budget_tokens is not None:
        print(f'Thinking budget_tokens: {args.thinking_budget_tokens}')
    print(f'Seed: {args.seed}')
    print(f'Batch Size: {args.batch_size}')
    print(f'Tokenizer Path: {args.tokenizer_path}')
    print(f'Truncation Tokens: {truncation_tokens}')
    print(f'Multi-run datasets: {multi_run}')
    print(f'Single-run datasets: {single_run}')
    print(f'Agent datasets: {agent}')
    print(f'Fingerprint datasets: {fingerprint}')
    print(f'Write summary: {args.write_summary}')
    print('=' * 60)

    def get_dataset_config(dataset_name: str) -> dict:
        """Return configured dataset config, or a default config for unknown benchmarks."""
        if dataset_name in dataset_configs:
            return dataset_configs[dataset_name]
        print(f'WARNING: {dataset_name} not in YAML config, using default generation_config '
              f'(temperature={DEFAULT_GENERATION_CONFIG["temperature"]}, '
              f'top_p={DEFAULT_GENERATION_CONFIG["top_p"]}, '
              f'max_tokens={DEFAULT_GENERATION_CONFIG["max_tokens"]})')
        return {'generation_config': deepcopy(DEFAULT_GENERATION_CONFIG)}

    def run_fingerprint_benchmark(dataset_name: str, benchmark_names: list,
                                  write_summary_flag: bool = True,
                                  summary_lock: threading.Lock = None):
        """Run one fingerprint benchmark via its standalone runner script.

        产出与 EvalScope 一致的 output/<folder>/<bench>/seed_<seed>/reports/*.json，
        并复用 active_time 计时与 perf 备份，保证 collect_results 可直接汇总。
        """
        work_dir = model_output_dir / dataset_name / f'seed_{args.seed}'
        report_path = work_dir / 'reports' / f'{dataset_name}.json'
        script = FINGERPRINT_SCRIPTS[dataset_name]

        common_cmd = [
            '--api-url', args.api_url,
            '--model', args.model,
            '--report-path', str(report_path),
            '--timeout', str(args.fingerprint_timeout),
        ]
        if dataset_name == 'llmmap':
            cmd = [args.llmmap_python, str(script), *common_cmd,
                   '--tools-root', args.tools_root]
            if args.expected_model:
                cmd += ['--expected-model', args.expected_model]
        elif dataset_name == 'fp_fusion':
            cmd = [sys.executable, str(script), *common_cmd,
                   '--tools-root', args.tools_root]
            if args.detector_reference:
                cmd += ['--reference', args.detector_reference]
        elif dataset_name == 'llm_verify':
            cmd = [args.verify_python, str(script), *common_cmd,
                   '--tools-root', args.tools_root]
        else:  # llm_fingerprint_detector
            cmd = [sys.executable, str(script), *common_cmd,
                   '--tools-root', args.tools_root,
                   '--preset', args.detector_preset,
                   '--concurrency', '4']
            if args.detector_node and args.detector_node != 'node':
                cmd += ['--node', args.detector_node]
            if args.detector_reference:
                cmd += ['--reference', args.detector_reference]

        print(f"\n{'='*60}")
        print(f'Running: {dataset_name} (fingerprint benchmark, seed={args.seed})')
        print(f"{'='*60}")

        start_ts = time.monotonic()
        try:
            proc = subprocess.run(cmd, capture_output=True, text=True,
                                  timeout=FP_OVERALL_TIMEOUT)
            tail = '\n'.join((proc.stdout or '').strip().splitlines()[-20:])
            if tail:
                print(tail)
            if proc.returncode != 0:
                err_tail = '\n'.join((proc.stderr or '').strip().splitlines()[-10:])
                print(f'ERROR in {dataset_name}: exit={proc.returncode}\n{err_tail}')
                fingerprint_common.write_report(
                    str(report_path), dataset_name, 0.0, num=0,
                    error=f'runner exited with code {proc.returncode}',
                    stderr_tail=err_tail[-800:],
                )
        except Exception as e:
            print(f'ERROR in {dataset_name}: {e}')
            fingerprint_common.write_report(
                str(report_path), dataset_name, 0.0, num=0,
                error=str(e),
            )
        finally:
            elapsed = time.monotonic() - start_ts
            perf_backup_module.record_active_time(str(model_output_dir), dataset_name,
                                                  args.model, elapsed)
            print(f'Active time for {dataset_name}: {elapsed:.1f}s')

        backup_after_run(str(model_output_dir), dataset_name, args.model, work_dir)
        if write_summary_flag:
            if summary_lock is not None:
                with summary_lock:
                    write_summary(str(model_output_dir), args.model, folder_name,
                                  benchmark_names=benchmark_names)
            else:
                write_summary(str(model_output_dir), args.model, folder_name,
                              benchmark_names=benchmark_names)

    def run_one(dataset_name, run_idx=0, benchmark_names=None, write_summary_flag=True,
                summary_lock=None):
        if dataset_name in FINGERPRINT_SCRIPTS:
            run_fingerprint_benchmark(dataset_name, benchmark_names or [dataset_name],
                                      write_summary_flag=write_summary_flag,
                                      summary_lock=summary_lock)
            return

        ds_cfg = get_dataset_config(dataset_name)
        task_cfg = build_task_config(
            dataset_name, ds_cfg, args.batch_size, enable_thinking, args.seed, limit,
            str(model_output_dir), args.model, args.api_url, args.api_key, args.dataset_dir, judge_model_args,
            run_idx=run_idx,
            thinking_max_tokens_scale=args.thinking_max_tokens_scale,
            max_tokens_add=args.max_tokens_add,
            thinking_budget_tokens=args.thinking_budget_tokens,
        )
        try:
            run_and_summarize(task_cfg, write_summary_flag, str(model_output_dir), args.model,
                              folder_name=folder_name, benchmark_names=benchmark_names,
                              summary_lock=summary_lock)
        except Exception as e:
            print(f'ERROR in {dataset_name} (run {run_idx + 1 if run_idx else 1}): {e}')

    # Benchmarks scheduled in this invocation: their summary rows are upserted;
    # other existing CSV/Excel rows are left in place.
    benchmark_names = []

    # Build a list of benchmark execution units.  Each unit runs one complete
    # benchmark (all its multi-runs).  This lets us run benchmarks either
    # serially or in parallel while keeping per-benchmark run parallelism
    # (``--parallel-runs``) intact.
    benchmark_units = []
    for dataset_name in multi_run:
        benchmark_names.append(dataset_name)
        benchmark_units.append((dataset_name, 'multi'))
    for dataset_name in single_run:
        benchmark_names.append(dataset_name)
        benchmark_units.append((dataset_name, 'single'))
    for dataset_name in agent:
        benchmark_names.append(dataset_name)
        benchmark_units.append((dataset_name, 'agent'))
    for dataset_name in fingerprint:
        benchmark_names.append(dataset_name)
        benchmark_units.append((dataset_name, 'fingerprint'))

    def run_benchmark_unit(dataset_name: str, kind: str, summary_lock=None):
        """Run one benchmark (all seeds/runs) and return its name."""
        if kind == 'multi':
            num_runs = MULTI_RUN_CONFIG.get(dataset_name, 1)
            parallel_runs = max(1, args.parallel_runs)
            if parallel_runs <= 1:
                for run_idx in range(num_runs):
                    print(f"\n{'='*60}")
                    print(f'Running: {dataset_name} (run {run_idx + 1}/{num_runs}, seed={args.seed})')
                    print(f"{'='*60}")
                    run_one(dataset_name, run_idx=run_idx, benchmark_names=benchmark_names,
                            write_summary_flag=True, summary_lock=summary_lock)
            else:
                print(f"\n{'='*60}")
                print(f'Running: {dataset_name} ({num_runs} runs, {parallel_runs} in parallel, '
                      f'seed={args.seed}, per-run batch_size={args.batch_size})')
                print(f"{'='*60}")
                from concurrent.futures import ThreadPoolExecutor, as_completed

                def _run_single(run_idx):
                    print(f'  -> start run {run_idx + 1}/{num_runs}')
                    run_one(dataset_name, run_idx=run_idx, benchmark_names=benchmark_names,
                            write_summary_flag=True, summary_lock=summary_lock)
                    print(f'  -> finish run {run_idx + 1}/{num_runs}')

                with ThreadPoolExecutor(max_workers=parallel_runs) as executor:
                    futures = [executor.submit(_run_single, run_idx) for run_idx in range(num_runs)]
                    for future in as_completed(futures):
                        try:
                            future.result()
                        except Exception as e:
                            print(f'ERROR in {dataset_name} parallel run: {e}')
        else:
            print(f"\n{'='*60}")
            print(f'Running: {dataset_name} (seed={args.seed})')
            print(f"{'='*60}")
            run_one(dataset_name, benchmark_names=benchmark_names, write_summary_flag=True,
                    summary_lock=summary_lock)
        return dataset_name

    # Lock to protect concurrent writes to the summary CSV/Excel when multiple
    # seeds/runs or benchmarks finish around the same time.
    summary_lock = threading.Lock()

    parallel_benchmarks = max(1, args.parallel_benchmarks)
    if parallel_benchmarks <= 1:
        for dataset_name, kind in benchmark_units:
            run_benchmark_unit(dataset_name, kind, summary_lock=summary_lock)
    else:
        print(f"\n{'='*60}")
        print(f'Running benchmarks in parallel: {len(benchmark_units)} units, '
              f'{parallel_benchmarks} at a time')
        print(f"{'='*60}")
        from concurrent.futures import ThreadPoolExecutor, as_completed

        with ThreadPoolExecutor(max_workers=parallel_benchmarks) as executor:
            futures = [executor.submit(run_benchmark_unit, dataset_name, kind, summary_lock)
                       for dataset_name, kind in benchmark_units]
            for future in as_completed(futures):
                try:
                    dataset_name = future.result()
                    print(f'  -> benchmark finished: {dataset_name}')
                except Exception as e:
                    print(f'ERROR in parallel benchmark: {e}')

    if args.write_summary:
        with summary_lock:
            write_summary(str(model_output_dir), args.model, folder_name, benchmark_names=benchmark_names)
    print('\nAll benchmarks done!')


if __name__ == '__main__':
    main()