"""题目加载与分层采样。

从 benchmark JSON 目录加载题目，提供按对错比例的分层采样。
对应训练循环中的 DataLoader 角色。
"""

from __future__ import annotations

import json
from typing import TYPE_CHECKING

from core.types import GeneratedQuestion

if TYPE_CHECKING:
    from pathlib import Path

_LEGACY_DEFAULT_DIFFICULTY = "medium"


def load_benchmark(questions_dir: Path) -> list[GeneratedQuestion]:
    """从 benchmark JSON 目录加载题目列表。

    每个 JSON 文件以文件名（不含扩展名）作为 video_id，
    文件内容为题目数组。

    参数:
        questions_dir: 包含 *.json 文件的目录路径。

    返回:
        按文件名排序加载的题目列表。
    """
    results: list[GeneratedQuestion] = []
    for path in sorted(questions_dir.glob("*.json")):
        video_id = path.stem
        with open(path, encoding="utf-8") as f:
            qa_list: list[dict] = json.load(f)
        for qa in qa_list:
            results.append(
                GeneratedQuestion(
                    question_id=qa["question_id"],
                    video_id=video_id,
                    task_type=qa["task_type"],
                    question=qa["question"],
                    options=tuple(qa["options"]),
                    answer=qa["answer"],
                    source_nodes=tuple(qa.get("source_nodes", ())),
                    difficulty=qa.get("difficulty", _LEGACY_DEFAULT_DIFFICULTY),
                )
            )
    return results