"""
LangGraph 条件边 - 让图学会做决定
"""
from langgraph.graph import StateGraph, START, END
from typing import TypedDict

# 1️⃣ 定义状态
class GraphState(TypedDict):
    message: str
    user_input: str

# 2️⃣ 定义节点
def receive_node(state: GraphState):
    """接收节点 - 分类用户输入"""
    text = state['user_input'].strip().lower()
    print(f"[接收] 用户说: {state['user_input']}")
    
    # 简单分类
    if any(word in text for word in ['你好', 'hello', 'hi', 'hey']):
        state['message'] = "问候"
    elif '?' in text or '？' in text or '怎么' in text or '什么' in text:
        state['message'] = "问题"
    else:
        state['message'] = "闲聊"
    
    print(f"[接收] 分类为: {state['message']}")
    return state

def respond_greeting(state: GraphState):
    """回应问候"""
    state['message'] = "你好呀！有什么可以帮你的吗？"
    print(f"[问候回应] {state['message']}")
    return state

def answer_question(state: GraphState):
    """回答问题"""
    state['message'] = "这是一个好问题！让我想想... (这里可以接入 LLM)"
    print(f"[问题回答] {state['message']}")
    return state

def chat_casual(state: GraphState):
    """闲聊"""
    state['message'] = "哈哈，聊得开心！"
    print(f"[闲聊] {state['message']}")
    return state

# 3️⃣ 条件路由函数 - 这是关键！
def route_input(state: GraphState):
    """根据分类决定下一步走哪个节点"""
    category = state['message']
    print(f"[路由] 决定走向: {category}")
    
    if category == "问候":
        return "respond_greeting"
    elif category == "问题":
        return "answer_question"
    else:
        return "chat_casual"

# 4️⃣ 构建图
graph = StateGraph(GraphState)

# 添加所有节点
graph.add_node("receive", receive_node)
graph.add_node("respond_greeting", respond_greeting)
graph.add_node("answer_question", answer_question)
graph.add_node("chat_casual", chat_casual)

# 添加边
graph.add_edge(START, "receive")

# 条件边！根据路由函数返回值决定走向
graph.add_conditional_edges(
    "receive",           # 从接收节点出发
    route_input,         # 用这个函数做判断
    {
        "respond_greeting": "respond_greeting",
        "answer_question": "answer_question",
        "chat_casual": "chat_casual",
    }
)

# 所有分支最后都汇聚到 END
graph.add_edge("respond_greeting", END)
graph.add_edge("answer_question", END)
graph.add_edge("chat_casual", END)

# 编译
app = graph.compile()

# 5️⃣ 测试三种情况
test_cases = [
    "你好！",
    "什么是 LangGraph？",
    "今天天气不错",
]

print("=" * 50)
print("条件边演示 - 图会根据输入走不同路径")
print("=" * 50)

for i, text in enumerate(test_cases, 1):
    print(f"\n--- 测试 {i}: {text} ---")
    result = app.invoke({"user_input": text, "message": ""})
    print(f"  最终输出: {result['message']}")