from langchain_core.messages import HumanMessage

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

conf = Config()


def semantic_transcription_node(state: AgentState) -> AgentState:
    """
    语义转写节点：
    根据历史对话内容（history_messages）和当前用户输入（input），
    生成一个更清晰、标准化的语义转写结果，用于下游意图识别或问题理解。

    这是图里唯一读 state["input"]（用户原话）的节点，往后所有节点读的都是
    本节点写出的 input_semantic_trans。因此这个键必须在每条路径上都写非空值 ——
    留空会让下游全部收到空输入，整张图返回空回答。
    """
    # 获取用户输入和历史上下文
    print("语义转写节点执行开始了")
    user_input = (state.get("input") or "").strip()
    history_messages = state.get("history_messages", [])
    print(f"历史消息：{history_messages}")

    # 本节点是每轮的第一个节点，也就是唯一能识别"新一轮开始"的地方，
    # 所以每轮只属于这一轮的键在这里清掉。
    # 加了 checkpointer 之后 state 会按 thread_id 跨轮保留，不清的话：
    #   - cypher_retry_count 只增不减，第 4 轮时已经到上限 3，这一轮的 cypher
    #     只要有一条不合法就直接放弃重试，额度被前几轮吃光了；
    #   - neo4j_answer / direct_out 只在各自那条分支写，走另一条分支时留着上一轮的值，
    #     路径判断（_summarize 用 neo4j_answer 判断走没走图谱）会看错；
    #   - output 万一哪条路径没写，会把上一轮的答案返回给这一轮的问题 —— 清空后
    #     退化成空串，_005_ 那边会正常兜底成错误文案。
    # history_messages 和 input_semantic_trans 不能清：前者就是跨轮记忆本身，后者下面就写。
    state["cypher_retry_count"] = 0
    for stale_key in (
        "neo4j_answer",
        "direct_out",
        "output",
        "cypher_query",
        "cypher_results",
        "is_all_validate_cypher",
    ):
        state[stale_key] = None

    # 空输入：没什么可转写的，也不必调大模型。下游靠 input_semantic_trans 判空短路
    if not user_input:
        state["input_semantic_trans"] = ""
        state["history_messages"] = history_messages
        print("语义转写节点：空输入，跳过")
        return state

    # 更新历史
    history_messages.append({"role": "user", "content": user_input})
    state["history_messages"] = history_messages
    # 只保留最近 history_num 轮。切片放在 append 之后、末尾去掉一条，
    # 去掉的正是刚加进去的这轮 —— 当前输入单独放在提示词的【当前用户输入】里，
    # 混进【历史对话】会让模型把它当成上一轮的话
    history_messages = history_messages[-(conf.history_num + 1):-1]

    # 第一轮没有历史，没有指代词可还原，转写只会把原问题改写一遍（甚至改坏），
    # 白花一次调用。直接用原话
    if not history_messages:
        state["input_semantic_trans"] = user_input
        print(f"语义转写节点：无历史，直接用原输入：{user_input}")
        return state

    # 构建提示词
    prompt = f"""
你是一名语言理解专家，负责将用户输入转写为清晰、简洁、语义明确的问题或意图表达。
你需要基于以下上下文（最近6轮对话），分析用户真正想表达的意思。
请不要编造不存在的信息，也不要直接回答问题，只需输出语义转写后的句子。

【历史对话】
{history_messages}

【当前用户输入】
{user_input}

请输出：
1. 一句自然语言，清楚表达用户当前真正的问题或意图；
2. 不要包含解释性文字；
3. 如果用户输入含糊不清，请基于上下文合理补全语义；
4. 输出格式：直接输出转写结果，无需加前缀。
"""

    # 调用大模型
    response = my_llm.invoke([HumanMessage(content=prompt)])
    result = response.content.strip()

    # 模型偶尔会返回空串。下游全靠这个键，空了等于整张图收不到输入，所以退回原话
    if not result:
        result = user_input
        print("语义转写节点：模型返回空，退回原输入")

    # 保存语义转写结果
    state["input_semantic_trans"] = result
    print(f"完成生成语义转写：{result}")
    return state


if __name__ == '__main__':
    # 每条覆盖一种情况：有历史要还原指代 / 中医多轮 / 无历史 / 空输入
    cases = [
        {
            "input": "那第二呢?",
            "history_messages": [
                {"role": "user", "content": "世界第一高峰是什么?"},
                {"role": "assistant", "content": "世界第一高峰是珠穆朗玛峰。"},
            ],
        },
        {
            "input": "它有什么副作用？",
            "history_messages": [
                {"role": "user", "content": "人参有什么功效？"},
                {"role": "assistant", "content": "人参能大补元气、复脉固脱、补脾益肺。"},
            ],
        },
        # 无历史：预期原样返回，不调用大模型
        {"input": "人参有什么功效？", "history_messages": []},
        # 空输入：预期空串，不调用大模型
        {"input": "", "history_messages": []},
    ]

    for case in cases:
        result_state = semantic_transcription_node(dict(case))
        print(f"\n原输入: {case['input']!r}")
        print(f"  转写后: {result_state['input_semantic_trans']!r}")
        print(f"  历史轮数: {len(result_state['history_messages'])}")