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

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

from __future__ import annotations

import json
import random
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


def stratified_sample(
    questions: list[GeneratedQuestion],
    correctness: dict[str, bool],
    size: int,
    correct_ratio: float | None,
    task_types: list[str] | None,
    seed: int,
    min_per_class: int | None,
) -> list[GeneratedQuestion]:
    """按题型过滤后采样 size 道题，可选按对错比例分层并按题型保底。

    参数:
        questions: 候选题目全集。
        correctness: question_id -> 基线是否答对。
        size: 采样总量。
        correct_ratio: 采样中"基线答对"题的占比；None 表示自然分布。
        task_types: 限定题型；None 表示不限。
        seed: 随机种子，保证可复现。
        min_per_class: 每个题型补足到的下限；None 表示不补足。

    返回:
        采样后的题目列表。

    异常:
        ValueError: 自然分布时池不足 size，或分层时某层题目不足。
    """
    rng = random.Random(seed)
    pool = [q for q in questions if task_types is None or q.task_type in task_types]

    if correct_ratio is None:
        if len(pool) < size:
            raise ValueError(f"自然分布采样不足: 需 {size} 道, 实有 {len(pool)} 道")
        sampled = rng.sample(pool, size)
    else:
        sampled = _ratio_stratified_sample(pool, correctness, size, correct_ratio, rng)

    if min_per_class is not None:
        sampled = _backfill_per_class(sampled, pool, min_per_class, rng)
    return sampled


def _ratio_stratified_sample(
    pool: list[GeneratedQuestion],
    correctness: dict[str, bool],
    size: int,
    correct_ratio: float,
    rng: random.Random,
) -> list[GeneratedQuestion]:
    """按对错比例分层采样：对题占 correct_ratio，其余为错题。

    参数:
        pool: 题型过滤后的候选题。
        correctness: question_id -> 基线是否答对。
        size: 采样总量。
        correct_ratio: 对题占比。
        rng: 随机数发生器。

    返回:
        采样后的题目列表（对题在前、错题在后）。

    异常:
        ValueError: 对题或错题层不足。
    """
    correct = [q for q in pool if correctness.get(q.question_id, False)]
    wrong = [q for q in pool if not correctness.get(q.question_id, False)]
    n_correct = round(size * correct_ratio)
    n_wrong = size - n_correct
    if len(correct) < n_correct or len(wrong) < n_wrong:
        raise ValueError(
            f"分层不足: 需对{n_correct}/错{n_wrong}, 实有对{len(correct)}/错{len(wrong)}"
        )
    return rng.sample(correct, n_correct) + rng.sample(wrong, n_wrong)


def _backfill_per_class(
    sampled: list[GeneratedQuestion],
    pool: list[GeneratedQuestion],
    min_per_class: int,
    rng: random.Random,
) -> list[GeneratedQuestion]:
    """对候选池中出现的每个题型，将采样结果补足到 min_per_class 道。

    遍历对象是候选池 pool 里出现的全部题型（非仅 sampled 命中的），
    保证任意稀疏题型都能拿到足额样本。

    参数:
        sampled: 主采样结果（不修改，返回新列表）。
        pool: 候选题全集（补足来源 + 题型枚举来源）。
        min_per_class: 每个题型的下限。
        rng: 随机数发生器。

    返回:
        补足后的题目列表。
    """
    selected_ids = {q.question_id for q in sampled}
    result = list(sampled)
    counts: dict[str, int] = {}
    for q in sampled:
        counts[q.task_type] = counts.get(q.task_type, 0) + 1
    ordered_task_types: dict[str, None] = {}
    for q in pool:
        ordered_task_types.setdefault(q.task_type, None)
    for task_type in ordered_task_types:
        deficit = min_per_class - counts.get(task_type, 0)
        if deficit <= 0:
            continue
        candidates = [
            q for q in pool if q.task_type == task_type and q.question_id not in selected_ids
        ]
        take = rng.sample(candidates, min(deficit, len(candidates)))
        for q in take:
            selected_ids.add(q.question_id)
            result.append(q)
    return result