"""
LangGraph 循环 - 智能体的核心能力
构建一个简单的"尝试-评估-重试"智能体
"""
from langgraph.graph import StateGraph, START, END
from typing import TypedDict

# 1️⃣ 定义状态 - 累积信息
class AgentState(TypedDict):
    task: str          # 任务
    attempt: int       # 尝试次数
    result: str        # 当前结果
    quality: str       # 质量评估: "good" / "needs_improvement" / "bad"
    feedback: str      # 改进建议

# 2️⃣ 定义节点
def plan_node(state: AgentState):
    """规划节点 - 制定初始方案"""
    print(f"\n[规划] 任务: {state['task']}")
    print(f"[规划] 制定初始方案...")
    
    # 模拟规划逻辑
    if "诗" in state['task']:
        state['result'] = "床前明月光"
        state['quality'] = "needs_improvement"
    elif "代码" in state['task'] or "code" in state['task'].lower():
        state['result'] = "print(hello)"  # 有 bug
        state['quality'] = "bad"
    else:
        state['result'] = "这是一个初步方案"
        state['quality'] = "needs_improvement"
    
    state['attempt'] = 1
    print(f"[规划] 第 {state['attempt']} 次尝试: {state['result']}")
    return state

def execute_node(state: AgentState):
    """执行节点 - 根据反馈改进"""
    state['attempt'] += 1
    print(f"\n[执行] 第 {state['attempt']} 次尝试...")
    print(f"[执行] 收到反馈: {state['feedback']}")
    
    # 模拟根据反馈改进
    if "诗" in state['task']:
        if state['attempt'] == 2:
            state['result'] = "床前明月光，疑是地上霜"
        else:
            state['result'] = "床前明月光，疑是地上霜。举头望明月，低头思故乡。"
    elif "代码" in state['task'] or "code" in state['task'].lower():
        if state['attempt'] == 2:
            state['result'] = "print('hello')"  # 修复了
        else:
            state['result'] = "print('hello, world!')  # 完美!"
    
    print(f"[执行] 改进后: {state['result']}")
    return state

def evaluate_node(state: AgentState):
    """评估节点 - 判断质量"""
    print(f"\n[评估] 检查第 {state['attempt']} 次结果...")
    
    # 模拟评估逻辑
    if state['attempt'] >= 3:
        state['quality'] = "good"
        state['feedback'] = "质量满意，通过！"
        print(f"[评估] 结果: {state['quality']} - {state['feedback']}")
    else:
        if "诗" in state['task']:
            state['feedback'] = "内容太短，需要更完整"
        elif "代码" in state['task'] or "code" in state['task'].lower():
            state['feedback'] = "代码有语法错误，需要修复"
        else:
            state['feedback'] = "需要更详细的内容"
        state['quality'] = "needs_improvement"
        print(f"[评估] 结果: {state['quality']} - {state['feedback']}")
    
    return state

# 3️⃣ 路由函数 - 决定是继续循环还是结束
def should_retry(state: AgentState):
    """⭐ 核心：决定是重试还是结束"""
    if state['quality'] == "good":
        return "end"          # 质量够了，结束
    elif state['attempt'] >= 3:
        return "end"          # 达到最大重试次数，强制结束
    else:
        return "retry"        # 继续改进

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

# 添加节点
graph.add_node("plan", plan_node)
graph.add_node("execute", execute_node)
graph.add_node("evaluate", evaluate_node)

# 添加边 - 注意这里的循环！
graph.add_edge(START, "plan")           # 开始 -> 规划
graph.add_edge("plan", "evaluate")      # 规划 -> 评估
graph.add_conditional_edges(
    "evaluate",
    should_retry,
    {
        "retry": "execute",   # 需要改进 -> 执行(改进)
        "end": END,           # 满意 -> 结束
    }
)
graph.add_edge("execute", "evaluate")   # 执行完再评估 -> 形成循环！

# 编译
app = graph.compile()

# 5️⃣ 测试
print("=" * 55)
print("循环演示 - 智能体不断尝试直到满意")
print("=" * 55)

# 测试 1: 写诗
print("\n" + "=" * 55)
print("测试 1: 写一首诗")
print("=" * 55)
result = app.invoke({
    "task": "写一首诗",
    "attempt": 0,
    "result": "",
    "quality": "",
    "feedback": ""
})
print(f"\n最终结果: {result['result']}")
print(f"尝试次数: {result['attempt']}")

# 测试 2: 写代码
print("\n" + "=" * 55)
print("测试 2: 写一段代码")
print("=" * 55)
result = app.invoke({
    "task": "写一段代码",
    "attempt": 0,
    "result": "",
    "quality": "",
    "feedback": ""
})
print(f"\n最终结果: {result['result']}")
print(f"尝试次数: {result['attempt']}")