"""
ReAct Agent（基于 DeepSeek Function Calling 实现）

ReAct = Reasoning + Acting：模型先"想"（决定是否调用工具、调用哪个），
再"做"（执行工具），拿到"观察结果"后继续下一轮思考，直到给出最终答案。
"""

import os
import json
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

client = OpenAI(
    api_key=os.environ.get("DEEPSEEK_API_KEY"),
    base_url="https://api.deepseek.com",
)

MODEL = "deepseek-chat"          # 换成你实际可用的模型名
MAX_STEPS = 8                    # 防止死循环的最大推理步数


# ============================================================
# 1. 工具实现层：每个工具就是一个普通的 Python 函数
# ============================================================
def get_weather(location: str) -> str:
    """模拟天气查询接口"""
    fake_db = {
        "杭州": "30℃，晴，东南风 2 级",
        "北京": "25℃，多云",
        "上海": "28℃，小雨",
    }
    return fake_db.get(location, f"{location}：暂无该城市天气数据")


def calculator(expression: str) -> str:
    """一个安全的四则运算计算器"""
    allowed = set("0123456789+-*/(). ")
    if not set(expression) <= allowed:
        return "错误：表达式包含非法字符"
    try:
        return str(eval(expression, {"__builtins__": {}}, {}))
    except Exception as e:
        return f"计算失败：{e}"


# 工具名 -> 函数 的注册表，Agent 靠它做分发
TOOL_REGISTRY = {
    "get_weather": get_weather,
    "calculator": calculator,
}


# ============================================================
# 2. 工具描述层：告诉模型有哪些工具可用（JSON Schema）
# ============================================================
TOOLS_SCHEMA = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "获取指定城市的当前天气",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {"type": "string", "description": "城市名，如杭州"},
                },
                "required": ["location"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "计算一个数学表达式，例如 '30 - 25'",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {"type": "string", "description": "要计算的数学表达式"},
                },
                "required": ["expression"],
            },
        },
    },
]


# ============================================================
# 3. 系统提示词：定义 Agent 的行为方式
# ============================================================
SYSTEM_PROMPT = """你是一个可以使用工具的智能助手。

工作要求：
1. 先思考问题需要哪些信息，判断是否需要调用工具。
2. 需要外部信息（天气、计算等）时，必须调用工具，不要凭空编造。
3. 可以连续多次调用工具；每一步只做当前最必要的事。
4. 拿到工具结果后，用简洁的中文给出最终答案。
"""


# ============================================================
# 4. 核心：ReAct 循环
# ============================================================
def run_react_agent(user_query: str, max_steps: int = MAX_STEPS, verbose: bool = True):
    """
    执行 ReAct 循环：
        思考(Thought) -> 行动(Action) -> 观察(Observation) -> 再思考 ... -> 最终答案
    返回 (最终答案, 完整消息历史)
    """
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_query},
    ]

    for step in range(1, max_steps + 1):
        if verbose:
            print(f"\n{'=' * 20} Step {step} {'=' * 20}")

        # ---------- Thought：让模型决定下一步做什么 ----------
        response = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            tools=TOOLS_SCHEMA,
            tool_choice="auto",
        )
        assistant_msg = response.choices[0].message

        # 把模型的回复写入历史（必须手动构造，content 为 None 时补成 ""）
        msg_dict = {"role": "assistant", "content": assistant_msg.content or ""}
        if assistant_msg.tool_calls:
            msg_dict["tool_calls"] = [tc.model_dump() for tc in assistant_msg.tool_calls]
        messages.append(msg_dict)

        if verbose and assistant_msg.content:
            print(f"[Thought] {assistant_msg.content}")

        # ---------- 终止条件：模型不再请求工具，直接给出答案 ----------
        if not assistant_msg.tool_calls:
            final_answer = assistant_msg.content or ""
            if verbose:
                print(f"[Final] {final_answer}")
            return final_answer, messages

        # ---------- Action + Observation：执行所有工具调用 ----------
        for tool_call in assistant_msg.tool_calls:
            name = tool_call.function.name

            # 注意：arguments 是 JSON 字符串，必须解析成 dict 再传参
            try:
                args = json.loads(tool_call.function.arguments or "{}")
            except json.JSONDecodeError:
                args = {}

            if verbose:
                print(f"[Action] {name}({json.dumps(args, ensure_ascii=False)})")

            func = TOOL_REGISTRY.get(name)
            if func is None:
                result = f"错误：未知工具 {name}"
            else:
                try:
                    result = func(**args)
                except Exception as e:
                    result = f"工具执行异常：{e}"

            result = str(result)
            if verbose:
                print(f"[Observation] {result}")

            # 工具结果必须带上 tool_call_id，一一对应
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.id,
                "content": result,
            })

    # 超过最大步数仍未收敛
    fallback = "抱歉，我在限定步骤内没能得出结论。"
    if verbose:
        print(f"[Final] {fallback}")
    return fallback, messages


# ============================================================
# 5. 运行入口
# ============================================================
if __name__ == "__main__":
    # 单轮测试
    for q in ["杭州天气怎么样？", "杭州和北京温差多少度？"]:
        print(f"\n\n########## 用户：{q} ##########")
        answer, history = run_react_agent(q)
        print(f"回答：{answer}")

    # 如果想做成多轮对话的交互式 Agent，可以这样：
    # while True:
    #     q = input("\n你：")
    #     if q in ("exit", "quit"):
    #         break
    #     answer, _ = run_react_agent(q)
    #     print("助手：", answer)