"""图谱答案生成节点

输入: state['input']            用户原始提问
      state['cypher_query']     执行过的语句，用来告诉模型每份结果是问什么的
      state['cypher_results']   run_cypher_node 查回来的行
输出: state['neo4j_answer']     根据查询结果组织的自然语言回答
      state['output']           同一份回答，整张图的统一出口

把 list[dict] 交给模型转写成人话，而不是在 Python 里拼字符串：返回列名由生成的 cypher
决定（herb / effect / formula ...），拼接逻辑得跟着 cypher 变，交给模型更省事。

约束模型只用查询结果作答，是这个节点的重点。图谱问答的价值就在于答案可溯源，
一旦允许模型补充自己的知识，就退化成了普通问答，还看不出哪句话是图里来的。
"""

import json

from langchain_core.messages import HumanMessage, SystemMessage

from _004_langgraph_more_nodes.agent_state import AgentState
from common.llm import my_llm

llm = my_llm

# 结果行数可能很多，全塞进 prompt 既慢又容易超长。每条语句最多给这么多行。
MAX_ROWS_PER_QUERY = 30

SYSTEM_PROMPT = (
    "你是一个中医知识问答助手。下面给你的是从中医知识图谱里查询出来的结果，"
    "请依据这些结果回答用户的提问。\n\n"
    "【要求】\n"
    "1. 只使用查询结果中的信息作答，不要补充你自己知道的中医知识，也不要推测。\n"
    "2. 查询结果为空时，直接说明图谱中没有查到相关信息，不要编造。\n"
    "3. 结果条目多时归类整理，用简洁的中文分点表述，不要罗列原始字段名。\n"
    "4. 不要提及 cypher、数据库、查询语句这些技术细节，直接回答内容。\n"
    "5. 涉及具体用药时，末尾提示需遵医嘱。"
)

EMPTY_ANSWER = "抱歉，图谱中没有查询到与您的问题相关的信息。"


def _build_result_block(state: AgentState) -> str:
    """把每条语句和它的结果配对，拼成给模型看的文本"""
    cyphers = state.get("cypher_query") or []
    results = state.get("cypher_results") or []

    blocks = []
    for i, rows in enumerate(results):
        # cypher_results 和 cypher_query 下标对齐，但防一手长度不一致
        cypher = cyphers[i] if i < len(cyphers) else ""
        shown = rows[:MAX_ROWS_PER_QUERY]
        omitted = len(rows) - len(shown)

        lines = [f"【查询 {i + 1}】{cypher}"]
        if not rows:
            lines.append("结果：无")
        else:
            lines.append(f"结果（共 {len(rows)} 条）：")
            for row in shown:
                lines.append(f"  {json.dumps(row, ensure_ascii=False)}")
            if omitted > 0:
                lines.append(f"  ...另有 {omitted} 条已省略")
        blocks.append("\n".join(lines))

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


def generate_neo4j_answer_node(state: AgentState) -> AgentState:
    print("图谱答案节点执行开始了")
    user_input = (state.get("input_semantic_trans") or "").strip()
    results = state.get("cypher_results") or []

    # 所有语句都没查回东西，不必调用大模型 —— 模型能说的也只是"没查到"
    if not user_input or not any(results):
        state["neo4j_answer"] = EMPTY_ANSWER
        state["output"] = EMPTY_ANSWER
        print("完成生成图谱答案: 无结果")
        return state

    messages = [
        SystemMessage(content=SYSTEM_PROMPT),
        HumanMessage(
            content=(
                f"用户提问：{user_input}\n\n"
                f"图谱查询结果：\n{_build_result_block(state)}"
            )
        ),
    ]
    answer = llm.invoke(messages).content.strip()

    state["neo4j_answer"] = answer
    state["output"] = answer
    history_messages = state.get("history_messages", [])
    history_messages.append({"role": "assistant", "content": answer})
    state["history_messages"] = history_messages

    print(f"完成生成图谱答案: {len(answer)} 字")
    return state


if __name__ == "__main__":
    cases = [
        {
            "input": "四君子汤由哪些药材组成？",
            "cypher_query": [
                "MATCH (f:Formula)-[:HAS_INGREDIENT]->(h:Herb) "
                "WHERE f.name = '四君子汤' RETURN h.name AS herb LIMIT 20"
            ],
            "cypher_results": [
                [{"herb": "甘草"}, {"herb": "人参"}, {"herb": "茯苓"}, {"herb": "白术"}]
            ],
        },
        {
            "input": "人参有什么功效？",
            "cypher_query": [
                "MATCH (h:Herb)-[:HAS_EFFECT]->(e:Effect) WHERE h.name IN ['人参'] "
                "RETURN h.name AS herb, e.name AS effect LIMIT 20"
            ],
            "cypher_results": [
                [
                    {"herb": "人参", "effect": "安神"},
                    {"herb": "人参", "effect": "益气温阳"},
                    {"herb": "人参", "effect": "补脾健肺"},
                    {"herb": "人参", "effect": "大补元气"},
                ]
            ],
        },
        # 查询结果全空：预期不调用大模型，直接给兜底话术
        {
            "input": "阿司匹林是什么药材？",
            "cypher_query": ["MATCH (h:Herb) WHERE h.name = '阿司匹林' RETURN h.name AS herb"],
            "cypher_results": [[]],
        },
    ]

    for case in cases:
        result = generate_neo4j_answer_node(dict(case))
        print(f"\n提问: {case['input']}")
        print(f"  {result['neo4j_answer'][:200]}")
        print(f"  output 一致: {result['output'] == result['neo4j_answer']}")