import os
from dotenv import load_dotenv
from langchain.agents import AgentExecutor, create_react_agent
from langchain_core.tools import Tool
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate  # ← 用于自定义 prompt


# 加载环境变量
load_dotenv()
# ========================
# 1. 配置 LLM（Ollama 兼容 OpenAI API）
# ========================
DEEPSEEK_API_KEY = "123"  # 替换为实际的 API Key
llm = ChatOpenAI(
    api_key=DEEPSEEK_API_KEY,
    base_url="http://192.168.0.100:8085/v1",
    model="qwen3.5-27b-awq",
    temperature=0.3,
    max_tokens=1024,
)


# ========================
# 2. 定义工具
# ========================
def get_current_time(*args, **kwargs):
    from datetime import datetime
    current_time = datetime.now().strftime("%Y年%m月%d日 %H:%M:%S")
    return f"当前时间：{current_time}"


def calculate_math_expression(expression: str):
    try:
        # 构建安全的白名单内置函数
        safe_funcs = {}
        for func_name in ("abs", "round", "min", "max", "pow"):
            if hasattr(__builtins__, func_name):
                safe_funcs[func_name] = getattr(__builtins__, func_name)

        # 执行安全计算
        result = eval(expression, {"__builtins__": {}}, safe_funcs)
        return f"数学表达式计算结果：{expression} = {result}"
    except Exception as e:
        return f"计算失败，错误信息：{str(e)}"

tools = [
    Tool(
        name="GetCurrentTime",
        func=get_current_time,
        description="获取当前系统时间，无需任何输入参数"
    ),
    Tool(
        name="CalculateMathExpression",
        func=calculate_math_expression,
        description="计算一个数学表达式，例如 '100 * (25 + 15) / 4 - 80'，仅支持 + - * / 和括号"
    )
]

# ========================
# 3. 自定义 ReAct 提示模板（中文 + 严格格式约束）
# ========================
react_prompt_template = """你是一个智能助手，必须严格遵守以下交互协议：

## 规则
1. 如果需要使用工具，请**严格按照以下三行格式输出**，不要添加任何额外文字、标点、说明或自然语言：
Thought: 我需要使用工具来解决问题。
Action: 工具名称
Action Input: 工具输入内容

2. 如果不需要工具，直接输出：
Thought: 我可以直接回答。
Final Answer: 你的最终答案

## 可用工具
{tools}

## 合法工具名称（必须完全匹配，区分大小写）：
{tool_names}

## 用户问题
{input}

## 历史执行记录（如有）
{agent_scratchpad}
"""

# 创建 PromptTemplate，确保包含所有必要变量
prompt = PromptTemplate(
    input_variables=["tools", "tool_names", "input", "agent_scratchpad"],
    template=react_prompt_template
)

# ========================
# 4. 创建 Agent 和执行器
# ========================
agent = create_react_agent(
    llm=llm,
    tools=tools,
    prompt=prompt
)

agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,   # 自动处理解析错误并重试
    max_iterations=5,             # 防止死循环
)

# ========================
# 5. 执行任务
# ========================
if __name__ == "__main__":
    print("===== 任务1：获取当前时间 =====")
    agent_executor.invoke({"input": "请告诉我现在的时间是多少？"})

    print("\n" + "=" * 50 + "\n")

    print("===== 任务2：数学计算 =====")
    agent_executor.invoke({"input": "请帮我计算 100 * (25 + 15) / 4 - 80 的结果是多少？"})