"""
LangGraph 第 5 步：真实 LLM 智能体
ReAct 模式：思考(Reason) - 行动(Act) - 观察(Observe)
配置见 config.py
"""

# ============================================================
# 导入模块
# ============================================================
import sys              # 命令行参数
import requests         # HTTP 请求（调用 LLM API）
from langgraph.graph import StateGraph, START, END  # LangGraph 核心组件
from typing import TypedDict  # 类型提示，用于定义状态结构

# ============================================================
# 1. 加载配置
# ============================================================
# 从 config.py 导入配置变量
# Python 的 import 语句可以直接导入另一个 .py 文件中的变量
from config import API_KEY, BASE_URL, MODEL, MAX_ITERATIONS, TEMPERATURE

# 检查用户是否修改了配置
if API_KEY == "sk-...":
    print("请先在 config.py 中配置 API_KEY")
    exit(1)  # 退出程序

# ============================================================
# 2. 定义状态（State）
# ============================================================
"""
TypedDict 是 Python 的类型提示工具，用来定义字典的结构。
LangGraph 用 State 在节点之间传递数据，每个节点可以读取和修改它。

为什么用 TypedDict 而不是普通 dict？
- 编辑器可以提供自动补全
- 可以检查键名是否正确
- 代码更清晰，一看就知道有哪些字段
"""
class AgentState(TypedDict):
    question: str          # 用户的问题
    thoughts: list         # 思考历史（list，记录每轮的思考）
    current_thought: str   # 当前这轮的思考
    action: str            # 要采取的行动（如 "calculator"）
    action_param: str      # 行动的参数（如 "123*456"）
    observation: str       # 行动后的观察结果
    final_answer: str      # 最终答案
    iteration: int         # 当前迭代轮次
    max_iterations: int    # 最大迭代次数（防止无限循环）

# ============================================================
# 3. 定义工具（Tools）
# ============================================================
"""
工具是智能体可以调用的函数。
这里定义了两个工具：计算器和知识搜索。

为什么工具是普通函数？
- LangGraph 中工具就是 Python 函数
- 智能体决定调用哪个工具、传什么参数
- 工具执行后返回结果，智能体再根据结果决定下一步
"""

def calculator(expression: str) -> str:
    """
    计算器工具
    
    参数: expression - 数学表达式字符串，如 "2+3*4"
    返回: 计算结果字符串
    
    为什么用 eval 而不是直接计算？
    - eval 可以解析任意数学表达式
    - 但 eval 有安全风险，所以传了空的 __builtins__ 限制权限
    """
    try:
        # eval() 把字符串当 Python 表达式执行
        # {"__builtins__": {}} 是安全限制，防止执行危险代码
        result = eval(expression, {"__builtins__": {}}, {})
        return f"计算结果: {expression} = {result}"
    except Exception as e:
        # 如果计算出错（如除零、无效表达式），返回错误信息
        return f"计算错误: {e}"

def search_knowledge(query: str) -> str:
    """
    知识库搜索工具（模拟）
    
    为什么用字典模拟而不是真实搜索？
    - 教学目的，先理解流程
    - 真实项目中可以替换为搜索引擎 API、向量数据库等
    """
    # 一个简单的知识库（字典）
    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}' 的精确信息"

# 把工具放进字典，方便通过名称查找
# 键是工具名（LLM 输出的），值是函数对象
tools = {
    "calculator": calculator,
    "search": search_knowledge,
}

# ============================================================
# 4. LLM 调用函数
# ============================================================
"""
为什么不用 OpenAI SDK 而用 requests？
- 某些 API 网关与 OpenAI SDK 不兼容
- requests 更灵活，可以直接控制请求格式
- 原理相同，只是底层 HTTP 调用方式不同
"""
def call_llm(messages, max_tokens=300):
    """
    调用 LLM API
    
    参数:
        messages: 消息列表，格式为 [{"role": "system/user/assistant", "content": "..."}]
        max_tokens: 最大生成 token 数
    
    返回: LLM 生成的文本
    """
    # 构建 API 请求
    url = f"{BASE_URL}/chat/completions"  # 拼接完整的 API 地址
    headers = {
        "Content-Type": "application/json",  # 告诉服务器发送 JSON
        "Authorization": f"Bearer {API_KEY}"  # API 认证
    }
    body = {
        "model": MODEL,           # 使用哪个模型
        "messages": messages,     # 对话消息
        "max_tokens": max_tokens, # 最大生成长度
        "temperature": TEMPERATURE,  # 创造性（0=确定，1=随机）
    }
    
    # 发送 HTTP POST 请求
    response = requests.post(url, headers=headers, json=body, timeout=30)
    response.raise_for_status()  # 如果 HTTP 状态码不是 200，抛出异常
    data = response.json()       # 把响应解析为 JSON
    
    # 从 JSON 响应中提取 LLM 生成的文本
    # 响应结构: {"choices": [{"message": {"content": "生成的文本"}}]}
    return data["choices"][0]["message"]["content"]

# ============================================================
# 5. 定义节点（Nodes）
# ============================================================
"""
节点是 LangGraph 图中的处理单元。
每个节点是一个函数，接收 State，返回更新后的 State。

为什么每个节点都要 state = state.copy()？
- Python 字典是可变对象，直接修改会影响原始数据
- .copy() 创建浅拷贝，确保每个节点操作的是自己的副本
- 这是 LangGraph 的最佳实践
"""

def think_node(state: AgentState):
    """
    思考节点 - 让 LLM 决定下一步
    
    这是 ReAct 模式的核心：
    1. 把问题和历史发给 LLM
    2. LLM 输出 [思考] 和 [行动] 或 [回答]
    3. 解析 LLM 的输出，提取行动或答案
    """
    state = state.copy()  # 创建副本，避免修改原始状态
    state['iteration'] += 1  # 迭代次数 +1

    # 构建系统提示词（System Prompt）
    # f-string 中的 {state['iteration']} 会被替换为实际值
    system_prompt = f"""你是一个智能助手，可以使用以下工具:
1. calculator - 数学计算，参数是数学表达式如 "2+3*4"
2. search - 搜索知识，参数是搜索关键词

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

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

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

    # 构建消息列表（Messages）
    # OpenAI API 的消息格式：角色 + 内容
    messages = [{"role": "system", "content": system_prompt}]
    messages.append({"role": "user", "content": state['question']})

    # 如果有上一次的观察结果，也发给 LLM
    # 这样 LLM 可以根据工具返回的结果做下一步决定
    if state.get('observation'):
        messages.append({"role": "assistant", "content": f"[观察] {state['observation']}"})

    # 调用 LLM
    thought_text = call_llm(messages)
    
    # 保存思考历史
    state['current_thought'] = thought_text
    state['thoughts'] = state.get('thoughts', []) + [thought_text]

    # 打印调试信息
    print(f"\n{'='*50}")  # {'='*50} 生成 50 个等号
    print(f"[思考] 第 {state['iteration']} 轮:")
    print(thought_text)

    # 解析 LLM 的输出
    # LLM 输出类似:
    # [思考] 我需要计算这个数学题
    # [行动] calculator|2+3*4
    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

    # 如果没解析到行动或回答，清空 action
    state['action'] = ""
    return state

def act_node(state: AgentState):
    """
    行动节点 - 执行工具
    
    1. 从 state 中获取工具名和参数
    2. 在 tools 字典中查找对应的函数
    3. 调用函数，保存结果到 observation
    """
    state = state.copy()
    action = state.get('action', '')      # 工具名，如 "calculator"
    param = state.get('action_param', '')  # 参数，如 "123*456"

    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):
    """
    回答节点 - 生成最终答案
    
    两种情况：
    1. LLM 已经在 think_node 中给出了 [回答]，直接用
    2. 达到最大迭代次数，让 LLM 总结现有信息
    """
    state = state.copy()

    # 情况 1：LLM 已经给出了最终答案
    if state.get('final_answer'):
        print(f"\n[回答] {state['final_answer']}")
        return state

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

    state['final_answer'] = call_llm(messages, max_tokens=200)
    print(f"\n[回答] {state['final_answer']}")
    return state

# ============================================================
# 6. 路由函数（Router）
# ============================================================
"""
路由函数决定图走到哪里。
它接收当前 State，返回下一步要去的节点名。

为什么需要路由函数？
- 固定边：永远走同一条路
- 条件边：根据 State 的内容动态决定走哪条路
- 智能体需要"判断"能力，所以用条件边
"""
def route(state: AgentState):
    """
    决定下一步去哪里
    
    返回 "act"  -> 去执行工具
    返回 "answer" -> 去生成答案（结束）
    """
    # 如果 LLM 已经给出了最终答案，直接结束
    if state.get('final_answer'):
        return "answer"
    
    # 如果达到最大迭代次数，强制结束（防止无限循环）
    if state['iteration'] >= state['max_iterations']:
        return "answer"
    
    # 如果有行动要执行，去 act 节点
    if state.get('action'):
        return "act"
    
    # 默认去 answer 节点
    return "answer"

# ============================================================
# 7. 构建图（Build Graph）
# ============================================================
"""
StateGraph(State) - 创建图，指定 State 类型
add_node(name, func) - 添加节点
add_edge(from, to) - 添加固定边
add_conditional_edges(node, router, mapping) - 添加条件边
compile() - 编译图，创建可执行的应用

图的结构：
    START -> think -> [条件边] -> act -> think (循环)
                    |
                    +-> answer -> END
"""
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")  # 开始 -> 思考

# 条件边：从 think 出发，根据 route 函数的返回值决定走向
graph.add_conditional_edges(
    "think",                                    # 从 think 节点出发
    route,                                      # 用 route 函数做判断
    {"act": "act", "answer": "answer"}          # 返回值 -> 节点名的映射
)

graph.add_edge("act", "think")   # 行动完回到思考（形成循环！）
graph.add_edge("answer", END)    # 回答完结束

# 编译图
app = graph.compile()

# ============================================================
# 8. 运行（Run）
# ============================================================
print("=" * 50)
print("LangGraph ReAct 智能体")
print("=" * 50)
print(f"API: {BASE_URL}")
print(f"模型: {MODEL}")
print(f"最大迭代: {MAX_ITERATIONS}")
print("\n图结构:")
print("  START -> think -> [有行动?] -> act -> think (循环)")
print("              |")
print("              +-> [无行动/达到限制] -> answer -> END")

# 命令行参数处理
# sys.argv 是命令行参数列表，sys.argv[0] 是脚本名
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}")

    # app.invoke() 运行图
    # 传入初始 State，图会一步步处理，最终返回完整 State
    result = app.invoke({
        "question": q,
        "thoughts": [],          # 空列表，记录思考历史
        "iteration": 0,          # 从 0 开始
        "max_iterations": MAX_ITERATIONS,
        "action": "",
        "observation": "",
        "final_answer": "",
        "action_param": "",
    })

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