# -*- 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")