#!/usr/bin/env python3
"""
k3_report_test.py

复现 Kimi K3 report（https://www.kimi.com/blog/kimi-k3）中可在本地 evalscope v1.9.1
上直接运行的 benchmark。对于需要外部 agent harness（Claude Code / Codex / Kimi Code）
的 benchmark，脚本会检查依赖并给出安装/配置提示。

用法:
    # 只看哪些能跑、哪些不能跑
    python3 k3_report_test.py --dry-run

    # 跑所有 evalscope 支持的 benchmark（limit 5 做冒烟）
    export EVAL_API_KEY="sk-xxxx"
    export EVAL_API_URL="https://api.example.com/v1"
    export EVAL_MODEL="kimi-k3"
    python3 k3_report_test.py --limit 5

    # 只跑指定类别
    python3 k3_report_test.py --categories Coding,Vision --limit 5

    # 只跑单个 benchmark
    python3 k3_report_test.py --datasets deep_swe --limit 1
"""

import argparse
import importlib.util
import os
import shutil
import subprocess
import sys
from pathlib import Path

def _find_repo_root() -> Path:
    """向上查找，直到目录下存在 bash/run.py。"""
    p = Path(__file__).resolve().parent
    while p != p.parent:
        if (p / "bash" / "run.py").exists():
            return p
        p = p.parent
    # 兜底：脚本所在目录的上级（兼容放在 bash/case/ 的情况）
    return Path(__file__).resolve().parent.parent


ROOT = _find_repo_root()

# ---------------------------------------------------------------------------
# Kimi K3 report -> evalscope dataset 映射
# ---------------------------------------------------------------------------
BENCHMARKS = {
    "Coding": [
        {
            "report_name": "DeepSWE",
            "dataset": "deep_swe",
            "deps": ["harbor"],
            "note": "需要 agent harness，推荐 Kimi Code 或 mini-SWE-agent；evalscope 提供 deep_swe adapter",
        },
        {
            "report_name": "Terminal Bench 2.1",
            "dataset": "terminal_bench_v2_1",
            "deps": ["harbor"],
            "note": "需要 harbor 框架；evalscope 已提供 TerminalBenchV2_1 adapter",
        },
    ],
    "Agentic": [
        {
            "report_name": "GDPval-AA v2",
            "dataset": "gdpval",
            "deps": [],
            "note": "Elo-score 需多模型结果聚合，单模型只能得到 raw score",
        },
        {
            "report_name": "BrowseComp",
            "dataset": "browsecomp",
            "deps": [],
        },
        {
            "report_name": "Toolathlon-Verified",
            "dataset": "toolathlon",
            "deps": [],
            "note": "toolathlon 公开子集；Verified 子集可能需额外配置",
        },
        {
            "report_name": "MCP Atlas",
            "dataset": "mcp_atlas",
            "deps": [],
            "note": "public 500-task subset",
        },
    ],
    "Reasoning & Knowledge": [
        {
            "report_name": "GPQA-Diamond",
            "dataset": "gpqa_diamond",
            "deps": [],
        },
        {
            "report_name": "HLE-Full",
            "dataset": "hle",
            "deps": [],
            "note": "w/ tools 变体没有独立 dataset，可用本地工具或 judge 扩展",
        },
    ],
    "Vision": [
        {
            "report_name": "MMMU-Pro",
            "dataset": "mmmu_pro",
            "deps": [],
        },
        {
            "report_name": "CharXiv (RQ)",
            "dataset": "charxiv",
            "deps": [],
        },
        {
            "report_name": "MathVision",
            "dataset": "math_vision",
            "deps": [],
        },
        {
            "report_name": "BabyVision w/ python",
            "dataset": "baby_vision",
            "deps": [],
            "note": "baby_vision 基础版可用；w/ python 需额外工具配置",
        },
        {
            "report_name": "ZeroBench_main (pass@5)",
            "dataset": "zerobench",
            "deps": [],
            "note": "pass@5 需设置 n_samples / temperature，详见 adapter 文档",
        },
        {
            "report_name": "WorldVQA ForceAnswer",
            "dataset": "world_vqa",
            "deps": [],
        },
        {
            "report_name": "OmniDocBench",
            "dataset": "omni_doc_bench",
            "deps": [],
        },
    ],
}

# 这些 benchmark 在当前 evalscope v1.9.1 里没有对应 dataset
UNSUPPORTED = [
    ("Coding", "Program Bench", "未开源/未接入 evalscope"),
    ("Coding", "FrontierSWE", "未接入 evalscope"),
    ("Coding", "SWE Marathon", "未接入 evalscope"),
    ("Coding", "PostTrain Bench", "未接入 evalscope"),
    ("Coding", "MLS Bench", "未接入 evalscope"),
    ("Coding", "Kimi Code Bench 2.0 (Internal)", "Kimi 内部 benchmark"),
    ("Agentic", "DeepSearchQA", "未接入 evalscope"),
    ("Agentic", "Automation Bench", "未接入 evalscope"),
    ("Agentic", "Job Bench", "未接入 evalscope"),
    ("Agentic", "AA-Briefcase", "未接入 evalscope"),
    ("Agentic", "APEX-Agents", "未接入 evalscope"),
    ("Agentic", "Office QA Pro", "evalscope 只有 OfficeQA，Pro 版未接入"),
    ("Agentic", "SpreadsheetBench 2", "未接入 evalscope"),
    ("Agentic", "DECK-Bench (Internal)", "Kimi 内部 benchmark"),
    ("Vision", "PerceptionBench", "未接入 evalscope"),
]


def check_dep(dep: str) -> bool:
    """检查 Python 包或系统命令是否存在。"""
    if dep == "harbor":
        return shutil.which("harbor") is not None or importlib.util.find_spec("harbor") is not None
    return shutil.which(dep) is not None or importlib.util.find_spec(dep) is not None


def print_support_matrix():
    print("=" * 70)
    print("Kimi K3 report benchmark 在 evalscope v1.9.1 中的支持情况")
    print("=" * 70)
    for category, items in BENCHMARKS.items():
        print(f"\n【{category}】")
        for item in items:
            missing = [d for d in item.get("deps", []) if not check_dep(d)]
            status = "✅ 可运行" if not missing else f"⚠️  缺依赖: {', '.join(missing)}"
            print(f"  {item['report_name']:30} -> {item['dataset']:25} {status}")
            if item.get("note"):
                print(f"      note: {item['note']}")

    print("\n【暂不支持 / 未接入】")
    for category, name, reason in UNSUPPORTED:
        print(f"  [{category}] {name}: {reason}")
    print("=" * 70)


def build_run_command(
    dataset: str,
    model: str,
    api_url: str,
    api_key: str,
    limit,
    output_dir: str,
    folder_name: str,
    config: str,
    thinking: bool,
    thinking_budget_tokens,
) -> list:
    cmd = [
        sys.executable,
        str(ROOT / "bash" / "run.py"),
        "--datasets", dataset,
        # "--model", model,
        # "--api-url", api_url,
        "--output-dir", output_dir,
        "--folder-name", folder_name,
        "--config", config,
        "--batch-size", "4",
    ]
    if api_key:
        cmd += ["--api-key", api_key]
    if limit is not None:
        cmd += ["--limit", str(limit)]
    if thinking:
        cmd.append("--thinking")
    if thinking_budget_tokens is not None:
        cmd += ["--thinking-budget-tokens", str(thinking_budget_tokens)]
    return cmd


def run_one(item: dict, args) -> int:
    report_name = item["report_name"]
    dataset = item["dataset"]
    print(f"\n>>> Running {report_name} ({dataset}) ...")

    missing = [d for d in item.get("deps", []) if not check_dep(d)]
    if missing:
        print(f"SKIP: 缺少依赖 {missing}；{item.get('note', '')}")
        return 0

    cmd = build_run_command(
        dataset=dataset,
        model=args.model,
        api_url=args.api_url,
        api_key=args.api_key,
        limit=args.limit,
        output_dir=args.output_dir,
        folder_name=args.folder_name,
        config=args.config,
        thinking=args.thinking,
        thinking_budget_tokens=args.thinking_budget_tokens,
    )

    print(" ", " ".join(cmd))
    if args.dry_run:
        return 0

    env = os.environ.copy()
    env.setdefault("PYTHONPATH", str(ROOT / "evalscope"))
    result = subprocess.run(cmd, cwd=ROOT, env=env)
    return result.returncode


def main():
    parser = argparse.ArgumentParser(description="复现 Kimi K3 report 中可本地运行的 benchmark")
    parser.add_argument("--model", default=os.getenv("EVAL_MODEL", "kimi-k3"), help="模型名")
    parser.add_argument("--api-url", default=os.getenv("EVAL_API_URL", "https://api.example.com/v1"), help="API URL")
    parser.add_argument("--api-key", default=os.getenv("EVAL_API_KEY", ""), help="API key")
    parser.add_argument("--limit", type=int, default=None, help="每个 benchmark 限制样本数，默认全量")
    parser.add_argument("--output-dir", default=str(ROOT / "output"), help="输出根目录")
    parser.add_argument("--folder-name", default="k3-report-test", help="输出文件夹名")
    parser.add_argument("--config", default=str(ROOT / "config" / "dpv4-int8_nothinking.yaml"), help="评测配置 YAML")
    parser.add_argument("--categories", default="", help="逗号分隔类别，如 Coding,Vision")
    parser.add_argument("--datasets", default="", help="逗号分隔 dataset，只跑指定几个")
    parser.add_argument("--thinking", action="store_true", help="启用 thinking 模式")
    parser.add_argument("--thinking-budget-tokens", type=int, default=None, help="thinking budget tokens")
    parser.add_argument("--dry-run", action="store_true", help="只打印命令，不执行")
    args = parser.parse_args()

    print_support_matrix()

    if args.dry_run:
        print("\n[Dry-run mode] 以下命令将被执行：\n")

    # 选择要跑的 benchmark
    selected = []
    categories = [c.strip() for c in args.categories.split(",") if c.strip()]
    explicit_datasets = [d.strip() for d in args.datasets.split(",") if d.strip()]

    for category, items in BENCHMARKS.items():
        if categories and category not in categories:
            continue
        for item in items:
            if explicit_datasets and item["dataset"] not in explicit_datasets:
                continue
            selected.append((category, item))

    if explicit_datasets:
        # 允许直接传 dataset 名，即使不在 BENCHMARKS 映射里
        known = {item["dataset"] for items in BENCHMARKS.values() for item in items}
        for d in explicit_datasets:
            if d not in known:
                selected.append(("Custom", {"report_name": d, "dataset": d, "deps": []}))

    if not selected:
        print("\n没有选中任何 benchmark，请调整 --categories 或 --datasets")
        return

    print(f"\n将运行 {len(selected)} 个 benchmark ...")
    failed = []
    for category, item in selected:
        rc = run_one(item, args)
        if rc != 0:
            failed.append(item["report_name"])

    print("\n" + "=" * 70)
    print("完成")
    if failed:
        print(f"失败: {failed}")
    else:
        print("全部成功")
    print(f"结果目录: {args.output_dir}/{args.folder_name}")


if __name__ == "__main__":
    main()