#!/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 sys
import time
from copy import deepcopy
from pathlib import Path

import yaml

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

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

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

DEFAULT_MODEL = 'DeepSeek-V4-Flash-Int8'
DEFAULT_API_URL = 'http://localhost:30000/v1'
DEFAULT_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']

# 分组基于 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'],
    },
    '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'],
    },
    # 腾讯/官方入库标准对比套件：取各厂商公开对比表的并集，约 11 个 benchmark
    # 注：swe_bench_verified/terminal_bench_v2/browsecomp/tau2_bench 需要额外镜像/工具
    'official': {
        'multi': ['aime25', 'aime26','live_code_bench'],
        'single': [
            'hle', 'mmlu_pro', 'gpqa_diamond','longbench_v2',
            # 'swe_bench_verified',
        ],
        'agent': [
            'tau2_bench'
        ],
    },
}

# ============================================================
# 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'}
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': {}
        }
    },
}


# ============================================================
# 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)')

    # 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('--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)')

    # 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('--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)')

    # 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 None:
        from transformers import AutoTokenizer
        try:
            _TOKENIZER = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
        except Exception:
            _TOKENIZER = AutoTokenizer.from_pretrained('deepseek-ai/DeepSeek-V4-Flash', trust_remote_code=True)
    return _TOKENIZER


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.longbench_v2.longbench_v2_adapter import LongBenchV2Adapter
    from evalscope.benchmarks.openai_mrcr.openai_mrcr_adapter import OpenAIMRCRAdapter

    _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) -> 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
    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,
    dataset_dir: str,
    judge_model_args: dict,
    run_idx: int = 0,
    thinking_max_tokens_scale: float = 1.0,
    max_tokens_add: int = 0,
) -> 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)
    work_dir = str(work_dir)

    generation_config = configure_thinking(deepcopy(ds_cfg['generation_config']), enable_thinking)
    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,
        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=work_dir,
        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
# ============================================================

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 = PROJECT_ROOT / 'results'
        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):
    """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. The summary table will be limited to
    these benchmarks so it does not pick up older results left in the output
    directory from previous runs.

    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.
    """
    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:
        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():
    parser = build_parser()
    args = parser.parse_args()

    # 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]
    else:
        suite = SUITES[args.suite]
        multi_run = list(suite['multi'])
        single_run = list(suite['single'])
        agent = list(suite['agent'])

    # 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]

    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)

    dataset_configs = load_dataset_configs(args.config)

    # 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 = PROJECT_ROOT / 'results'
    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'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}')
    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'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_one(dataset_name, run_idx=0, benchmark_names=None):
        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.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,
        )
        try:
            run_and_summarize(task_cfg, args.write_summary, str(model_output_dir), args.model,
                              folder_name=folder_name, benchmark_names=benchmark_names)
        except Exception as e:
            print(f'ERROR in {dataset_name} (run {run_idx + 1 if run_idx else 1}): {e}')

    # Keep the ordered list of benchmarks actually run in this invocation so
    # the summary table only includes them, not stale results from earlier runs.
    benchmark_names = []

    for dataset_name in multi_run:
        benchmark_names.append(dataset_name)
        num_runs = MULTI_RUN_CONFIG.get(dataset_name, 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)

    for dataset_name in single_run:
        benchmark_names.append(dataset_name)
        print(f"\n{'='*60}")
        print(f'Running: {dataset_name} (seed={args.seed})')
        print(f"{'='*60}")
        run_one(dataset_name, benchmark_names=benchmark_names)

    for dataset_name in agent:
        benchmark_names.append(dataset_name)
        print(f"\n{'='*60}")
        print(f'Running: {dataset_name} (seed={args.seed})')
        print(f"{'='*60}")
        run_one(dataset_name, benchmark_names=benchmark_names)

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


if __name__ == '__main__':
    main()
