# -*- coding: utf-8 -*-
"""
Plan-and-Execute Agent using DeepSeek via LangChain
- Planner: 拆解任务为子步骤
- Executor: 逐个执行子步骤（可调用工具）
- 支持多轮对话（通过手动注入历史）
------使用默认提示词
"""

import os
from datetime import datetime
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain.agents import Tool
from langchain.memory import ConversationBufferMemory
from langchain_experimental.plan_and_execute import (
    PlanAndExecute,
    load_agent_executor,
    load_chat_planner,
)

load_dotenv()

# ========================
# 1. 初始化 DeepSeek LLM   
# ========================
DEEPSEEK_API_KEY = "123"  # 替换为实际的 API Key
llm = ChatOpenAI(
    api_key=DEEPSEEK_API_KEY,
    base_url="https://xxxxx.cn/v1",  # Deepseek 的 API 基础地址
    model="deepseek-v3:671b",  # Deepseek 对话模型（可选：deepseek-chat-pro 等高级模型）
    temperature=0.7,  # 温度参数（0-1，越低越稳定）
    max_tokens=1024  # 最大生成 tokens
)

# ========================
# 2. 定义本地工具
# ========================
def calculate(expr: str) -> str:
    allowed = set("0123456789+-*/(). ")
    if not all(c in allowed for c in expr):
        return "错误：非法字符"
    try:
        return str(eval(expr, {"__builtins__": {}}, {}))
    except Exception as e:
        return f"计算失败: {e}"


def get_current_time(_) -> str:
    return datetime.now().strftime("%Y年%m月%d日 %H:%M:%S")


def answer_simple_question(query: str) -> str:
    knowledge = {
        "你是谁": "我是基于 DeepSeek 的 Plan-and-Execute 智能助手。",
        "你能做什么": "我可以制定计划并执行多步任务，比如计算、查时间、逻辑推理等。",
    }
    return knowledge.get(query.strip("？? "), "这个问题我暂时无法回答。")


tools = [
    Tool(name="Calculator", func=calculate, description="执行数学计算，输入如 '2 + 3 * 4'"),
    Tool(name="CurrentTime", func=get_current_time, description="获取当前日期和时间"),
    Tool(name="SimpleQA", func=answer_simple_question, description="回答简单预设问题"),
]

# ========================
# 3. 创建 Plan-and-Execute Agent
# ========================
planner = load_chat_planner(llm)
executor = load_agent_executor(llm, tools, verbose=True)
agent = PlanAndExecute(planner=planner, executor=executor, verbose=True)

# 添加对话记忆（PlanAndExecute 默认不支持 memory，需手动处理）
memory = ConversationBufferMemory(memory_key="history", input_key="input")


# ========================
# 4. 执行函数（带记忆）
# ========================
def run_with_memory(user_input: str) -> str:
    history = memory.load_memory_variables({}).get("history", "")
    if history:
        full_input = f"之前的对话:\n{history}\n\n现在用户问:\n{user_input}"
    else:
        full_input = user_input

    # 使用 invoke 替代 run
    response = agent.invoke({"input": full_input})["output"]

    memory.save_context({"input": user_input}, {"output": response})
    return response


# ========================
# 5. 主循环
# ========================
if __name__ == "__main__":
    print("🧠 Plan-and-Execute Agent 启动！")
    print("💡 示例：'先算 12*5，再告诉我现在几点，最后说你好'")
    print("🛑 输入 'quit' 退出。\n")

    while True:
        user_input = input("👤 你: ").strip()
        if not user_input or user_input.lower() in ["quit", "exit"]:
            print("👋 再见！")
            break

        try:
            response = run_with_memory(user_input)
            print(f"🤖 助手: {response}\n")
        except Exception as e:
            print(f"❌ 错误: {e}\n")

🧠 Plan-and-Execute Agent 启动！
💡 示例：'先算 12*5，再告诉我现在几点，最后说你好'
🛑 输入 'quit' 退出。

👤 你: 先算 12*5，再告诉我现在几点，最后说你好

> Entering new PlanAndExecute chain...
steps=[Step(value='Calculate the product of 12 multiplied by 5.'), Step(value='Determine the current time.'), Step(value='Greet the user with "你好".'), Step(value="Given the above steps taken, please respond to the user's original question.\n")]

> Entering new AgentExecutor chain...
Thought: To calculate the product of 12 multiplied by 5, I will use the Calculator tool.
Action:
```
{
  "action": "Calculator",
  "action_input": "12 * 5"
}
```

Observation: 60
Thought:Action:
```
{
  "action": "Final Answer",
  "action_input": "The product of 12 multiplied by 5 is 60."
}
```

> Finished chain.
*****

Step: Calculate the product of 12 multiplied by 5.

Response: The product of 12 multiplied by 5 is 60.

> Entering new AgentExecutor chain...
Action:
```
{
  "action": "CurrentTime",
  "action_input": ""
}
```
Observation: 2026年04月14日 18:43:00
Thought:{
  "action": "Final Answer",
  "action_input": "The current time is 2026年04月14日 18:43:00."
}

> Finished chain.
*****

Step: Determine the current time.

Response: {
  "action": "Final Answer",
  "action_input": "The current time is 2026年04月14日 18:43:00."
}

> Entering new AgentExecutor chain...
```
{
  "action": "Final Answer",
  "action_input": "你好"
}
```

> Finished chain.
*****

Step: Greet the user with "你好".

Response: 你好

> Entering new AgentExecutor chain...
{
  "action": "Final Answer",
  "action_input": "你好\n\nThe product of 12 multiplied by 5 is 60.\n\nThe current time is 2026年04月14日 18:43:00."
}

> Finished chain.
*****

Step: Given the above steps taken, please respond to the user's original question.

Response: {
  "action": "Final Answer",
  "action_input": "你好\n\nThe product of 12 multiplied by 5 is 60.\n\nThe current time is 2026年04月14日 18:43:00."
}
> Finished chain.
🤖 助手: {
  "action": "Final Answer",
  "action_input": "你好\n\nThe product of 12 multiplied by 5 is 60.\n\nThe current time is 2026年04月14日 18:43:00."
}