"""世谱预算裁剪与注入分组。

按 token 预算裁剪（priority 值小优先保留），再按 insertion_order 排定顺序，
最后按 position 分组为可注入的文本块。纯逻辑，无 IO。
"""

from __future__ import annotations

from collections.abc import Callable

from app.modules.lorebook.types import MatchedEntry

TokenCounter = Callable[[str], int]

_POSITIONS = ("before_char", "after_char")


def _default_token_counter(text: str) -> int:
    """P0 以字符数近似 token；后续接神识 tokenizer。"""
    return len(text)


def select_within_budget(
    matched: list[MatchedEntry],
    token_budget: int | None = None,
    token_counter: TokenCounter | None = None,
) -> list[MatchedEntry]:
    """按 priority 选入预算内条目，再按 insertion_order 排定注入顺序。"""
    counter = token_counter or _default_token_counter
    if token_budget is None:
        return sorted(matched, key=lambda m: m.entry.insertion_order)

    by_priority = sorted(matched, key=lambda m: m.entry.priority)
    kept: list[MatchedEntry] = []
    total = 0
    for m in by_priority:
        cost = counter(m.entry.content)
        if cost == 0:
            continue
        if total + cost > token_budget:
            continue
        kept.append(m)
        total += cost
    kept.sort(key=lambda m: m.entry.insertion_order)
    return kept


def build_injections(matched: list[MatchedEntry]) -> dict[str, str]:
    """按 position 分组，合并同组条目 content 为单段文本。

    返回 ``{"before_char": "...", "after_char": "..."}``，空组不出现在结果中。
    集成层据此把 before_char 插在道体 personality_summary 之前、after_char 之后。
    """
    groups: dict[str, list[str]] = {pos: [] for pos in _POSITIONS}
    for m in matched:
        pos = m.entry.position if m.entry.position in groups else "after_char"
        groups[pos].append(m.entry.content)
    return {pos: "\n".join(parts) for pos, parts in groups.items() if parts}