Agent/examples/demand/tree_service.py

from __future__ import annotations

from collections.abc import Iterable
from pathlib import Path

from sqlalchemy import text

from examples.demand.db_manager import DatabaseManager

db = DatabaseManager()


def _fmt_node(name: str, level: int | None) -> str:
    if level is None:
        return f"{name}[L?]"
    return f"{name}[L{int(level)}]"


def _fetch_roots_by_name(session, execution_id: int, name: str):
    return session.execute(
        text(
            """
            SELECT id, name, level
            FROM topic_pattern_category
            WHERE execution_id = :execution_id AND name = :name
            ORDER BY id
            """
        ),
        {"execution_id": execution_id, "name": name},
    ).fetchall()


def _fetch_children_by_parent(session, execution_id: int, parent_id: int):
    return session.execute(
        text(
            """
            SELECT id, name, level
            FROM topic_pattern_category
            WHERE execution_id = :execution_id AND parent_id = :parent_id
            ORDER BY id
            """
        ),
        {"execution_id": execution_id, "parent_id": parent_id},
    ).fetchall()


def _collect_subtree_lines(
        session,
        execution_id: int,
        node_id: int,
        name: str,
        level: int | None,
        depth: int,
) -> list[str]:
    indent = "  " * depth
    lines = [f"{indent}{_fmt_node(name, level)}"]
    for row in _fetch_children_by_parent(session, execution_id, node_id):
        cid, cname, clevel = row[0], row[1], row[2]
        lines.extend(
            _collect_subtree_lines(session, execution_id, cid, cname, clevel, depth + 1)
        )
    return lines


def build_category_trees_text(execution_id: int, names: Iterable[str]) -> str:
    """
    按 execution_id 与名称列表，在 topic_pattern_category 中从每个 name 对应行出发，
    沿 parent_id 向下遍历整棵子树，返回多棵树拼接的文本（子节点每深一层缩进 2 空格）。

    每行格式：name[L{level}]（level 为表字段；若为空则显示 L?）。
    输入列表中的每个 name 输出一块；若同名匹配多行，则该块内为多棵子树，树之间空一行。
    某 name 无匹配时，仅输出一行 name[L?]。
    """
    clean_names = [str(n).strip() for n in names if n is not None and str(n).strip()]
    if not clean_names:
        return ""

    blocks: list[str] = []
    session = db.get_session()
    try:
        for name in clean_names:
            roots = _fetch_roots_by_name(session, execution_id, name)
            if not roots:
                blocks.append(_fmt_node(name, None))
                continue
            tree_texts: list[str] = []
            for row in roots:
                rid, rname, rlevel = row[0], row[1], row[2]
                tree_texts.append(
                    "\n".join(
                        _collect_subtree_lines(session, execution_id, rid, rname, rlevel, 0)
                    )
                )
            blocks.append("\n\n".join(tree_texts))
    finally:
        session.close()

    return "\n\n".join(blocks)


if __name__ == "__main__":
    execution_id = 58
    names = ["健康养生"
        , "公共安全"
        , "地标景观"
        , "防诈反骗"
        , "身体健康"
        , "饮食调理"
        , "地形水域"
        , "建筑地标"
        , "自然地标"
        , "医疗卫生"
        , "综合风光"
        , "居家生活"
        , "办事指南"
        , "政策法规"
        , "自然奇景"
        , "中医养生"
        , "烹饪技巧"
        , "安全防护"
        , "皮肤护理"
        , "地质奇观"
        , "社会保障"
        , "心理认知"
        , "办事规程"
        , "养老规划"
        , "金融理财"
        , "农业农村"
        , "补贴福利"
        , "禁毒防毒"
        , "国家安全"
        , "设施安全"]
    text_out = build_category_trees_text(execution_id, names)
    out_path = Path(__file__).resolve().parent / "new_result" / f"category_tree_{execution_id}.txt"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(text_out, encoding="utf-8")
    print(out_path)