Video-Tree-TRM5/app/question_gen/synthesizer.py

"""赛题合成核心逻辑 — 节点采样、prompt 构造、去重。

纯函数为主，异步编排仅 generate_one。
通过 DI 接收 VLMProvider / EmbeddingProvider，不 import adapters/。
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import random

    from app.tree.index import L2Node, L3Node, TreeIndex


@dataclass(frozen=True)
class AnchorContext:
    """锚节点上下文——生成单道题所需的全部素材。

    属性:
        node_id: 锚节点 ID。
        card_text: 锚节点 card 序列化文本。
        frame_paths: 帧图片路径列表。
        subtitle: 对应字幕（可空）。
        distractor_texts: 同视频其他节点摘要（供 VLM 生成干扰项）。
    """

    node_id: str
    card_text: str
    frame_paths: list[str]
    subtitle: str
    distractor_texts: list[str]


@dataclass(frozen=True)
class TaskTypeSpec:
    """题型的生成规格。

    属性:
        level: 锚定层级（"L3" / "L2" / "L1" / "L1-L2"）。
        needs_frames: 是否必须提供帧图。
        frame_count: 帧数范围描述（如 "1", "2-3", "0-1"）。
        context_fields: 需要提取的 card 字段元组。
    """

    level: str
    needs_frames: bool
    frame_count: str
    context_fields: tuple[str, ...]


# ---------------------------------------------------------------------------
# 12 种 Video-MME 题型 → 树层级 + 生成规格映射
# ---------------------------------------------------------------------------

TASK_TYPE_LEVEL_MAP: dict[str, TaskTypeSpec] = {
    # --- L3 单帧题型 ---
    "Object Recognition": TaskTypeSpec("L3", True, "1", ("frame_summary",)),
    "Attribute Perception": TaskTypeSpec("L3", True, "1", ("frame_summary",)),
    "OCR Problems": TaskTypeSpec("L3", True, "1", ("frame_summary",)),
    "Spatial Reasoning": TaskTypeSpec("L3", True, "1", ("frame_summary", "spatial_layout")),
    "Spatial Perception": TaskTypeSpec("L3", True, "1", ("frame_summary",)),
    # --- L2 多帧 / 事件级题型 ---
    "Action Recognition": TaskTypeSpec("L2", True, "2-3", ("event_description",)),
    "Action Reasoning": TaskTypeSpec("L2", True, "2-3", ("event_description",)),
    "Counting Problem": TaskTypeSpec("L2", True, "2-3", ("event_description",)),
    "Temporal Perception": TaskTypeSpec("L2", False, "0-1", ("event_description", "time_range")),
    # --- L1 / 跨层级题型 ---
    "Temporal Reasoning": TaskTypeSpec("L1", True, "per-L2", ("scene_summary",)),
    "Information Synopsis": TaskTypeSpec("L1", True, "per-L2", ("scene_summary",)),
    "Object Reasoning": TaskTypeSpec("L1-L2", True, "per-L2", ("event_description",)),
}


# ---------------------------------------------------------------------------
# 内部辅助函数
# ---------------------------------------------------------------------------


def _serialize_l3_card(l3: L3Node, context_fields: tuple[str, ...]) -> str:
    """将 L3 节点 card 按 context_fields 序列化为可读文本。

    参数:
        l3: L3 节点。
        context_fields: 需提取的字段名元组。

    返回:
        多行 "field: value" 格式的文本。
    """
    parts: list[str] = []
    for fld in context_fields:
        val = getattr(l3.card, fld, None)
        if val is not None:
            parts.append(f"{fld}: {val}")
    return "\n".join(parts)


def _l2_time_range_str(l2: L2Node) -> str:
    """将 L2 的 time_range 格式化为可读字符串。

    参数:
        l2: L2 节点。

    返回:
        "time_range: (start, end)" 格式，或 "time_range: unknown"。
    """
    if l2.time_range is not None:
        return f"time_range: ({l2.time_range[0]:.2f}, {l2.time_range[1]:.2f})"
    return "time_range: unknown"


def _representative_frame(l2: L2Node) -> str | None:
    """取 L2 的代表帧路径——第一个有 frame_path 的 L3 子节点。

    参数:
        l2: L2 节点。

    返回:
        帧路径字符串，或 None（无可用帧时）。
    """
    for l3 in l2.children:
        if l3.frame_path:
            return l3.frame_path
    return None


def _sort_l2_by_time(l2_nodes: list[L2Node]) -> list[L2Node]:
    """按 time_range 升序排列 L2 节点（None 排末尾）。

    参数:
        l2_nodes: 待排序的 L2 节点列表。

    返回:
        排序后的新列表（不修改原列表）。
    """
    return sorted(
        l2_nodes,
        key=lambda n: n.time_range[0] if n.time_range is not None else float("inf"),
    )


# ---------------------------------------------------------------------------
# 各层级采样策略
# ---------------------------------------------------------------------------


def _sample_l3(
    tree: TreeIndex,
    task_type: str,
    spec: TaskTypeSpec,
    used_node_ids: set[str],
    rng: random.Random,
) -> AnchorContext:
    """L3 层级锚节点采样。

    收集全部 L3 节点，排除已用节点，随机选取一个。

    参数:
        tree: 三层树索引。
        task_type: 题型名称。
        spec: 题型规格。
        used_node_ids: 已用节点 ID 集合。
        rng: 随机数生成器。

    返回:
        AnchorContext 实例。

    异常:
        ValueError: 候选 L3 节点不足。
    """
    # Phase 1: 收集所有 L3 候选（必须有 frame_path）
    candidates: list[tuple[L3Node, L2Node]] = []
    for root in tree.roots:
        for l2 in root.children:
            for l3 in l2.children:
                if l3.id not in used_node_ids and l3.frame_path:
                    candidates.append((l3, l2))

    if not candidates:
        raise ValueError(f"锚节点不足: {task_type} 无可用 L3 节点（需具备 frame_path）")

    # Phase 2: 随机选取
    chosen_l3, parent_l2 = rng.choice(candidates)

    # Phase 3: 构造上下文（frame_path 已在候选过滤中保证非 None）
    card_text = _serialize_l3_card(chosen_l3, spec.context_fields)
    frame_paths = [chosen_l3.frame_path]  # type: ignore[list-item]
    subtitle = chosen_l3.subtitle or ""

    # Phase 4: 干扰项——整棵树中其他 L3 的 frame_summary
    distractor_texts = [
        l3.card.frame_summary
        for root in tree.roots
        for l2 in root.children
        for l3 in l2.children
        if l3.id != chosen_l3.id
    ]

    return AnchorContext(
        node_id=chosen_l3.id,
        card_text=card_text,
        frame_paths=frame_paths,
        subtitle=subtitle,
        distractor_texts=distractor_texts,
    )


def _sample_l2(
    tree: TreeIndex,
    task_type: str,
    spec: TaskTypeSpec,
    used_node_ids: set[str],
    rng: random.Random,
) -> AnchorContext:
    """L2 层级锚节点采样（含 Temporal Perception 特殊处理）。

    普通 L2 题型：随机选 1 个 L2，取 2-3 个子 L3 帧。
    Temporal Perception：0-1 帧，card_text 必含 time_range。

    参数:
        tree: 三层树索引。
        task_type: 题型名称。
        spec: 题型规格。
        used_node_ids: 已用节点 ID 集合。
        rng: 随机数生成器。

    返回:
        AnchorContext 实例。

    异常:
        ValueError: 候选 L2 节点不足。
    """
    # Phase 1: 收集所有 L2 候选
    all_l2: list[L2Node] = []
    for root in tree.roots:
        for l2 in root.children:
            if l2.id not in used_node_ids:
                all_l2.append(l2)

    if not all_l2:
        raise ValueError(f"锚节点不足: {task_type} 无可用 L2 节点")

    # Phase 2: 随机选取
    chosen_l2 = rng.choice(all_l2)

    is_temporal_perception = task_type == "Temporal Perception"

    # Phase 3: 帧路径
    if is_temporal_perception:
        # 0-1 帧：有子节点则取 1 帧，否则 0 帧
        frame_paths: list[str] = []
        if chosen_l2.children:
            first_frame = chosen_l2.children[0].frame_path
            if first_frame:
                frame_paths = [first_frame]
    else:
        # 普通 L2：随机采样 2-3 个 L3 帧
        children_with_frames = [l3 for l3 in chosen_l2.children if l3.frame_path]
        if len(children_with_frames) < 2:
            raise ValueError(
                f"锚节点不足: {task_type} 需要 >=2 个子帧，"
                f"但 {chosen_l2.id} 仅有 {len(children_with_frames)} 个可用帧"
            )
        n_frames = min(rng.randint(2, 3), len(children_with_frames))
        sampled = rng.sample(children_with_frames, n_frames)
        frame_paths = [l3.frame_path for l3 in sampled if l3.frame_path]

    # Phase 4: card_text
    card_text = f"event_description: {chosen_l2.card.event_description}"
    if is_temporal_perception:
        card_text += f"\n{_l2_time_range_str(chosen_l2)}"

    # Phase 5: 字幕（L2 无自身字幕，取首个子 L3 字幕）
    subtitle = ""
    if chosen_l2.children and chosen_l2.children[0].subtitle:
        subtitle = chosen_l2.children[0].subtitle

    # Phase 6: 干扰项——整棵树中其他 L2 的 event_description
    distractor_texts = [
        l2.card.event_description
        for root in tree.roots
        for l2 in root.children
        if l2.id != chosen_l2.id
    ]

    return AnchorContext(
        node_id=chosen_l2.id,
        card_text=card_text,
        frame_paths=frame_paths,
        subtitle=subtitle,
        distractor_texts=distractor_texts,
    )


def _sample_l1(
    tree: TreeIndex,
    task_type: str,
    spec: TaskTypeSpec,
    used_node_ids: set[str],
    rng: random.Random,
) -> AnchorContext:
    """L1 层级锚节点采样（Temporal Reasoning / Information Synopsis）。

    Information Synopsis：使用目标 L1 下全部 L2 子节点。
    Temporal Reasoning：严格要求 >=3 个 L2 子节点，不足则抛 ValueError。
    L2 按 time_range 升序排列，每个 L2 取一帧代表。

    参数:
        tree: 三层树索引。
        task_type: 题型名称。
        spec: 题型规格。
        used_node_ids: 已用节点 ID 集合。
        rng: 随机数生成器。

    返回:
        AnchorContext 实例。

    异常:
        ValueError: 候选 L1 节点不足，或 Temporal Reasoning 的 L2 子节点 <3。
    """
    # Phase 1: 收集可用 L1
    candidates = [r for r in tree.roots if r.id not in used_node_ids]
    if not candidates:
        raise ValueError(f"锚节点不足: {task_type} 无可用 L1 节点")

    # Phase 2: 随机选取
    chosen_l1 = rng.choice(candidates)

    # Phase 3: 选定 L2 子集
    if task_type == "Information Synopsis":
        # 必须使用全部 L2
        selected_l2 = list(chosen_l1.children)
    else:
        # Temporal Reasoning：严格要求 >=3 个 L2
        if len(chosen_l1.children) < 3:
            raise ValueError(
                f"锚节点不足: {task_type} 需要 >=3 个 L2 子节点，"
                f"但 {chosen_l1.id} 仅有 {len(chosen_l1.children)} 个"
            )
        if len(chosen_l1.children) == 3:
            selected_l2 = list(chosen_l1.children)
        else:
            selected_l2 = rng.sample(chosen_l1.children, rng.randint(3, len(chosen_l1.children)))

    # Phase 4: 按 time_range 升序排列
    selected_l2 = _sort_l2_by_time(selected_l2)

    # Phase 5: card_text（场景摘要）
    card_text = f"scene_summary: {chosen_l1.card.scene_summary}"

    # Phase 6: 帧路径——每个 L2 取一帧代表
    frame_paths: list[str] = []
    for l2 in selected_l2:
        rep = _representative_frame(l2)
        if rep:
            frame_paths.append(rep)

    # Phase 7: 字幕（L1 无字幕）
    subtitle = ""

    # Phase 8: 干扰项——其他 L1 的 scene_summary
    distractor_texts = [r.card.scene_summary for r in tree.roots if r.id != chosen_l1.id]

    return AnchorContext(
        node_id=chosen_l1.id,
        card_text=card_text,
        frame_paths=frame_paths,
        subtitle=subtitle,
        distractor_texts=distractor_texts,
    )


def _sample_l1_l2(
    tree: TreeIndex,
    task_type: str,
    spec: TaskTypeSpec,
    used_node_ids: set[str],
    rng: random.Random,
) -> AnchorContext:
    """L1-L2 跨层级锚节点采样（Object Reasoning）。

    从全部 L2 中随机选 2-3 个，按 time_range 排序，
    card_text 为各 L2 的 event_description 拼接。

    参数:
        tree: 三层树索引。
        task_type: 题型名称。
        spec: 题型规格。
        used_node_ids: 已用节点 ID 集合。
        rng: 随机数生成器。

    返回:
        AnchorContext 实例。

    异常:
        ValueError: 候选 L2 节点不足。
    """
    # Phase 1: 收集全部 L2
    all_l2: list[L2Node] = []
    for root in tree.roots:
        for l2 in root.children:
            if l2.id not in used_node_ids:
                all_l2.append(l2)

    if len(all_l2) < 2:
        raise ValueError(
            f"锚节点不足: {task_type} 需要 >=2 个 L2 节点，但仅有 {len(all_l2)} 个可用"
        )

    # Phase 2: 随机选 2-3 个
    n_pick = min(rng.randint(2, 3), len(all_l2))
    selected = rng.sample(all_l2, n_pick)

    # Phase 3: 按 time_range 升序排列
    selected = _sort_l2_by_time(selected)

    # Phase 4: card_text = 各 L2 event_description 拼接
    card_text = "\n".join(f"event_description: {l2.card.event_description}" for l2 in selected)

    # Phase 5: 帧路径——每个 L2 取一帧代表
    frame_paths: list[str] = []
    for l2 in selected:
        rep = _representative_frame(l2)
        if rep:
            frame_paths.append(rep)

    # Phase 6: 字幕
    subtitle = ""

    # Phase 7: 干扰项——未被选中的 L2 的 event_description
    selected_ids = {l2.id for l2 in selected}
    distractor_texts = [l2.card.event_description for l2 in all_l2 if l2.id not in selected_ids]

    # 使用第一个被选中节点的 ID 作为锚节点 ID
    anchor_id = selected[0].id

    return AnchorContext(
        node_id=anchor_id,
        card_text=card_text,
        frame_paths=frame_paths,
        subtitle=subtitle,
        distractor_texts=distractor_texts,
    )


# ---------------------------------------------------------------------------
# 公开接口
# ---------------------------------------------------------------------------


def sample_anchor(
    tree: TreeIndex,
    task_type: str,
    used_node_ids: set[str],
    rng: random.Random,
) -> AnchorContext:
    """根据题型从视频树中采样锚节点及上下文素材。

    依据 TASK_TYPE_LEVEL_MAP 中的层级规格，分发到对应的层级采样策略。
    每种层级有不同的帧选取、card 序列化和干扰项收集逻辑。

    参数:
        tree: 三层树索引。
        task_type: 12 种 Video-MME 题型之一。
        used_node_ids: 本轮已用节点 ID 集合（避免重复采样）。
        rng: 可控随机数生成器（保证可复现）。

    返回:
        AnchorContext 实例，包含锚节点 ID、card 文本、帧路径、字幕和干扰项。

    异常:
        KeyError: task_type 不在 TASK_TYPE_LEVEL_MAP 中。
        ValueError: 候选节点不足（全部被 used_node_ids 排除）。
    """
    spec = TASK_TYPE_LEVEL_MAP[task_type]

    if spec.level == "L3":
        return _sample_l3(tree, task_type, spec, used_node_ids, rng)
    elif spec.level == "L2":
        return _sample_l2(tree, task_type, spec, used_node_ids, rng)
    elif spec.level == "L1":
        return _sample_l1(tree, task_type, spec, used_node_ids, rng)
    elif spec.level == "L1-L2":
        return _sample_l1_l2(tree, task_type, spec, used_node_ids, rng)
    else:
        raise ValueError(f"未知层级: {spec.level}")