"""三元组文本化 + 上下文构建。

将图谱子图（实体 + 关系）转为自然语言描述，
并与向量检索片段合并为 LLM 可消费的上下文文本。
"""

from __future__ import annotations

from app.module.kg_graphrag.models import (
    EntitySummary,
    RelationSummary,
    VectorChunk,
)

# 关系类型 -> 中文模板映射
RELATION_TEMPLATES: dict[str, str] = {
    "HAS_FIELD": "{source}包含字段{target}",
    "DERIVED_FROM": "{source}来源于{target}",
    "USES_DATASET": "{source}使用了数据集{target}",
    "PRODUCES": "{source}产出了{target}",
    "ASSIGNED_TO": "{source}分配给了{target}",
    "BELONGS_TO": "{source}属于{target}",
    "TRIGGERS": "{source}触发了{target}",
    "DEPENDS_ON": "{source}依赖于{target}",
    "IMPACTS": "{source}影响了{target}",
    "SOURCED_FROM": "{source}的知识来源于{target}",
}

# 通用模板（未在映射中的关系类型）
_DEFAULT_TEMPLATE = "{source}与{target}存在{relation}关系"


def textualize_subgraph(
    entities: list[EntitySummary],
    relations: list[RelationSummary],
) -> str:
    """将图谱子图转为自然语言描述。

    Args:
        entities: 子图中的实体列表。
        relations: 子图中的关系列表。

    Returns:
        文本化后的图谱描述，每条关系/实体一行。
    """
    lines: list[str] = []

    # 记录有关系的实体名称
    mentioned_entities: set[str] = set()

    # 1. 对每条关系生成一句话
    for rel in relations:
        source_label = f"{rel.source_type}'{rel.source_name}'"
        target_label = f"{rel.target_type}'{rel.target_name}'"
        template = RELATION_TEMPLATES.get(rel.relation_type, _DEFAULT_TEMPLATE)
        line = template.format(
            source=source_label,
            target=target_label,
            relation=rel.relation_type,
        )
        lines.append(line)
        mentioned_entities.add(rel.source_name)
        mentioned_entities.add(rel.target_name)

    # 2. 对独立实体（无关系）生成描述句
    for entity in entities:
        if entity.name not in mentioned_entities:
            desc = entity.description or ""
            if desc:
                lines.append(f"{entity.type}'{entity.name}': {desc}")
            else:
                lines.append(f"存在{entity.type}'{entity.name}'")

    return "\n".join(lines)


def build_context(
    vector_chunks: list[VectorChunk],
    graph_text: str,
    vector_weight: float = 0.6,
    graph_weight: float = 0.4,
) -> str:
    """合并向量检索片段和图谱文本化内容为 LLM 上下文。

    Args:
        vector_chunks: 向量检索到的文档片段列表。
        graph_text: 文本化后的图谱描述。
        vector_weight: 向量分数权重（当前用于日志/调试，不影响上下文排序）。
        graph_weight: 图谱相关性权重。

    Returns:
        合并后的上下文文本，分为「相关文档」和「知识图谱上下文」两个部分。
    """
    sections: list[str] = []

    # 向量检索片段
    if vector_chunks:
        doc_lines = ["## 相关文档"]
        for i, chunk in enumerate(vector_chunks, 1):
            doc_lines.append(f"[{i}] {chunk.text}")
        sections.append("\n".join(doc_lines))

    # 图谱文本化内容
    if graph_text:
        sections.append(f"## 知识图谱上下文\n{graph_text}")

    if not sections:
        return "（未检索到相关上下文信息）"

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