import yaml
from copy import deepcopy
from pathlib import Path
import sys
import os

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

# ============================================================
# ★★★ 必改参数 (USER CONFIG) ★★★
# 每次评测新模型前，只需要检查/修改以下参数。
# 也可以通过命令行覆盖：--model / --api-url / --dataset-dir /
# --output-dir / --limit / --config
# ============================================================

# 模型名（served model name）
MODEL = 'DeepSeek-V4-Flash-Int8'

# 模型服务地址（OpenAI 兼容 API）
API_URL = 'http://localhost:30000/v1'

# 本地数据集根目录（提前下载好的 datasets 目录）
DATASET_DIR = str(PROJECT_ROOT / 'datasets')

# 评测输出目录（每个 benchmark 单独子目录）
OUTPUT_DIR = str(PROJECT_ROOT / 'output')

# 采样数量上限；None 表示跑全部样本
LIMIT = None

# 每个 benchmark 的生成参数配置文件（max_tokens / temperature 等）
CONFIG = None

# 是否开启 thinking 模式（sglang chat_template_kwargs.thinking）
ENABLE_THINKING = False

# 测试哪些 benchmark：见下方 DATASETS = _multi_run_order + _single_run_order + _agent_order
# 想只跑部分 benchmark 时，注释掉对应行即可。
MODEL_PATH = '/data1/models/DeepSeek-V4-Flash-INT8'

# Judge model for benchmarks that require LLM-as-judge (simple_qa, tau2_bench, etc.)
# Vectron endpoint with DeepSeek-V4-Pro
JUDGE_MODEL_ARGS = {
    'model_id': 'DeepSeek/DeepSeek-V4-Pro',
    'api_url': 'https://api.vectron.meta-stone.com/v1',
    'api_key': 'sk-dbd8a665f7634081b87ec409c7636500',
    'eval_type': 'openai_api',
    'generation_config': {
        'temperature': 0.0,
        'max_tokens': 10240,
    },
}

TRUNCATION_CONFIG = {
    'longbench_v2': 32768*4,
    'openai_mrcr': 32768*4,
}
# Combine all datasets in order (shortest estimated time first)
# 1. Multi-run small benchmarks (temperature=1.0 for sampling diversity)
_multi_run_order = [
    # 'humaneval',        # ~164 samples x 3 runs = 492
    # 'live_code_bench',  # ~100 samples x 5 runs = 500
    # 'aime26',           # ~30 samples x 12 runs = 360
    # 'aime24',           # ~30 samples x 12 runs = 360
    # 'aime25',           # ~30 samples x 12 runs = 360
    # 'gpqa_diamond',     # ~198 samples x 2 runs = 396
    # 'imo_answerbench',  # ~120 samples x 4 runs = 480
    # 'hmmt26',           # ~30 samples x 12 runs = 360
]

# 2. Single-run benchmarks (temperature=0.0 greedy)
_single_run_order = [
    # 'arc',
    # 'bfcl_v3',          # function calling: greedy decoding
    # 'winogrande',
    # 'competition_math',
    # 'gsm8k',
    # 'hellaswag',
    # 'bigcodebench',
    # 'drop',
    # 'bbh',
    # 'openai_mrcr',
    # 'longbench_v2',
    # 'mmlu',
    # 'cmmlu',
    # 'simple_qa',
    # 'mmlu_pro',
    'hle',
    # 'trivia_qa',
]

# 3. Agent / tool benchmarks (last) 
_agent_order = [
    # 'tau2_bench',
    # 'general_fc',
]
# ============================================================
# Command line argument overrides (不需要修改)
# ============================================================

for i, arg in enumerate(sys.argv):
    if arg == '--limit' and i + 1 < len(sys.argv):
        limit_val = sys.argv[i + 1]
        if limit_val.lower() == 'none' or limit_val.lower() == 'all':
            LIMIT = None
        else:
            LIMIT = int(limit_val)
    elif arg == '--model' and i + 1 < len(sys.argv):
        MODEL = sys.argv[i + 1]
    elif arg == '--api-url' and i + 1 < len(sys.argv):
        API_URL = sys.argv[i + 1]
    elif arg == '--dataset-dir' and i + 1 < len(sys.argv):
        DATASET_DIR = sys.argv[i + 1]
    elif arg == '--output-dir' and i + 1 < len(sys.argv):
        OUTPUT_DIR = sys.argv[i + 1]
    elif arg == '--config' and i + 1 < len(sys.argv):
        CONFIG = sys.argv[i + 1]

# Benchmarks that require sandboxed code execution
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': {}
        }
    },
}

# ============================================================
# User-tunable parameters
# ============================================================

# Fixed seed for reproducibility. With temperature > 0 the model still samples
# randomly, so running N times with the same seed still yields variance.
SEED = 42

# Benchmarks to run multiple times with temperature=1.0.
# Format: {benchmark_name: num_runs}
# Number of runs chosen so total samples ≈ 400-500 per benchmark.
MULTI_RUN_CONFIG = {
    # ~30 samples each -> 12 runs = ~360 samples
    'aime24': 12,
    'aime25': 12,
    'aime26': 12,
    'hmmt26': 12,
    # ~100 samples -> 5 runs = ~500 samples
    'live_code_bench': 5,
    # ~120 samples -> 4 runs = ~480 samples
    'imo_answerbench': 4,
    # ~164 samples -> 3 runs = ~492 samples
    'humaneval': 3,
    # ~198 samples -> 2 runs = ~396 samples
    'gpqa_diamond': 2,
}

# Single-run benchmarks (temperature=0.0 greedy)
SINGLE_RUN_DATASETS = [
    'bigcodebench',
    'bfcl_v3',          # function calling: greedy decoding
    'competition_math',
    'gsm8k',
    'hle',
    'super_gpqa',
    'arc',
    'bbh',
    'cmmlu',
    'drop',
    'hellaswag',
    'mmlu',
    'mmlu_pro',
    'simple_qa',
    'trivia_qa',
    'winogrande',
    'openai_mrcr',      # long-context, run once
    'longbench_v2',     # long-context, run once
]

# Agent / tool benchmarks (temperature=0.0 greedy, run once)
AGENT_DATASETS = [
    'tau2_bench',
    'general_fc',
]



DATASETS = _multi_run_order + _single_run_order + _agent_order

# ENABLE_THINKING 已移至文件头部「必改参数」区
BATCH_SIZE_LIST = [4]
# LIMIT is parsed from command line: --limit N
SHUFFLE = True

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

CONFIG_PATH = str(Path(CONFIG) if CONFIG else SCRIPT_DIR / 'config' / 'dpv4-int8_nothinking.yaml')
if not Path(CONFIG_PATH).exists():
    raise FileNotFoundError(f'Config file not found: {CONFIG_PATH}')


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{{}}."
)

with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
    DATASET_CONFIGS = yaml.safe_load(f)


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

_TOKENIZER = None


def get_tokenizer():
    global _TOKENIZER
    if _TOKENIZER is None:
        from transformers import AutoTokenizer
        try:
            # Try local path first
            _TOKENIZER = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
        except Exception:
            # Fallback to HuggingFace model ID
            _TOKENIZER = AutoTokenizer.from_pretrained('deepseek-ai/DeepSeek-V4-Flash', trust_remote_code=True)
    return _TOKENIZER


def truncate_middle(text: str, max_tokens: int) -> str:
    if max_tokens <= 0:
        return text
    tokenizer = get_tokenizer()
    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():
    from evalscope.benchmarks.longbench_v2.longbench_v2_adapter import LongBenchV2Adapter
    from evalscope.benchmarks.openai_mrcr.openai_mrcr_adapter import OpenAIMRCRAdapter

    # Patch LongBenchV2Adapter.format_prompt_template
    _orig_longbench_format = LongBenchV2Adapter.format_prompt_template
    def _patched_longbench_format(self, sample):
        max_tok = TRUNCATION_CONFIG.get('longbench_v2')
        if max_tok and sample.metadata and 'context' in sample.metadata:
            sample.metadata['context'] = truncate_middle(sample.metadata['context'], max_tok)
        return _orig_longbench_format(self, sample)
    LongBenchV2Adapter.format_prompt_template = _patched_longbench_format

    # Patch OpenAIMRCRAdapter.record_to_sample with needle-aware truncation.
    # MRCR is a long chat history with needles hidden at desired_msg_index.
    # We keep the head, tail, and a window around the needle, and truncate
    # each kept message if it is still too long. This preserves the retrieval
    # task while fitting GPU memory.
    _orig_mrcr_record = OpenAIMRCRAdapter.record_to_sample
    def _patched_mrcr_record(self, record):
        import json
        max_total_tok = TRUNCATION_CONFIG.get('openai_mrcr')
        per_msg_max_tok = 8192
        if max_total_tok and '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()
                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 <= max_total_tok:
                    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()
                # Head and tail context
                keep.update(range(min(2, n)))
                keep.update(range(max(0, n - 2), n))
                # Window around the needle
                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)
                    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


_patch_adapters_for_truncation()


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

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, run_idx: int = 0) -> TaskConfig:
    # Each run gets its own work_dir so use_cache does not reuse predictions
    # across repeated samples. This is required for temperature=1.0 multi-run
    # benchmarks to actually measure variance.
    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)
    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 loop
# ============================================================

def main():
    print(f"Config: {CONFIG_PATH}")
    print(f"Model: {MODEL}")
    print(f"API URL: {API_URL}")
    print(f"Dataset Dir: {DATASET_DIR}")
    print(f"Output Dir: {OUTPUT_DIR}")
    print(f"Limit: {LIMIT if LIMIT is not None else 'ALL'}")
    print(f"Total benchmarks: {len(DATASETS)}")
    print("="*60)

    for batch_size in BATCH_SIZE_LIST:
        # 1. Run multi-run benchmarks (temperature=1.0, same seed, variance from sampling)
        for dataset_name in _multi_run_order:
            if dataset_name not in DATASET_CONFIGS:
                print(f"WARNING: {dataset_name} not in YAML config, skipping")
                continue
            ds_cfg = DATASET_CONFIGS[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={SEED})")
                print(f"{'='*60}")
                task_cfg = build_task_config(dataset_name, ds_cfg, batch_size, ENABLE_THINKING, SEED, run_idx=run_idx)
                try:
                    run_task(task_cfg)
                except Exception as e:
                    print(f"ERROR in {dataset_name} (run {run_idx + 1}): {e}")
                    continue

        # 2. Run single-run benchmarks (temperature=0.0 greedy)
        for dataset_name in _single_run_order:
            if dataset_name not in DATASET_CONFIGS:
                print(f"WARNING: {dataset_name} not in YAML config, skipping")
                continue
            ds_cfg = DATASET_CONFIGS[dataset_name]
            print(f"\n{'='*60}")
            print(f"Running: {dataset_name} (seed={SEED})")
            print(f"{'='*60}")
            task_cfg = build_task_config(dataset_name, ds_cfg, batch_size, ENABLE_THINKING, SEED)
            try:
                run_task(task_cfg)
            except Exception as e:
                print(f"ERROR in {dataset_name}: {e}")
                continue

        # 3. Run agent benchmarks
        for dataset_name in _agent_order:
            if dataset_name not in DATASET_CONFIGS:
                print(f"WARNING: {dataset_name} not in YAML config, skipping")
                continue
            ds_cfg = DATASET_CONFIGS[dataset_name]
            print(f"\n{'='*60}")
            print(f"Running: {dataset_name} (seed={SEED})")
            print(f"{'='*60}")
            task_cfg = build_task_config(dataset_name, ds_cfg, batch_size, ENABLE_THINKING, SEED)
            try:
                run_task(task_cfg)
            except Exception as e:
                print(f"ERROR in {dataset_name}: {e}")
                continue

    print("\nAll benchmarks done!")


if __name__ == '__main__':
    main()