"""
LangGraph 第 5 步：真实 LLM 智能体
ReAct 模式：思考(Reason) - 行动(Act) - 观察(Observe)
支持自定义 OpenAI 兼容 API
"""
import os
import sys
from langgraph.graph import StateGraph, START, END
from typing import TypedDict

# 1️⃣ 配置 API
API_KEY = os.environ.get("OPENAI_API_KEY", "")
BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
MODEL = os.environ.get("MODEL_NAME", "gpt-4o-mini")

if not API_KEY:
    print("请先设置环境变量:")
    print("  $env:OPENAI_API_KEY = 'your-api-key'")
    print("  $env:OPENAI_BASE_URL = 'your-api-base-url'  # 可选，默认 OpenAI")
    print("  $env:MODEL_NAME = 'gpt-4o-mini'  # 可选")
    exit(1)

from openai import OpenAI
client = OpenAI(api_key=API_KEY, base_url=BASE_URL)

# 2️⃣ 定义状态
class AgentState(TypedDict):
    question: str
    thoughts: list
    current_thought: str
    action: str
    action_param: str
    observation: str
    final_answer: str
    iteration: int
    max_iterations: int

# 3️⃣ 定义工具
def calculator(expression: str) -> str:
    """计算器工具"""
    try:
        result = eval(expression, {"__builtins__": {}}, {})
        return f"计算结果: {expression} = {result}"
    except Exception as e:
        return f"计算错误: {e}"

def search_knowledge(query: str) -> str:
    """知识库搜索工具"""
    knowledge = {
        "langgraph": "LangGraph 是 LangChain 团队开发的框架，用于构建有状态、基于图的 AI 应用。",
        "python": "Python 是一种高级编程语言，广泛用于 AI、Web 开发、数据分析等领域。",
        "langchain": "LangChain 是构建 LLM 应用的框架，提供 Prompt 管理、Chain、Agent 等组件。",
        "ai": "人工智能 (AI) 是计算机科学的一个分支，致力于创建能执行智能任务的系统。",
    }
    for key, value in knowledge.items():
        if key in query.lower():
            return f"搜索到: {value}"
    return f"未找到关于 '{query}' 的精确信息"

tools = {
    "calculator": calculator,
    "search": search_knowledge,
}

# 4️⃣ 定义节点
def think_node(state: AgentState):
    """思考节点 - 让 LLM 决定下一步"""
    state = state.copy()
    state['iteration'] += 1

    system_prompt = f"""你是一个智能助手，可以使用以下工具:
1. calculator - 数学计算，参数是数学表达式如 "2+3*4"
2. search - 搜索知识，参数是搜索关键词

当前是第 {state['iteration']}/{state['max_iterations']} 轮。

请严格按照以下格式回复:
[思考] 你的思考过程
[行动] 工具名称|参数
例如:
[思考] 我需要计算这个数学题
[行动] calculator|2+3*4

如果可以直接回答，请这样回复:
[思考] 我已经知道答案了
[回答] 你的最终答案"""

    messages = [{"role": "system", "content": system_prompt}]
    messages.append({"role": "user", "content": state['question']})

    if state.get('observation'):
        messages.append({"role": "assistant", "content": f"[观察] {state['observation']}"})

    response = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        max_tokens=300,
        temperature=0.3,
    )

    thought_text = response.choices[0].message.content
    state['current_thought'] = thought_text
    state['thoughts'] = state.get('thoughts', []) + [thought_text]

    print(f"\n{'='*50}")
    print(f"[思考] 第 {state['iteration']} 轮:")
    print(thought_text)

    # 解析行动
    for line in thought_text.split('\n'):
        if '[行动]' in line:
            parts = line.replace('[行动]', '').strip().split('|')
            if len(parts) == 2:
                state['action'] = parts[0].strip()
                state['action_param'] = parts[1].strip()
                return state
        if '[回答]' in line:
            state['final_answer'] = line.replace('[回答]', '').strip()
            return state

    state['action'] = ""
    return state

def act_node(state: AgentState):
    """行动节点 - 执行工具"""
    state = state.copy()
    action = state.get('action', '')
    param = state.get('action_param', '')

    print(f"\n[行动] 执行 {action}({param})")

    if action in tools:
        result = tools[action](param)
        state['observation'] = result
        print(f"[观察] {result}")
    else:
        state['observation'] = f"未知工具: {action}"
        print(f"[观察] 未知工具: {action}")

    return state

def answer_node(state: AgentState):
    """回答节点 - 生成最终答案"""
    state = state.copy()

    if state.get('final_answer'):
        print(f"\n[回答] {state['final_answer']}")
        return state

    messages = [
        {"role": "system", "content": "请根据以下信息给出简洁的最终答案"},
        {"role": "user", "content": f"问题: {state['question']}\n\n思考过程:\n" + "\n".join(state.get('thoughts', []))}
    ]

    response = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        max_tokens=200,
    )

    state['final_answer'] = response.choices[0].message.content
    print(f"\n[回答] {state['final_answer']}")
    return state

# 5️⃣ 路由函数
def route(state: AgentState):
    """决定下一步"""
    if state.get('final_answer'):
        return "answer"
    if state['iteration'] >= state['max_iterations']:
        return "answer"
    if state.get('action'):
        return "act"
    return "answer"

# 6️⃣ 构建图
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", route, {"act": "act", "answer": "answer"})
graph.add_edge("act", "think")
graph.add_edge("answer", END)

app = graph.compile()

# 7️⃣ 运行
print("=" * 50)
print("LangGraph ReAct 智能体")
print("=" * 50)
print(f"API: {BASE_URL}")
print(f"模型: {MODEL}")
print("\n图结构:")
print("  START -> think -> [有行动?] -> act -> think (循环)")
print("              |")
print("              +-> [无行动/达到限制] -> answer -> END")

# 如果有命令行参数，用命令行参数作为问题
if len(sys.argv) > 1:
    questions = [" ".join(sys.argv[1:])]
else:
    questions = [
        "计算一下 123 * 456",
        "什么是 LangGraph？",
    ]

for q in questions:
    print(f"\n{'#'*50}")
    print(f"问题: {q}")
    print(f"{'#'*50}")

    result = app.invoke({
        "question": q,
        "thoughts": [],
        "iteration": 0,
        "max_iterations": 4,
        "action": "",
        "observation": "",
        "final_answer": "",
        "action_param": "",
    })

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