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

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

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

from __future__ import annotations

import contextlib
import json
import re
from dataclasses import dataclass
from typing import TYPE_CHECKING

import numpy as np
from loguru import logger

if TYPE_CHECKING:
    import random
    from collections.abc import Callable

    from app.tree.index import L2Node, L3Node, TreeIndex
    from core.protocols import VLMProvider
    from core.types import GeneratedQuestion


@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.card.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 = chosen_l2.card.subtitle or ""
    if not subtitle and chosen_l2.children and chosen_l2.children[0].card.subtitle:
        subtitle = chosen_l2.children[0].card.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}")


# ---------------------------------------------------------------------------
# Prompt 构造与 VLM 响应解析
# ---------------------------------------------------------------------------

_VALID_ANSWERS = frozenset({"A", "B", "C", "D"})


def build_generation_prompt(
    task_type: str,
    anchor: AnchorContext,
    exemplars: list[GeneratedQuestion],
) -> tuple[list[dict[str, str]], list[str]]:
    """组装 VLM 出题 prompt。

    构造 OpenAI 格式的 messages 列表和帧图片路径列表，
    供 VLMProvider.chat_with_images 直接消费。

    参数:
        task_type: 题型名称（如 "Object Recognition"）。
        anchor: 锚节点上下文（card_text, subtitle, distractor_texts, frame_paths）。
        exemplars: 少样本示例列表（可为空）。

    返回:
        (messages, image_paths) — messages 为 OpenAI 格式消息列表，
        image_paths 为帧图片路径列表，直接喂给 VLMProvider.chat_with_images。
    """
    # Phase 1: 构造 system message
    system_parts: list[str] = [
        "你是一个视频理解题目生成器。",
        f"题型: {task_type}",
        "约束:",
        "- 题目必须基于提供的节点内容",
        "- 干扰选项应来自其他节点的信息",
        "- 生成风格应与示例保持一致",
        '- 以 JSON 格式返回: {"question": "...", "options": ["A. ...", "B. ...", "C. ...", "D. ..."], "answer": "A/B/C/D"}',
    ]

    # Phase 2: 加入 few-shot 示例
    if exemplars:
        system_parts.append("\n示例:")
        for i, ex in enumerate(exemplars, 1):
            system_parts.append(f"  示例 {i}:")
            system_parts.append(f"    question: {ex.question}")
            system_parts.append(f"    options: {list(ex.options)}")
            system_parts.append(f"    answer: {ex.answer}")

    system_content = "\n".join(system_parts)

    # Phase 3: 构造 user message
    user_parts: list[str] = [f"节点内容:\n{anchor.card_text}"]

    if anchor.subtitle:
        user_parts.append(f"\n字幕:\n{anchor.subtitle}")

    if anchor.distractor_texts:
        user_parts.append("\n干扰项来源节点摘要:")
        for dt in anchor.distractor_texts:
            user_parts.append(f"- {dt}")

    user_content = "\n".join(user_parts)

    messages = [
        {"role": "system", "content": system_content},
        {"role": "user", "content": user_content},
    ]

    return messages, list(anchor.frame_paths)


def parse_vlm_response(
    raw: str,
    video_id: str,
    task_type: str,
    seq: int,
) -> dict:
    """解析 VLM 返回的 JSON → 部分字段字典。

    尝试直接解析 JSON；若失败，从 markdown 代码块中提取后重试。
    校验必需字段、选项数量和答案合法性。

    参数:
        raw: VLM 原始返回文本。
        video_id: 所属视频标识。
        task_type: 题型名称（用于错误消息）。
        seq: 序列号，用于生成 question_id。

    返回:
        {"question_id": "gen-{video_id}-{seq:03d}", "question": ..., "options": [...], "answer": ...}
        调用方（generate_one）补齐 source_nodes/difficulty 后构造 GeneratedQuestion。

    异常:
        ValueError: JSON 解析失败、缺必需字段、options 非 4 项、answer 不在 A-D。
    """
    # Phase 1: 尝试直接解析 JSON
    data = None
    with contextlib.suppress(json.JSONDecodeError):
        data = json.loads(raw)

    # Phase 2: 从 markdown 代码块提取 JSON
    if data is None:
        match = re.search(r"```(?:json)?\s*\n?(.*?)\n?\s*```", raw, re.DOTALL)
        if match:
            with contextlib.suppress(json.JSONDecodeError):
                data = json.loads(match.group(1))

    if data is None:
        raise ValueError(f"VLM 返回无法解析为 JSON: {raw[:200]}")

    # Phase 3: 校验必需字段
    required = ("question", "options", "answer")
    missing = [f for f in required if f not in data]
    if missing:
        raise ValueError(f"VLM 返回缺少必需字段 {missing}: {raw[:200]}")

    # Phase 4: options 必须恰好 4 项
    options = data["options"]
    if not isinstance(options, list) or len(options) != 4:
        raise ValueError(
            f"options 必须恰好 4 项，实际 {len(options) if isinstance(options, list) else type(options).__name__}: {raw[:200]}"
        )

    # Phase 5: answer 必须是 A-D
    answer = data["answer"]
    if answer not in _VALID_ANSWERS:
        raise ValueError(f"answer 必须是 A/B/C/D 之一，实际 '{answer}': {raw[:200]}")

    return {
        "question_id": f"gen-{video_id}-{task_type.lower().replace(' ', '_')}-{seq:03d}",
        "question": data["question"],
        "options": list(options),
        "answer": answer,
    }


# ---------------------------------------------------------------------------
# Embedding 去重
# ---------------------------------------------------------------------------


def is_duplicate(
    question_text: str,
    pool_embeddings: np.ndarray,
    embed_fn: Callable[[str | list[str]], np.ndarray],
    threshold: float,
) -> bool:
    """embedding 去重判定。

    参数:
        question_text: 待检查的题目文本。
        pool_embeddings: 已有题目的 embedding 矩阵 [N, D]（L2 归一化）。
        embed_fn: 文本嵌入函数，返回 [N, D] ndarray（L2 归一化）。
        threshold: 余弦相似度阈值。

    返回:
        True 表示与池中某题重复。空池永远返回 False。
    """
    if pool_embeddings.shape[0] == 0:
        return False

    query = embed_fn(question_text)  # [1, D]
    query = query.squeeze(0)  # [D]
    similarities = pool_embeddings @ query  # [N]
    return bool(np.max(similarities) >= threshold)


# ---------------------------------------------------------------------------
# 单题生成
# ---------------------------------------------------------------------------


async def generate_one(
    vlm: VLMProvider,
    tree: TreeIndex,
    video_id: str,
    task_type: str,
    seq: int,
    *,
    exemplars: list[GeneratedQuestion],
    used_node_ids: set[str],
    max_retries: int,
    rng: random.Random,
    session_id: str,
) -> GeneratedQuestion | None:
    """生成单道候选题（不含去重——去重在调用方汇总点原子执行）。

    循环最多 max_retries 次尝试生成。每次尝试：
    1. 采样锚节点
    2. 构造 prompt
    3. 调用 VLM
    4. 解析响应
    5. 构造 GeneratedQuestion

    返回 None 表示耗尽重试。

    参数:
        vlm: VLM 调用端口。
        tree: 三层树索引。
        video_id: 所属视频标识。
        task_type: 12 种 Video-MME 题型之一。
        seq: 序列号，用于生成 question_id。
        exemplars: 少样本示例列表。
        used_node_ids: 已用节点 ID 集合。
        max_retries: 最大重试次数。
        rng: 可控随机数生成器。
        session_id: 会话 ID（传递给 VLM 遥测）。

    返回:
        GeneratedQuestion 实例，或 None（耗尽重试）。
    """
    from core.types import GeneratedQuestion as _GeneratedQuestion

    for attempt in range(max_retries):
        try:
            # Phase 1: 采样锚节点
            anchor = sample_anchor(tree, task_type, used_node_ids, rng)

            # Phase 2: 构造 prompt
            messages, images = build_generation_prompt(task_type, anchor, exemplars)

            # Phase 3: 调用 VLM
            response = await vlm.chat_with_images(
                messages,
                images,
                session_id=session_id,
            )

            # Phase 4: 解析响应
            parsed = parse_vlm_response(response.content, video_id, task_type, seq)

            # Phase 5: 构造 GeneratedQuestion
            return _GeneratedQuestion(
                question_id=parsed["question_id"],
                video_id=video_id,
                task_type=task_type,
                question=parsed["question"],
                options=tuple(parsed["options"]),
                answer=parsed["answer"],
                source_nodes=(anchor.node_id,),
                difficulty="medium",
            )
        except (ValueError, KeyError) as exc:
            logger.warning(
                "generate_one 尝试 {}/{} 失败 ({}): {}",
                attempt + 1,
                max_retries,
                task_type,
                exc,
            )
            continue

    logger.warning(
        "generate_one 耗尽 {} 次重试 (video={}, task_type={})",
        max_retries,
        video_id,
        task_type,
    )
    return None