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

从 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

    from core.types import QuestionUnit

_LEGACY_DEFAULT_DIFFICULTY = "medium"


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

    video_id 优先使用题目 JSON 中的 ``video_id`` 字段；若缺失则回退到
    文件名（不含扩展名）。Video-MME benchmark 按视频拆文件（文件名即
    video_id），v2 生成题把多视频题目合并在单个 JSON 中（每条记录自带
    ``video_id``），两种格式均兼容。

    pair 契约字段（``pair_id`` / ``question_role`` / ``flip_axis`` / ``unit_id``）
    用 ``.get`` 读取：旧 benchmark 无这些键时退化为 single（``question_role``
    默认 "single"，``unit_id`` 留空由 __post_init__ 回填为 question_id），
    保证历史题库可无缝加载。

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

    返回:
        按文件名排序加载的题目列表。
    """
    results: list[GeneratedQuestion] = []
    for path in sorted(questions_dir.glob("*.json")):
        fallback_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=qa.get("video_id", fallback_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),
                    family=qa.get("family"),
                    skill_target=qa.get("skill_target"),
                    difficulty_steps=qa.get("difficulty_steps"),
                    sub_pattern=qa.get("sub_pattern"),
                    # pair 契约字段：旧 benchmark 无这些键时按 single 默认兜底，
                    # unit_id 留空交由 GeneratedQuestion.__post_init__ 回填。
                    pair_id=qa.get("pair_id"),
                    question_role=qa.get("question_role", "single"),
                    flip_axis=qa.get("flip_axis"),
                    unit_id=qa.get("unit_id", ""),
                )
            )
    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: 候选题目全集（single 与孪生对成员可混含）。
        correctness: question_id -> 基线是否答对（单元级正确性取成员 AND）。
        size: 采样单元总量（single 计 1、pair 计 1）。
        correct_ratio: 采样中"基线答对"单元的占比；None 表示自然分布。
        task_types: 限定题型；None 表示不限。
        seed: 随机种子，保证可复现。
        min_per_class: 每个题型补足到的单元下限；None 表示不补足。

    返回:
        采样后的题目列表（pair 单元展开为原始的两道题）。

    异常:
        ValueError: 自然分布时单元池不足 size，或分层时某层单元不足。

    关键实现:
        以 **QuestionUnit 为采样原子**（single 计 1、pair 计 1），size /
        correct_ratio / min_per_class 均按 unit 计数，孪生对两题永不被劈开。
        采样完成后 flatten_units 展开回逐题列表。纯 single 输入时 build_units
        与题目一一对应、顺序不变，rng 消耗与旧逐题实现完全一致（字节级回归）。

        build_units / flatten_units 采用函数内延迟导入：loader 属 question_gen，
        question_units 属 harness，而 harness 包初始化会反向 import question_gen，
        模块级导入将触发循环依赖（沿用 adversarial_filter 的既有做法）。
    """
    from app.harness.question_units import build_units, flatten_units

    rng = random.Random(seed)
    units = build_units(questions)
    pool = [u for u in units if task_types is None or u.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 flatten_units(sampled)


def _unit_correct(unit: QuestionUnit, correctness: dict[str, bool]) -> bool:
    """单元级正确性：成员全部答对才算对（缺失按 False，宽松口径）。

    参数:
        unit: 目标单元（single 1 题，pair 2 题）。
        correctness: question_id -> 基线是否答对。

    返回:
        pair 走双向 AND、single 即单题正确性；任一成员缺失或答错即 False。
    """
    return all(correctness.get(q.question_id, False) for q in unit.questions)


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

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

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

    异常:
        ValueError: 对单元或错单元层不足。
    """
    correct = [u for u in pool if _unit_correct(u, correctness)]
    wrong = [u for u in pool if not _unit_correct(u, correctness)]
    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[QuestionUnit],
    pool: list[QuestionUnit],
    min_per_class: int,
    rng: random.Random,
) -> list[QuestionUnit]:
    """对候选池中出现的每个题型，将采样单元补足到 min_per_class 个。

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

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

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