StarAsh042-ComfyUI_RollingArtist/RollingCharacter.py

"""ComfyUI RollingCharacter 节点。

从角色 CSV（默认 ``danbooru_character_001.csv``）中随机抽取角色，
按 ``角色标签,作品标签,core_tags`` 逐行输出（多个角色用「逗号 + 换行」分隔），
同时输出结构化 JSON 与已测数量。

与 RollingArtist 的关系：抽样与进度输出的做法一致；差异来自数据与产物：

- **不做权重分配**：角色标签直接输出，不套 ``(名称:权重)``，因此没有权重相关参数
- **名字取自 ``trigger`` 列的第 1 段**：该列经 34,416 行全量校验是恒定的
  「角色标签,作品标签」两段结构，第 1 段已转义、可直接使用，第 2 段是作品标签
- **每个角色一行**：``角色标签,作品标签,core_tags``（空的部分自动省略）；
  多个角色之间用「逗号 + 换行」分隔（行尾补逗号，拼成一行时标签不会粘在一起）
- **可按作品过滤**：``copyright_pick`` 下拉列出**全部作品**，可打字筛选
- **不重复靠一个简单的内存窗口**：最近 10 条输出记录里出现过的角色名会被回避
  （见 :data:`_RECENT_RECORDS`），没有去重模式参数
- **不做穷举**：角色是单词条，穷举意义不大，因此没有 ``mode`` 参数
"""

import json
import random
import threading
from collections import deque
from typing import Any, Deque, Dict, List, Optional, Sequence, Set, Tuple

from .ra_core.artists import (
    available_names,
    compute_top_count,
    get_repository,
    select_artists,
    split_names,
)
from .ra_core.characters import (
    COPYRIGHT_ANY,
    CharacterIndex,
    CharacterRecord,
    build_character_prompt,
    build_payload,
    copyright_options,
    load_index,
    load_records,
)
from .ra_core.constants import (
    DEFAULT_CHARACTER_CSV,
    DEFAULT_CHARACTER_DB,
    KIND_CHARACTER,
    LOGGER,
)
from .ra_core.db import get_db
from .ra_core.params import as_bool as _as_bool
from .ra_core.params import as_int as _as_int
from .ra_core.params import as_text as _text

__all__ = ["RollingCharacter", "NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

# 「最近多少条输出记录内角色不重复」：这里的一条 = 一次生成（可能含多个角色）。
# 窗口放在节点实例的内存里，因此重启 ComfyUI 会清零，多个实例也各算各的。
_RECENT_RECORDS = 10


def _key_weights(count: int) -> List[float]:
    """记录用的权重占位值：角色输出不带权重，键里统一按 1.0 计算。"""
    return [1.0] * count


# 作品下拉菜单的选项：按「默认 CSV 的文件签名」缓存（3.4 万行，首次约 0.5 秒）。
# 用户替换 CSV 后签名变化，菜单会自动重建，不需要重启 ComfyUI。
# 用户填的 ``custom_csv_path`` 在节点注册阶段拿不到，因此不参与构建。
_COPYRIGHT_MENU: Optional[List[str]] = None
_COPYRIGHT_MENU_KEY: Optional[Tuple[str, int, int]] = None
_COPYRIGHT_LOCK = threading.RLock()


def _copyright_menu() -> List[str]:
    """作品下拉菜单选项：``[COPYRIGHT_ANY, 全部作品名...]``（按 count 总和降序）。

    默认 CSV 缺失或解析失败时退化为只有「不限」一项，绝不让菜单构建失败影响节点注册。
    """
    global _COPYRIGHT_MENU, _COPYRIGHT_MENU_KEY
    repository = get_repository(DEFAULT_CHARACTER_CSV)
    key = repository.stat_key(repository.resolve_path(""))
    with _COPYRIGHT_LOCK:
        if _COPYRIGHT_MENU is not None and key == _COPYRIGHT_MENU_KEY:
            return _COPYRIGHT_MENU
        options: List[str] = [COPYRIGHT_ANY]
        try:
            options.extend(copyright_options(load_index(repository, "", "trigger").records))
        except Exception as error:  # pragma: no cover - 菜单构建必须保持健壮
            LOGGER.error("[RollingCharacter] 构建作品下拉菜单失败: %s", error)
        if len(options) == 1:
            LOGGER.warning(
                "[RollingCharacter] 作品下拉菜单只有“不限”一项（默认 CSV 缺失或没有读到作品列），"
                "请检查 danbooru_character_001.csv 是否就位"
            )
        _COPYRIGHT_MENU = options
        _COPYRIGHT_MENU_KEY = key
        return _COPYRIGHT_MENU


# 作品筛选报错时给出的热门作品示例数（帮用户快速对照「这个数据集里有什么」）
_PICK_HINT_SAMPLE = 5


def _copyright_pick_reject_reason(pick: str, index: CharacterIndex, custom_csv_path: str) -> str:
    """作品筛选结果为空时的报错原因：指明作品名不在哪个数据集里，并给出可用示例。

    用户最常见的困惑：作品下拉只列默认 CSV 的作品，填了自定义 CSV 路径后
    选了个默认 CSV 的作品名 -> 在自定义 CSV 里筛不出角色。这时必须说清楚
    「不在你的自定义 CSV 里」，而不是笼统的「没有可用角色」。
    """
    copyrights = index.copyrights
    if copyrights:
        sample = "、".join(copyrights[:_PICK_HINT_SAMPLE])
        total = f"共 {len(copyrights)} 部作品，热门示例：{sample}"
    else:
        total = "该数据集没有任何作品列数据"
    if _text(custom_csv_path).strip():
        return (
            f"作品选择={pick!r} 不在自定义 CSV 的作品列表中（{total}）。"
            "作品下拉仅列出默认 CSV 的作品，自定义 CSV 的作品名不会出现在下拉里，"
            "请直接填写该 CSV 中的作品名（原始写法 fate_(series) 亦可）"
        )
    return f"作品选择={pick!r} 不在默认 CSV 的作品列表中（{total}）。请改成“(不限)”或换一个作品"


class RollingCharacter:
    """RollingCharacter 节点类。

    功能特性：
    - 从角色 CSV 加载角色（默认 ``trigger`` 列，已转义、可直接使用）
    - 支持 Top 池优先、排除/强制包含名单、按作品筛选（下拉列出全部作品）
    - 逐行输出 ``角色标签,作品标签,core_tags``（不带权重；角色描述可开关）；
      多个角色之间用「逗号 + 换行」分隔
    - 结构化输出 characters_json，便于下游处理
    - 最近 10 条输出记录内不重复（内存窗口，无去重模式参数）

    线程安全：
    - 角色 CSV 解析结果按文件签名缓存，多个实例共享
    - 「选角色 -> 记录」整段处于数据库事务中；不重复判定用的是节点实例自己的内存窗口
    - 记录带 ``kind=character`` 标记：即使被指向画师节点的数据库也不会互相冒充
    """

    @classmethod
    def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]:
        """定义节点的输入参数类型和界面显示。

        界面顺序即下方字典顺序，按「抽样 → Top 池 → 种子 → 抽样偏好 → 名单 →
        作品 → 输出内容 → 高级（CSV 与数据库路径）」归类；
        很少改动的 CSV / 路径标记为 ``advanced``，由前端折叠到“高级输入”。

        角色节点不做权重分配、不做穷举；输出格式固定为 ``角色标签,作品标签,core_tags``，
        只提供「开启角色描述」这一个内容开关。
        """
        return {
            "required": {
                # ---------------- 抽样：抽几个角色、从哪里抽 ----------------
                "character_count": (
                    "INT",
                    {
                        "default": 1,
                        "min": 1,
                        "max": 10,
                        "step": 1,
                        "display": "slider",
                        "tooltip": "每次生成的角色数量（1-10），每个角色输出一行。角色是单词条，没有穷举模式；重复回避窗口为最近 10 次生成，存于节点实例内存（重启清零、多实例各自独立）",
                    },
                ),
                "character_top_ratio": (
                    "FLOAT",
                    {
                        "default": 0.01,
                        "min": 0.01,
                        "max": 1.0,
                        "step": 0.01,
                        "display": "slider",
                        "tooltip": "人气池占比（0.01-1.0）：按 CSV 行序取前 ceil(候选数×比例) 个"
                        "（至少 1 个）作为 Top 池，行序越靠前的角色作品数越多。"
                        "候选数指“作品筛选之后”剩下的角色数",
                    },
                ),
                "character_top_count": (
                    "INT",
                    {
                        "default": 1,
                        "min": 1,
                        "max": 10,
                        "step": 1,
                        "display": "slider",
                        "tooltip": "每次至少抽取几个人气池角色（1-10）；"
                        "仅在“Top 优先抽样”开启且人气池非空时生效",
                    },
                ),
                "seed": (
                    "INT",
                    {
                        "default": 1234,
                        "min": 0,
                        "max": 4294967295,
                        "tooltip": "随机种子。同一种子在同一组参数下结果可复现；"
                        "最近 10 条输出记录的回避窗口也会影响实际结果",
                    },
                ),
            },
            "optional": {
                # ---------------- 抽样偏好 ----------------
                "use_top_priority": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "开启时优先从人气池抽角色，并保证含 character_top_count 个人气角色；"
                        "关闭时从完整候选列表均匀抽样",
                    },
                ),
                # ---------------- 名单 ----------------
                "exclude_characters": (
                    "STRING",
                    {"default": "", "tooltip": "要排除的角色名，英文逗号分隔；被排除的名字不会出现在结果中"},
                ),
                "force_include": (
                    "STRING",
                    {
                        "default": "",
                        "tooltip": "必须出现的角色名，英文逗号分隔；优先级高于“排除角色”，"
                        "数量达到角色数量时直接从中随机抽取。"
                        "不在 CSV 里的名字也能用，只是没有作品与角色描述信息",
                    },
                ),
                # ---------------- 作品 ----------------
                "copyright_pick": (
                    _copyright_menu(),
                    {
                        "default": COPYRIGHT_ANY,
                        "tooltip": "从下拉里挑一个作品。下拉列出 CSV 里的全部作品"
                        "（默认数据约 3460 个，按作品热度降序：该作品下所有角色的 count 之和），"
                        "并且可以直接打字筛选、忽略大小写（输入 a 只留 a 开头的，"
                        "继续输入 ab 会进一步收窄）。“(不限)”表示不筛选。"
                        "同样接受 copyright 列的原始写法 fate_(series) 与 trigger 里"
                        "转义后的写法 fate_\\(series\\)。注意：下拉项来自默认 CSV，"
                        "换“自定义 CSV 路径”时下拉不会跟着变——但筛选会在该 CSV 上生效，"
                        "其作品名不会出现在下拉里，请直接填写",
                    },
                ),
                # ---------------- 输出内容 ----------------
                "describe_character": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "开启时在每行末尾拼上角色的固定外观标签（core_tags，即角色描述），"
                        "格式为「角色标签,作品标签,角色描述」；"
                        "关闭时只输出「角色标签,作品标签」",
                    },
                ),
                # ---------------- 高级：CSV 与数据库（界面默认折叠） ----------------
                "custom_csv_path": (
                    "STRING",
                    {
                        "default": "",
                        "advanced": True,
                        "tooltip": "自定义角色 CSV 的路径；留空使用节点目录下的 danbooru_character_001.csv。"
                        "作品筛选（“作品”参数）在该 CSV 上生效，但其作品名不会出现在作品下拉里，"
                        "请直接填写该 CSV 中的作品名",
                    },
                ),
                "tested_db_path": (
                    "STRING",
                    {
                        "default": "",
                        "advanced": True,
                        "tooltip": "记录数据库（SQLite）路径；留空使用节点目录下的 "
                        "rollingcharacter.sqlite（与画师节点分开）——默认路径的库每次启动会被清空"
                        "（不跨会话累积），想跨会话保留请把这里指向自己的文件。"
                        "这里只存历史输出记录，更换该路径不会重置“最近 10 条不重复”的窗口"
                        "（那个窗口在内存里，重启清零、多实例各自独立）",
                    },
                ),
                # ---------------- 高级：CSV 读取上限（追加在最后，见下方兼容性说明） ----------------
                "csv_max_rows": (
                    "INT",
                    {
                        "default": 0,
                        "min": 0,
                        "max": 1000000,
                        "step": 1000,
                        "advanced": True,
                        "tooltip": "CSV 读取上限（条记录，含表头；0 = 不限制）。"
                        "超过上限的记录不参与加载，候选池、作品菜单与热度池会相应变小；"
                        "默认 0 保持与旧版完全一致的全量读取。"
                        "表很大且只想用前一部分时可设为正整数以降低读取耗时与内存",
                    },
                ),
            },
        }

    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("prompt", "characters_json", "tested_count")
    FUNCTION = "generate_characters"
    CATEGORY = "RollingArtist"
    # 作为输出节点，即使输出未被下游使用也会执行
    OUTPUT_NODE = True
    DESCRIPTION = (
        "从角色 CSV 中随机抽取角色，按「角色标签,作品标签,角色描述」逐行输出"
        "（多个角色用「逗号+换行」分隔、不带权重），并输出结构化 JSON 与已测数量。"
    )
    OUTPUT_TOOLTIPS = (
        "角色列表：每个角色一行，格式为「角色标签,作品标签,角色描述」，不带权重；"
        "关闭“开启角色描述”后为「角色标签,作品标签」",
        "结构化结果 JSON（角色 / 作品 / 是否 Top / 热度 / 角色描述）",
        "当前数据库里已记录的角色组合数量（字符串形式，可直接接到文本节点查看）",
    )

    @classmethod
    def IS_CHANGED(cls, custom_csv_path: str = "", **kwargs: Any) -> str:
        """把角色 CSV 的文件签名纳入缓存键：文件被替换时强制重算。

        与 RollingArtist 同源（共用 :func:`ra_core.storage.file_fingerprint`）：
        参数不变而 CSV 内容变了时，默认缓存会一直返回旧角色列表。
        文件与参数都不变时指纹稳定，缓存行为不变。
        """
        return get_repository(DEFAULT_CHARACTER_CSV).cache_fingerprint(custom_csv_path)

    def __init__(self) -> None:
        """初始化节点实例：准备角色仓库、状态容器与“最近记录”窗口。

        真正的加载发生在 ``generate_characters`` 中（按文件签名缓存），
        这里只做一次默认 CSV 的预热，且任何异常都不会影响节点注册。
        """
        self._repo = get_repository(DEFAULT_CHARACTER_CSV)
        self._lock = threading.RLock()
        self.characters: List[CharacterRecord] = []
        # 最近若干条输出记录（每条 = 一次生成的角色名），用于回避重复
        self._recent: Deque[Tuple[str, ...]] = deque(maxlen=_RECENT_RECORDS)
        try:
            self.characters = load_records(self._repo, "", "trigger")
        except Exception as error:  # pragma: no cover - 初始化必须保持健壮
            LOGGER.error("[RollingCharacter] 初始化加载默认角色 CSV 失败: %s", error)

    # ------------------------------------------------------------------
    # 最近记录窗口
    # ------------------------------------------------------------------
    def recent_names(self) -> Set[str]:
        """最近 :data:`_RECENT_RECORDS` 条输出记录里出现过的角色名（并集）。"""
        with self._lock:
            names: Set[str] = set()
            for record_names in self._recent:
                names.update(record_names)
            return names

    def _remember(self, names: Sequence[str]) -> None:
        """把本次输出的角色名记进窗口。"""
        with self._lock:
            self._recent.append(tuple(names))

    def reset_recent(self) -> None:
        """清空“最近记录”窗口（手动重置重复回避时使用）。"""
        with self._lock:
            self._recent.clear()

    # ------------------------------------------------------------------
    # 主流程
    # ------------------------------------------------------------------
    def generate_characters(
        self,
        character_count: int,
        character_top_ratio: float,
        character_top_count: int,
        seed: int,
        use_top_priority: bool = True,
        exclude_characters: str = "",
        force_include: str = "",
        copyright_pick: str = COPYRIGHT_ANY,
        describe_character: bool = True,
        custom_csv_path: str = "",
        tested_db_path: str = "",
        csv_max_rows: int = 0,
    ) -> Tuple[str, str, str]:
        """生成角色列表。

        形参顺序与界面参数顺序一致（ComfyUI 按关键字传参，顺序只影响外部脚本按位置调用）。
        读取的 CSV 列固定为 ``trigger``（「角色标签,作品标签」两段结构），不可配置。

        参数:
            csv_max_rows: CSV 读取上限（条记录，含表头；``0`` = 不限制）。
                新增参数**追加在形参表末尾**，因此按位置调用的外部脚本不受影响。

        不重复的规则很简单：把最近 :data:`_RECENT_RECORDS` 条输出记录里出现过的角色名
        从候选里排除；只有当窗口把候选全部吃掉时才退回「允许重复」并输出 WARNING。

        返回:
            ``(prompt, characters_json, tested_count)``；``tested_count`` 是字符串，
            方便直接接到文本节点查看
        """
        count = _as_int(character_count, 1)
        if count < 1:
            raise ValueError(f"[RollingCharacter] character_count 必须 >= 1，当前值: {character_count!r}")

        # 0) 读取上限：0 = 不限制（映射为 None）；其他值必须是正整数
        limit = _as_int(csv_max_rows, 0)
        max_rows = None if limit == 0 else limit
        if max_rows is not None and max_rows < 0:
            raise ValueError(f"[RollingCharacter] csv_max_rows 必须 >= 0，当前值: {csv_max_rows!r}")

        # 1) 加载角色记录（内容未变化时直接命中缓存，不重复解析 CSV）
        #    列固定为 trigger：该列恒为「角色标签,作品标签」两段结构
        #    派生结构（名字列表 / 名字->记录 / 作品筛选下标）随记录一起缓存，
        #    避免每次执行都在 3.4 万条数据上重建一遍（实测约 19 ms）
        index = load_index(self._repo, custom_csv_path, "trigger", max_rows=max_rows)
        records = index.records
        if not records:
            raise ValueError(
                "[RollingCharacter] 未能从 CSV 加载角色列表，请检查文件是否存在、"
                f"编码是否为 UTF-8、是否含有 trigger 列: "
                f"{self._repo.resolve_path(_text(custom_csv_path))}"
                "（若本节点是从旧工作流载入的，请检查“自定义 CSV 路径”是否已清空）"
            )
        with self._lock:
            self.characters = records

        # 2) 作品筛选（下拉单选，“(不限)”表示不筛选）+ 名单
        force_list = split_names(force_include)
        exclude_set = set(split_names(exclude_characters))
        pick = _text(copyright_pick).strip()
        copyright_filter = None if pick == COPYRIGHT_ANY else pick
        allowed = index.indices_for(copyright_filter)
        if not allowed and not force_list:
            raise ValueError(
                "[RollingCharacter] " + _copyright_pick_reject_reason(pick, index, custom_csv_path)
            )

        # 名字唯一（默认数据 34,416 行实测无重名），因此可以按名字直接取回记录
        record_map = index.record_map
        names = index.names_for(copyright_filter)
        top_count = compute_top_count(len(names), character_top_ratio)

        rng_seed = _as_int(seed, 0)
        top_priority = top_count if _as_bool(use_top_priority, True) else 0
        with_core = _as_bool(describe_character, True)

        # 3) 回避最近 _RECENT_RECORDS 条输出记录里出现过的角色
        #    分两步走：先算"排除名单后还剩哪些候选"，再决定是否启用回避窗口。
        #    这样"窗口把候选吃光时退回允许重复"的降级是显式的，而不是事后修正一个
        #    已经拼好的屏蔽集。
        recent = self.recent_names()
        candidates = available_names(names, exclude_set)
        blocked = exclude_set
        if recent:
            survivors = available_names(candidates, recent)
            if survivors:
                blocked = exclude_set | recent
            elif candidates and not force_list:
                # 候选本来就少、且已被窗口全部覆盖：退回允许重复，并明确告警
                LOGGER.warning(
                    "[RollingCharacter] 最近 %d 条输出记录已覆盖全部 %d 个候选角色，"
                    "本轮不再回避重复（可用角色太少时属于正常现象）",
                    _RECENT_RECORDS,
                    len(candidates),
                )

        # 4) 记录库：角色专用 SQLite（记录带 kind 标记，与画师互不干扰）
        db = get_db(tested_db_path, DEFAULT_CHARACTER_DB)

        # 5) 「选角色 -> 记录」处于同一个事务内，进程内 / 跨进程都不会重复落盘
        with db.generation_lock():
            rng = random.Random(rng_seed)
            selected = select_artists(
                rng,
                count,
                _as_int(character_top_count, 1),
                force_list,
                names,
                top_priority,
                blocked,
            )
            if not selected:
                raise ValueError(
                    "[RollingCharacter] 未选出任何角色，请检查 character_count、排除名单与作品选择"
                )
            if len(selected) < count:
                LOGGER.warning(
                    "[RollingCharacter] 可用角色不足：请求 %d 个，实际选中 %d 个"
                    "（请检查 character_count、作品选择、排除名单或 CSV 规模）",
                    count,
                    len(selected),
                )

            selected_records = [
                record_map.get(name) or CharacterRecord(name, "", "", "", 0) for name in selected
            ]
            prompt = build_character_prompt(selected_records, with_core)

            # 记录里的权重列统一用等长的 1.0 占位（角色不带权重）
            if not db.record(selected, _key_weights(len(selected)), prompt, KIND_CHARACTER):
                LOGGER.error(
                    "[RollingCharacter] 记录写入数据库失败(%s)，本次组合未落盘:\n%s",
                    db.path,
                    prompt,
                )
            # 在事务内取计数：与本次写入处于同一快照，避免事务外再查导致不一致
            tested_total = db.count(KIND_CHARACTER)

        # 记进“最近记录”窗口（放在事务之后：只有真正输出的组合才占用窗口）
        self._remember(selected)

        # 人气池覆盖全部候选时，所有人都是 Top，无需构建集合
        top_names: Optional[Set[str]] = None
        if top_count < len(names):
            top_names = set(names[:top_count])

        payload = build_payload(selected_records, top_names)
        return (
            prompt,
            json.dumps(payload, ensure_ascii=False),
            str(tested_total),
        )


# 注册节点类
NODE_CLASS_MAPPINGS = {"RollingCharacter": RollingCharacter}
# 设置节点显示名称
NODE_DISPLAY_NAME_MAPPINGS = {"RollingCharacter": "Rolling Character"}