"""题族规格声明 — 定义 5 大问题家族及其采样、泄漏检测、提示模板约束。

每个 QuestionFamilySpec 对应一种失败机制（skill_target M1–M5），
由 get_family_for_slot 在出题时按权重分配。
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import random


@dataclass(frozen=True)
class LeakTestProfile:
    """泄漏测试配置 — 定义快捷答题捷径类型与通过阈值。

    Attributes:
        shortcut_type: 捷径类型标识。
        probe_template: store/prompts/question_gen/ 下的探测模板文件名。
        pass_threshold: 通过阈值（0–1），低于此值视为存在泄漏。
    """

    shortcut_type: str
    probe_template: str
    pass_threshold: float


@dataclass(frozen=True)
class SamplingConstraint:
    """采样约束 — 对树节点的最低要求。

    Attributes:
        min_subtitles: 最少字幕段数。
        min_l3_nodes: 最少 L3 节点数。
        require_frames: 是否要求帧图像可用。
        cross_l2_span: 是否要求跨 L2 段采样。
    """

    min_subtitles: int
    min_l3_nodes: int
    require_frames: bool
    cross_l2_span: bool


@dataclass(frozen=True)
class QuestionFamilySpec:
    """问题家族规格 — 一个家族的完整声明。

    Attributes:
        name: 家族标识名（如 "RETRIEVAL"）。
        skill_target: 目标失败机制编号（M1–M5）。
        sampling: 采样约束。
        legal_task_types: 该家族合法的任务类型集合（frozenset）。
        leak_profile: 泄漏测试配置。
        prompt_template: 出题 prompt 模板文件名。
    """

    name: str
    skill_target: str
    sampling: SamplingConstraint
    legal_task_types: frozenset[str]
    leak_profile: LeakTestProfile
    prompt_template: str


# ---------------------------------------------------------------------------
# 5 大家族实例
# ---------------------------------------------------------------------------

RETRIEVAL_FAMILY = QuestionFamilySpec(
    name="RETRIEVAL",
    skill_target="M1",
    sampling=SamplingConstraint(
        min_subtitles=2,
        min_l3_nodes=3,
        require_frames=False,
        cross_l2_span=False,
    ),
    legal_task_types=frozenset(
        [
            "Action Recognition",
            "Action Reasoning",
            "Action Prediction",
            "Action Sequence",
            "Object Recognition",
            "Object Reasoning",
            "Object Interaction",
            "Scene Understanding",
            "Event Reasoning",
            "Causal Reasoning",
            "Temporal Reasoning",
            "Spatial Reasoning",
        ]
    ),
    leak_profile=LeakTestProfile(
        shortcut_type="temporal_proximity",
        probe_template="gate_leak_retrieval.md",
        pass_threshold=0.6,
    ),
    prompt_template="retrieval.md",
)

REASONING_FAMILY = QuestionFamilySpec(
    name="REASONING",
    skill_target="M2",
    sampling=SamplingConstraint(
        min_subtitles=3,
        min_l3_nodes=4,
        require_frames=False,
        cross_l2_span=True,
    ),
    legal_task_types=frozenset(
        [
            "Action Reasoning",
            "Object Reasoning",
            "Event Reasoning",
            "Causal Reasoning",
            "Temporal Reasoning",
            "Spatial Reasoning",
        ]
    ),
    leak_profile=LeakTestProfile(
        shortcut_type="frequency",
        probe_template="gate_leak_reasoning.md",
        pass_threshold=0.5,
    ),
    prompt_template="reasoning.md",
)

ENUMERATION_FAMILY = QuestionFamilySpec(
    name="ENUMERATION",
    skill_target="M3",
    sampling=SamplingConstraint(
        min_subtitles=2,
        min_l3_nodes=5,
        require_frames=False,
        cross_l2_span=False,
    ),
    legal_task_types=frozenset(
        [
            "Action Sequence",
            "Object Recognition",
            "Object Interaction",
            "Scene Understanding",
        ]
    ),
    leak_profile=LeakTestProfile(
        shortcut_type="option_length",
        probe_template="gate_leak_enumeration.md",
        pass_threshold=0.6,
    ),
    prompt_template="enumeration.md",
)

VISUAL_FAMILY = QuestionFamilySpec(
    name="VISUAL",
    skill_target="M4",
    sampling=SamplingConstraint(
        min_subtitles=0,
        min_l3_nodes=3,
        require_frames=True,
        cross_l2_span=False,
    ),
    legal_task_types=frozenset(
        [
            "Object Recognition",
            "Scene Understanding",
            "Action Recognition",
            "Spatial Reasoning",
        ]
    ),
    leak_profile=LeakTestProfile(
        shortcut_type="visual_salience",
        probe_template="gate_leak_visual.md",
        pass_threshold=0.5,
    ),
    prompt_template="visual.md",
)

SPATIAL_FAMILY = QuestionFamilySpec(
    name="SPATIAL",
    skill_target="M5",
    sampling=SamplingConstraint(
        min_subtitles=0,
        min_l3_nodes=3,
        require_frames=True,
        cross_l2_span=False,
    ),
    legal_task_types=frozenset(
        [
            "Spatial Reasoning",
            "Object Interaction",
            "Scene Understanding",
        ]
    ),
    leak_profile=LeakTestProfile(
        shortcut_type="spatial_default",
        probe_template="gate_leak_spatial.md",
        pass_threshold=0.5,
    ),
    prompt_template="spatial.md",
)

ALL_FAMILIES: tuple[QuestionFamilySpec, ...] = (
    RETRIEVAL_FAMILY,
    REASONING_FAMILY,
    ENUMERATION_FAMILY,
    VISUAL_FAMILY,
    SPATIAL_FAMILY,
)

# 按 name 索引，用于 get_family_for_slot 快速查找
_FAMILY_BY_NAME: dict[str, QuestionFamilySpec] = {f.name: f for f in ALL_FAMILIES}


def get_family_for_slot(
    task_type: str,
    family_ratios: dict[str, float],
    rng: random.Random,
) -> QuestionFamilySpec:
    """根据任务类型和家族权重比例，随机选择一个合法的问题家族。

    Args:
        task_type: 任务类型字符串（必须是 12 种合法类型之一）。
        family_ratios: 家族名称到权重的映射（如 {"RETRIEVAL": 0.30, ...}）。
        rng: 随机数生成器实例（确保可复现）。

    Returns:
        被选中的 QuestionFamilySpec。

    Raises:
        ValueError: task_type 不在任何家族的 legal_task_types 中。
        ValueError: 给定 task_type 下没有合法家族（所有合法族权重为 0 或不在 ratios 中）。
    """
    # 检查 task_type 是否被任一家族接受
    all_legal_types: set[str] = set()
    for family in ALL_FAMILIES:
        all_legal_types.update(family.legal_task_types)
    if task_type not in all_legal_types:
        msg = f"task_type '{task_type}' 不在任何家族的合法类型中"
        raise ValueError(msg)

    # 过滤出接受该 task_type 且在 ratios 中有正权重的家族
    candidates: list[QuestionFamilySpec] = []
    weights: list[float] = []
    for family_name, weight in family_ratios.items():
        family = _FAMILY_BY_NAME.get(family_name)
        if family is None:
            continue
        if task_type in family.legal_task_types and weight > 0:
            candidates.append(family)
            weights.append(weight)

    if not candidates:
        msg = f"task_type '{task_type}' 下没有合法家族可选（检查 family_ratios）"
        raise ValueError(msg)

    # 归一化权重 + 加权随机选择
    chosen = rng.choices(candidates, weights=weights, k=1)[0]
    return chosen