"""把 nodes/ 下的节点串成一张 LangGraph 图

    START → 语义转写节点 → 中医意图识别节点
              ├─ 非中医 ─→ 直接回答节点 ─────────────────────────→ END
              └─ 中医  ─→ 实体抽取节点 → 实体匹配节点 → cypher生成节点 ←──┐
                              → cypher检查节点                          │
                                   ├─ 全部合法 ─→ cypher执行节点         │
                                   │                  → 图谱答案节点 → END│
                                   ├─ 不合法且未超 3 次 ─→ 重新生成 ──────┘
                                   └─ 不合法且已 3 次 ─→ 直接回答节点 → END

语义转写排在最前面，是因为它后面每一个节点都依赖它的产出：多轮对话里"那第二呢？"
这种话单看一句没有任何信息，意图识别会判成非中医、实体抽取一个词也抽不出来。转写节点
结合 history_messages 把它补成"珠穆朗玛峰之后第二高的山峰是什么？"，下游才有东西可用。

所以 `input` 只有转写节点自己读（它的输入就是原话），其余节点一律读
`input_semantic_trans`。两者的区别是"用户打了什么"和"用户想问什么"。

两个条件边，两处兜底，一个环：

- 意图识别之后：跟中医无关的提问查图谱只会白跑一圈，直接交给大模型。
- cypher 检查之后：不合法先回到生成节点重写，最多生成 3 次（首次 + 2 次重试）；
  还是不行才退回大模型。计数在 `cypher_retry_count`，生成节点里自增 —— 路由函数
  只读 state，没法自己记账，所以计数必须由节点写。
- `直接回答节点` 被两条兜底路径共用（LangGraph 允许多个前驱指向同一节点）。

重试时会把上一轮不合法的语句回喂给模型（见生成节点里的 `retry_block`），否则模型
看到的输入和上次一模一样，大概率原样再生成一遍。

还有一处兜底没有做成条件边：cypher 生成为空（一个实体都没匹配上）。这种情况
`cypher_query` 是 `[]`，而 `check_cypher_node` 对空列表判 False，于是会走重试
那条边 —— 但生成节点在没有实体时直接返回不调大模型，所以这两次重试是空转、不花钱，
最后由计数刹住落到兜底。省下了为空列表单独写一条分支。

节点名用中文，和 `_000_demo/04_langgraph_demo.py` 保持一致。
"""
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START
from langgraph.graph import END, StateGraph

from _004_langgraph_more_nodes.agent_state import AgentState
from _004_langgraph_more_nodes.nodes.check_cypher_node import check_cypher_node
from _004_langgraph_more_nodes.nodes.extract_entity_from_user_input_node import (
    extract_entity_from_user_input_node,
)
from _004_langgraph_more_nodes.nodes.generate_neo4j_cypher_node import (
    generate_neo4j_cypher_node,
)
from _004_langgraph_more_nodes.nodes.llm_direct_out_node import llm_direct_out_node
from _004_langgraph_more_nodes.nodes.match_entity_from_neo4j_node import (
    match_entity_from_neo4j_node,
)
from _004_langgraph_more_nodes.nodes.neo4j_answer_generate_node import (
    neo4j_answer_generate_node,
)
from _004_langgraph_more_nodes.nodes.run_cypher_node import run_cypher_node
from _004_langgraph_more_nodes.nodes.semantic_transcription_node import (
    semantic_transcription_node,
)
from _004_langgraph_more_nodes.nodes.zhongyi_intent_node import zhongyi_intent_node

# 节点名集中在这里，路由函数和加边都引用它，避免字符串写错了却要到运行时才报错
SEMANTIC_TRANS = "语义转写节点"
INTENT = "中医意图识别节点"
EXTRACT_ENTITY = "实体抽取节点"
MATCH_ENTITY = "实体匹配节点"
GENERATE_CYPHER = "cypher生成节点"
CHECK_CYPHER = "cypher检查节点"
RUN_CYPHER = "cypher执行节点"
NEO4J_ANSWER = "图谱答案节点"
DIRECT_OUT = "直接回答节点"

# cypher 最多生成 3 次：首次 + 2 次重试。超了就别再烧钱了，退回大模型直接回答。
MAX_CYPHER_RETRY = 3


# ======================
# 路由函数
# ======================
def zhongyi_intent_control(state: AgentState) -> str:
    """意图分流：和中医相关走图谱，否则直接回答"""
    if state.get("is_zhongyi_intent"):
        return EXTRACT_ENTITY
    return DIRECT_OUT


def check_cypher_control(state: AgentState) -> str:
    """cypher 分流，三条出路：

    - 全部合法       -> 去查库
    - 不合法且还有额度 -> 回到生成节点重写（这是图里唯一的环）
    - 不合法且额度用尽 -> 退回大模型直接回答

    环的刹车靠 cypher_retry_count，它在生成节点里自增。只看 is_all_validate_cypher
    的话，模型一直生成不出合法语句就会无限打转。
    """
    if state.get("is_all_validate_cypher"):
        return RUN_CYPHER

    retry_count = state.get("cypher_retry_count") or 0
    if retry_count < MAX_CYPHER_RETRY:
        print(f"cypher 不合法，重新生成（已生成 {retry_count} 次，上限 {MAX_CYPHER_RETRY} 次）")
        return GENERATE_CYPHER

    print(f"cypher 已生成 {retry_count} 次仍不合法，转由大模型直接回答")
    return DIRECT_OUT


# ======================
# 建图
# ======================
def build_langgraph():
    """把节点和边装好并编译，返回可以 invoke 的图"""
    graph_builder = StateGraph(AgentState)

    # 加节点：名字 -> 函数。函数签名都是 (state) -> state，直接传进来即可
    graph_builder.add_node(SEMANTIC_TRANS, semantic_transcription_node)
    graph_builder.add_node(INTENT, zhongyi_intent_node)
    graph_builder.add_node(EXTRACT_ENTITY, extract_entity_from_user_input_node)
    graph_builder.add_node(MATCH_ENTITY, match_entity_from_neo4j_node)
    graph_builder.add_node(GENERATE_CYPHER, generate_neo4j_cypher_node)
    graph_builder.add_node(CHECK_CYPHER, check_cypher_node)
    graph_builder.add_node(RUN_CYPHER, run_cypher_node)
    graph_builder.add_node(NEO4J_ANSWER, neo4j_answer_generate_node)
    graph_builder.add_node(DIRECT_OUT, llm_direct_out_node)

    # 入口：先转写，再判意图。顺序不能反 —— 意图识别读的是转写结果
    graph_builder.add_edge(START, SEMANTIC_TRANS)
    graph_builder.add_edge(SEMANTIC_TRANS, INTENT)

    # 条件边 1：意图识别之后分两路。
    # 第三个参数（路由表）必须给：不给的话 LangGraph 不知道这条边可能通向哪些节点，
    # 画出来的图会只剩 START -> 意图识别 -> END，后面的边全看不到。
    graph_builder.add_conditional_edges(
        INTENT,
        zhongyi_intent_control,
        {EXTRACT_ENTITY: EXTRACT_ENTITY, DIRECT_OUT: DIRECT_OUT},
    )

    # 中医这一路的固定顺序
    graph_builder.add_edge(EXTRACT_ENTITY, MATCH_ENTITY)
    graph_builder.add_edge(MATCH_ENTITY, GENERATE_CYPHER)
    graph_builder.add_edge(GENERATE_CYPHER, CHECK_CYPHER)

    # 条件边 2：cypher 合法去查库，不合法回生成节点重试，重试用尽走兜底。
    # 路由表里的 GENERATE_CYPHER 就是那条回边，图里唯一的环。
    graph_builder.add_conditional_edges(
        CHECK_CYPHER,
        check_cypher_control,
        {
            RUN_CYPHER: RUN_CYPHER,
            GENERATE_CYPHER: GENERATE_CYPHER,
            DIRECT_OUT: DIRECT_OUT,
        },
    )

    graph_builder.add_edge(RUN_CYPHER, NEO4J_ANSWER)

    # 三个出口
    graph_builder.add_edge(NEO4J_ANSWER, END)
    graph_builder.add_edge(DIRECT_OUT, END)
    # 添加记忆
    memory = InMemorySaver()
    # compile 之后才是可执行的图
    return graph_builder.compile(memory)


# 模块级单例：_005_ 的接口层导入这个，不用每个请求重建一次图
graph = build_langgraph()


def ask(question: str) -> dict:
    """问一句，返回跑完的完整 state"""
    return graph.invoke({"input": question})


# ======================
# 测试用例
# ======================
# (提问, 期望走图谱那一路)。每条覆盖一种路径或一类实体：
TEST_CASES = [
    # 图谱路：单个药材 + 属性
    ("人参有什么功效？", True),
    # 图谱路：方剂 → 组成药材，全图最主要的关系
    ("四君子汤由哪些药材组成？", True),
    # 图谱路：出处反查，考的是关系方向没写反
    ("《伤寒论》里记载了哪些方剂？", True),
    # 图谱路：口语症状词，考实体标准化（拉肚子 -> 腹泻）
    ("我最近总是拉肚子，有什么方剂可以调理？", True),
    # 兜底路：非中医提问，意图识别就该分流走
    ("帮我写一段 Python 快排代码", False),
    # 兜底路：像中医但图里没有的实体，匹配不到 -> cypher 为空 -> 落到兜底
    ("阿司匹林的功效是什么？", False),
]


def _summarize(state: dict) -> str:
    """一行讲清这一轮走了哪条路"""
    if state.get("neo4j_answer"):
        rows = sum(len(r) for r in state.get("cypher_results") or [])
        return f"图谱路（{len(state.get('cypher_query') or [])} 条 cypher / {rows} 行）"
    return "兜底路（大模型直接回答）"


if __name__ == "__main__":
    # 图结构：不调用大模型也能看，改完边先跑这段确认接线对不对
    # print("=== 图结构 ===")
    # print(graph.get_graph().draw_ascii())

    print("\n=== 端到端测试 ===")
    wrong = 0
    for question, expect_graph_path in TEST_CASES:
        print(f"\n{'-' * 70}\n提问: {question}")
        state = ask(question)

        went_graph_path = bool(state.get("neo4j_answer"))
        flag = "OK " if went_graph_path == expect_graph_path else "MISS"
        if went_graph_path != expect_graph_path:
            wrong += 1

        print(f"[{flag}] {_summarize(state)}")
        print(f"  意图: {state.get('is_zhongyi_intent')}")
        if state.get("is_zhongyi_intent"):
            matched = {
                key.replace("matched_", ""): value
                for key, value in state.items()
                if key.startswith("matched_") and value
            }
            print(f"  匹配实体: {matched or '<无>'}")
            print(f"  cypher 合法: {state.get('is_all_validate_cypher')}")
            for i, cypher in enumerate(state.get("cypher_query") or [], 1):
                print(f"    {i}. {cypher[:100]}")
        print(f"  回答: {(state.get('output') or '')[:180]}")

    print(f"\n{'=' * 70}\n{len(TEST_CASES)} 条用例，路径与预期不符 {wrong} 条")