"""
LangGraph 实战 - 接入 OpenAI LLM 的智能体
构建一个能"思考-行动-反思"的真实 AI 智能体
"""
import os
from langgraph.graph import StateGraph, START, END
from typing import TypedDict, Annotated
from functools import reduce

# 检查 API Key
OPENAI_KEY = os.environ.get("OPENAI_API_KEY", "")
if not OPENAI_KEY:
    print("请先设置 OPENAI_API_KEY:")
    print("  $env:OPENAI_API_KEY = 'sk-...'  # PowerShell")
    print()
    print("暂时使用模拟模式演示...")
    USE_MOCK = True
else:
    USE_MOCK = False

# 1️⃣ 定义状态
class AgentState(TypedDict):
    query: str
    thought: str
    action: str
    observation: str
    answer: str
    max_iterations: int
    iteration: int

# 2️⃣ 模拟 LLM 调用
def mock_llm(prompt: str) -> str:
    """模拟 LLM 响应"""
    responses = {
        "思考": "我需要先分析问题，然后给出解答。",
        "行动": "搜索相关信息并整理答案。",
        "回答": "根据分析，LangGraph 是一个用于构建有状态 AI 应用的框架。",
    }
    for key, value in responses.items():
        if key in prompt:
            return value
    return "这是一个有趣的问题。"

def real_llm(prompt: str) -> str:
    """真实 LLM 调用"""
    from openai import OpenAI
    client = OpenAI()
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=200,
        temperature=0.7,
    )
    return response.choices[0].message.content

# 3️⃣ 定义节点
def think_node(state: AgentState):
    """思考节点 - 分析当前情况"""
    state = state.copy()
    state['iteration'] += 1
    print(f"\n{'='*40}")
    print(f"[思考] 第 {state['iteration']}/{state['max_iterations']} 轮")
    
    if USE_MOCK:
        state['thought'] = mock_llm(f"思考: {state['query']}")
    else:
        state['thought'] = real_llm(f"思考: {state['query']}")
    
    print(f"  想法: {state['thought']}")
    return state

def act_node(state: AgentState):
    """行动节点 - 执行操作"""
    state = state.copy()
    print(f"[行动] 执行操作...")
    
    if USE_MOCK:
        state['action'] = mock_llm("行动")
        state['observation'] = f"观察到: {state['action']}的结果"
    else:
        state['action'] = real_llm(f"针对问题 '{state['query']}' 我该采取什么行动？")
        state['observation'] = f"执行了: {state['action']}"
    
    print(f"  行动: {state['action']}")
    print(f"  观察: {state['observation']}")
    return state

def answer_node(state: AgentState):
    """回答节点 - 给出最终答案"""
    state = state.copy()
    print(f"\n[回答] 生成最终答案...")
    
    if USE_MOCK:
        state['answer'] = mock_llm("回答")
    else:
        prompt = f"""基于以下信息回答问题:
问题: {state['query']}
思考过程: {state['thought']}
行动: {state['action']}
观察: {state['observation']}
请给出简洁的答案:"""
        state['answer'] = real_llm(prompt)
    
    print(f"  答案: {state['answer']}")
    return state

# 4️⃣ 路由函数
def should_continue(state: AgentState):
    """决定是否继续迭代"""
    if state['iteration'] >= state['max_iterations']:
        return "answer"
    observation = state.get('observation', '')
    if observation and len(observation) > 10:
        return "answer"
    return "act"

# 5️⃣ 构建图
graph = StateGraph(AgentState)
graph.add_node("think", think_node)
graph.add_node("act", act_node)
graph.add_node("answer", answer_node)

# 连接
graph.add_edge(START, "think")
graph.add_conditional_edges("think", should_continue, {"act": "act", "answer": "answer"})
graph.add_edge("act", "answer")
graph.add_edge("answer", END)

# 编译
app = graph.compile()

# 6️⃣ 运行
print("=" * 40)
print("LangGraph AI 智能体演示")
print("=" * 40)

result = app.invoke({
    "query": "什么是 LangGraph？",
    "thought": "",
    "action": "",
    "observation": "",
    "answer": "",
    "max_iterations": 2,
    "iteration": 0,
})

print(f"\n{'='*40}")
print(f"最终答案: {result['answer']}")
print(f"{'='*40}")