"""混合 mini-batch 切分：以 QuestionUnit 为最小调度粒度，大类打散、小类整锁。

供 runner 每 step 处理一个 batch。孪生对（AR pair）作为 2 题单元整锁不拆、按单元级
正确性分桶；非 AR single 单元的抽样/洗牌 draw 流与"引入 QuestionUnit 前"的旧逐题算法
逐字节一致（AR 折叠不干扰非 AR draw 流）。
"""

from __future__ import annotations

import hashlib
import math
import random
from typing import TYPE_CHECKING

from app.harness.question_units import build_units, flatten_units, unit_correctness

if TYPE_CHECKING:
    from core.types import GeneratedQuestion, QuestionUnit


def _rng_ns(seed: int, ns: str) -> random.Random:
    """由 (seed, 命名空间) 稳定派生独立随机数发生器。

    用 SHA-256 派生而非 Python 内置 ``hash()``——后者受 hash randomization 影响，
    跨进程不可复现。不同命名空间的 draw 流互不干扰，使 AR 单元折叠不扰动非 AR 抽样。

    参数:
        seed: 实验随机种子。
        ns: 命名空间标签（如 "AR"）。
    返回:
        以 SHA-256(f"{ns}:{seed}") 前 8 字节为种子的 ``random.Random``。
    """
    digest = hashlib.sha256(f"{ns}:{seed}".encode()).digest()
    return random.Random(int.from_bytes(digest[:8], "big"))


def build_batches(
    items: list[GeneratedQuestion],
    correctness: dict[str, bool],
    batch_size: int,
    min_class_per_batch: int,
    seed: int,
    correct_ratio: float = 0.0,
) -> tuple[list[list[GeneratedQuestion]], int]:
    """把诊断池里的题目切成多个混合 mini-batch（以 QuestionUnit 为原子调度单元）。

    single 题为 1 题单元，AR pair 孪生对为 2 题单元；同一 pair 的两题整锁进同一 batch，
    按单元级正确性（双向 AND）分桶。当 ``correct_ratio > 0`` 时，按题型为每组错误单元配比
    一定数量的正确单元（"动量"机制）；``correct_ratio <= 0`` 时退化为纯错误单元模式。

    参数:
        items: 候选题目全集（可混含 single 与孪生对成员）。
        correctness: question_id -> 基线是否答对。
        batch_size: 单个 batch 的题目数上限（> 0，pair 占 2）。
        min_class_per_batch: 小类判定阈值——单元题目总数 ≤ 此值的题型整组锁进单一
            batch（> 0）。
        seed: 随机种子，保证相同输入产出完全一致的切分。
        correct_ratio: 正确题占比（0.0 ~ 1.0）。0.0 = 纯错误单元；0.5 = 错:正 = 1:1。
    返回:
        (非空 mini-batch 列表, selected_count)；无错误单元时返回 ([], 0)。
        selected_count 是所有 batch 中题目（展开后）总数。
    异常:
        ValueError: batch_size 或 min_class_per_batch < 1, 或
            min_class_per_batch >= batch_size（破坏小类整组装箱不超容的前提）。
    关键实现细节:
        非 AR（single）与 AR（pair）各用独立稳定派生的 rng：非 AR 用 ``random.Random(seed)``
        （复现旧逐题算法的确切 draw 序列，保证纯非 AR 输入逐字节一致），AR 用
        ``_rng_ns(seed, "AR")``；二者 draw 流互不干扰，故加入/移除 pair 不改变非 AR 的
        抽样/洗牌序列。抽样在合并前按流分别进行（``_select_mixed_by_task_type`` 各跑一次），
        大类洗牌按单元 kind 拆分后各用对应流。装箱顺序「先小类后大类」：小类整组
        first-fit-decreasing（容量按单元 ``size`` 计，pair 占 2）装入首个容得下的 batch，
        装不下新开 bin；大类洗牌后 round-robin 分发，遇碎片（size-2 单元放不进任一现存
        batch 的剩余容量）新开 bin 兜底而非报错。最终每个 batch 展开回题目列表。
        题型按名称排序处理以保证跨运行确定性。
    """
    _validate_params(batch_size, min_class_per_batch)

    # 非 AR 复现旧版 random.Random(seed) 的确切序列以满足黄金 byte-identity；
    # AR 走独立命名空间派生流，二者互不干扰。
    rng_nonar = random.Random(seed)
    rng_ar = _rng_ns(seed, "AR")

    grouped = _group_units_by_task_type(items, correctness, correct_ratio, rng_nonar, rng_ar)

    total = sum(_group_load(g) for g in grouped.values())
    if total == 0:
        return [], 0

    nb = max(1, math.ceil(total / batch_size))
    batches: list[list[QuestionUnit]] = [[] for _ in range(nb)]

    small, large = _split_by_size(grouped, min_class_per_batch)
    for group in _small_groups_decreasing(small):
        _pack_small_class(batches, group, batch_size)
    _distribute_large_classes(batches, large, batch_size, rng_nonar, rng_ar)

    result = [flatten_units(b) for b in batches if b]
    selected_count = sum(len(b) for b in result)
    return result, selected_count


def _validate_params(batch_size: int, min_class_per_batch: int) -> None:
    """校验切分参数，非法值直接报错而非用默认值掩盖。

    除各自 >= 1 外，强制 min_class_per_batch < batch_size：小类组题目总数 ≤
    min_class_per_batch，唯有此前提成立才能保证小类整组放入单一 batch 而不超容；否则
    _pack_small_class 新开的 bin 会装入超 batch_size 的整组，静默违反容量合约。此约束
    与 config._validate_minibatch 一致，是 build_batches 对自身前提的防御性自校验（P5）。
    """
    if batch_size < 1:
        raise ValueError(f"batch_size 必须 >= 1, 实为 {batch_size}")
    if min_class_per_batch < 1:
        raise ValueError(f"min_class_per_batch 必须 >= 1, 实为 {min_class_per_batch}")
    if min_class_per_batch >= batch_size:
        raise ValueError(
            f"min_class_per_batch 必须严格 < batch_size, 否则无法保证小类整组放入单一 "
            f"batch 不超容; 实为 min_class_per_batch={min_class_per_batch}, "
            f"batch_size={batch_size}"
        )


def _group_units_by_task_type(
    items: list[GeneratedQuestion],
    correctness: dict[str, bool],
    correct_ratio: float,
    rng_nonar: random.Random,
    rng_ar: random.Random,
) -> dict[str, list[QuestionUnit]]:
    """把题目聚合为单元并按题型分组：非 AR 与 AR 各走独立 draw 流后合并。

    参数:
        items: 候选题目全集。
        correctness: question_id -> 基线是否答对。
        correct_ratio: 正确题占比。
        rng_nonar: 非 AR（single 单元）抽样用 rng。
        rng_ar: AR（pair 单元）抽样用 rng。
    返回:
        task_type -> 混合后的单元列表（single 单元在前、pair 单元在后）。
    """
    units = build_units(items)
    singles = [u for u in units if u.kind == "single"]
    pairs = [u for u in units if u.kind == "pair"]
    grouped_nonar = _select_mixed_by_task_type(singles, correctness, correct_ratio, rng_nonar)
    grouped_ar = _select_mixed_by_task_type(pairs, correctness, correct_ratio, rng_ar)
    return _merge_grouped(grouped_nonar, grouped_ar)


def _group_load(group: list[QuestionUnit]) -> int:
    """一组单元展开后的题目总数（single 计 1，pair 计 2），即占用的 batch 容量。"""
    return sum(u.size for u in group)


def _batch_load(batch: list[QuestionUnit]) -> int:
    """一个 batch 内单元展开后的题目总数，用于容量判断。"""
    return sum(u.size for u in batch)


def _merge_grouped(
    grouped_nonar: dict[str, list[QuestionUnit]],
    grouped_ar: dict[str, list[QuestionUnit]],
) -> dict[str, list[QuestionUnit]]:
    """按 task_type 合并非 AR 与 AR 两条流的分组（single 在前、pair 在后）。

    参数:
        grouped_nonar: 非 AR（single 单元）分组。
        grouped_ar: AR（pair 单元）分组。
    返回:
        task_type -> 合并后的单元列表；每类 single 单元在前、pair 单元在后，顺序稳定。
    """
    merged: dict[str, list[QuestionUnit]] = {}
    for task_type in sorted({*grouped_nonar, *grouped_ar}):
        merged[task_type] = grouped_nonar.get(task_type, []) + grouped_ar.get(task_type, [])
    return merged


def _split_by_size(
    grouped: dict[str, list[QuestionUnit]],
    min_class_per_batch: int,
) -> tuple[dict[str, list[QuestionUnit]], dict[str, list[QuestionUnit]]]:
    """按题目总数（单元展开）把题型分为小类（≤ 阈值）与大类（> 阈值）两组。"""
    small = {t: g for t, g in grouped.items() if _group_load(g) <= min_class_per_batch}
    large = {t: g for t, g in grouped.items() if _group_load(g) > min_class_per_batch}
    return small, large


def _classify_unit(unit: QuestionUnit, correctness: dict[str, bool]) -> str | None:
    """判定单元落入哪个桶：error / correct / None（未知，跳过）。

    参数:
        unit: 目标单元。
        correctness: question_id -> 是否答对（缺键视为未知）。
    返回:
        "error"（单元级正确性为 False）、"correct"（双向 AND 为 True）；单元内任一题
        未知（correctness 缺该键）返回 None，与旧逐题算法把未知题排除在错/对两桶之外
        的语义一致。
    关键实现:
        先探测是否有未知题（get 返回 None ⟺ 键缺失，因 correctness 值恒为 bool），
        全部已知后交由 unit_correctness 计双向 AND（此时 KeyError 不可达）。
    """
    if any(correctness.get(q.question_id) is None for q in unit.questions):
        return None
    return "correct" if unit_correctness(unit, correctness) else "error"


def _select_mixed_by_task_type(
    units: list[QuestionUnit],
    correctness: dict[str, bool],
    correct_ratio: float,
    rng: random.Random,
) -> dict[str, list[QuestionUnit]]:
    """按题型分组，为每组错误单元按比例采样正确单元混入（单元粒度）。

    只对有错误单元的题型做混合——无错误单元的题型不进 batch，即使有正确单元。
    ``correct_ratio <= 0`` 时退化为纯错误单元模式。本函数只处理单一 draw 流（全 single
    或全 pair），使非 AR 与 AR 的抽样互不干扰。

    参数:
        units: 同一流的候选单元（全 single 或全 pair）。
        correctness: question_id -> 基线是否答对。
        correct_ratio: 正确题占比（0.0 ~ 1.0）。
        rng: 本流专用随机数发生器，用于采样正确单元。
    返回:
        task_type -> 该题型的混合单元列表（错误单元全部 + 按比例采样的正确单元）。
    关键实现:
        n_correct 按错误单元「题目总数」而非单元数计，与旧逐题语义对齐（纯 single 时
        单元数 == 题目数，采样序列逐字节一致）。
    """
    errors_by_type: dict[str, list[QuestionUnit]] = {}
    correct_by_type: dict[str, list[QuestionUnit]] = {}
    for unit in units:
        bucket = _classify_unit(unit, correctness)
        if bucket == "error":
            errors_by_type.setdefault(unit.task_type, []).append(unit)
        elif bucket == "correct":
            correct_by_type.setdefault(unit.task_type, []).append(unit)

    if correct_ratio <= 0:
        return errors_by_type

    grouped: dict[str, list[QuestionUnit]] = {}
    for task_type in sorted(errors_by_type):
        errs = errors_by_type[task_type]
        n_err = _group_load(errs)
        n_correct = round(n_err * correct_ratio / (1 - correct_ratio))
        available = correct_by_type.get(task_type, [])
        sampled = (
            list(available) if len(available) <= n_correct else rng.sample(available, n_correct)
        )
        grouped[task_type] = errs + sampled

    return grouped


def _small_groups_decreasing(
    small: dict[str, list[QuestionUnit]],
) -> list[list[QuestionUnit]]:
    """按组题目总数降序、同大小按 task_type 升序排出小类组（first-fit-decreasing 顺序）。

    参数:
        small: task_type -> 小类单元列表。
    返回:
        排好序的小类组列表；降序处理可降低碎片，确定性 tie-break 保证跨运行一致。
    """
    return [small[t] for t in sorted(small, key=lambda t: (-_group_load(small[t]), t))]


def _pack_small_class(
    batches: list[list[QuestionUnit]],
    group: list[QuestionUnit],
    batch_size: int,
) -> None:
    """用 first-fit 把一个小类整组放入首个容得下的 batch，装不下则新开 bin（就地修改）。

    因小类组题目总数 ≤ min_class_per_batch < batch_size，新开的空 batch 必能容纳整组，
    故此函数永不抛 ValueError，且整组（含内部 pair 单元）不拆。

    参数:
        batches: 当前各 batch（就地追加，必要时 append 新空 batch）。
        group: 待锁定的小类单元组（整组不拆）。
        batch_size: 单 batch 题目容量上限。
    """
    load = _group_load(group)
    for b in batches:
        if _batch_load(b) + load <= batch_size:
            b.extend(group)
            return
    batches.append(list(group))


def _distribute_large_classes(
    batches: list[list[QuestionUnit]],
    large: dict[str, list[QuestionUnit]],
    batch_size: int,
    rng_nonar: random.Random,
    rng_ar: random.Random,
) -> None:
    """将各大类单元洗牌后 round-robin 分发到所有现存 batch（就地修改）。

    参数:
        batches: 当前各 batch（含小类装箱可能新开的 bin，就地追加）。
        large: task_type -> 大类单元列表。
        batch_size: 单 batch 题目容量上限。
        rng_nonar: 非 AR（single 单元）洗牌用 rng。
        rng_ar: AR（pair 单元）洗牌用 rng。
    关键实现细节:
        每组按单元 kind 拆成 single 子列与 pair 子列，分别用 rng_nonar / rng_ar 洗牌后
        拼接（single 在前），使非 AR 洗牌 draw 流不受 pair 存在与否影响（纯 single 时
        single 子列即整组，复现旧版单一 rng.shuffle 的序列）。全局指针在所有大类单元间
        持续轮转，遇满箱跳过、遇碎片新开 bin。题型按名称排序以保证分发顺序确定。
    """
    pointer = 0
    for task_type in sorted(large):
        group = large[task_type]
        singles = [u for u in group if u.kind == "single"]
        pairs = [u for u in group if u.kind == "pair"]
        rng_nonar.shuffle(singles)
        rng_ar.shuffle(pairs)
        for unit in singles + pairs:
            pointer = _place_round_robin(batches, unit, pointer, batch_size)


def _place_round_robin(
    batches: list[list[QuestionUnit]],
    unit: QuestionUnit,
    pointer: int,
    batch_size: int,
) -> int:
    """从 pointer 起找第一个容量够放 unit 的 batch 放入，返回下一次起始指针。

    参数:
        batches: 当前各 batch（就地追加）。
        unit: 待放置的单元（占用 unit.size 个容量）。
        pointer: 本次轮转起始 batch 下标。
        batch_size: 单 batch 题目容量上限。
    返回:
        下一次轮转的起始指针（已前移一位）。
    关键实现:
        单个单元容量 ≤ batch_size 是前提（pair 占 2，而 batch_size > min_class ≥ 1 ⇒
        batch_size ≥ 2），故此处断言防御。扫描一轮所有现存 batch 都放不下（size-2 单元
        遇满地碎片）时新开 bin 兜底而非报错——聚合容量足够但单箱剩余不足是合法碎片场景。
        纯 single（size 1）永不触发新开分支，故与旧逐题 round-robin 逐字节一致。
    """
    assert unit.size <= batch_size, f"单元 size={unit.size} 超过 batch_size={batch_size}"
    nb = len(batches)
    for offset in range(nb):
        idx = (pointer + offset) % nb
        if _batch_load(batches[idx]) + unit.size <= batch_size:
            batches[idx].append(unit)
            return (idx + 1) % nb
    batches.append([unit])
    return len(batches) % len(batches)
