import json
import re
from dataclasses import dataclass
from typing import Callable


@dataclass
class Tool:
    name: str
    description: str
    handler: Callable[[dict], str]


def calculate(args: dict) -> str:
    expression = str(args.get("expression", ""))
    if not re.fullmatch(r"[0-9+\-*/(). ]+", expression):
        return "计算失败：表达式只允许数字和基础运算符"
    try:
        return str(eval(expression, {"__builtins__": {}}, {}))
    except Exception as exc:
        return f"计算失败：{exc}"


def save_note(args: dict) -> str:
    title = str(args.get("title", "untitled"))
    body = str(args.get("body", ""))
    return f"已保存笔记《{title}》，长度 {len(body)} 字"


TOOLS = {
    "calculate": Tool("calculate", "计算基础数学表达式", calculate),
    "save_note": Tool("save_note", "保存一条学习笔记", save_note),
}


class PocketAgent:
    def __init__(self, tools: dict[str, Tool], max_steps: int = 4) -> None:
        self.tools = tools
        self.max_steps = max_steps
        self.memory: list[str] = []

    def build_prompt(self, task: str) -> str:
        tools = "\n".join(f"- {tool.name}: {tool.description}" for tool in self.tools.values())
        history = "\n".join(self.memory[-4:]) or "暂无历史"
        return (
            "你是一个只能输出 JSON 的任务助手。\n"
            f"用户目标：{task}\n"
            f"可用工具：\n{tools}\n"
            f"最近观察：{history}\n"
            '输出格式：{"type":"tool","tool":"工具名","arguments":{}} '
            '或 {"type":"finish","answer":"最终答案"}'
        )

    def call_model(self, prompt: str, task: str) -> str:
        if not any("calculate 返回" in item for item in self.memory):
            return json.dumps(
                {"type": "tool", "tool": "calculate", "arguments": {"expression": "128 * 3 + 42"}},
                ensure_ascii=False,
            )
        if not any("save_note 返回" in item for item in self.memory):
            return json.dumps(
                {
                    "type": "tool",
                    "tool": "save_note",
                    "arguments": {"title": "第一次 Agent 运行记录", "body": "Agent 先计算，再保存结果。"},
                },
                ensure_ascii=False,
            )
        return json.dumps({"type": "finish", "answer": "任务完成：计算结果是 426。"}, ensure_ascii=False)

    def parse_decision(self, raw: str) -> dict:
        decision = json.loads(raw)
        if decision.get("type") not in {"tool", "finish"}:
            raise ValueError("type 必须是 tool 或 finish")
        return decision

    def run(self, task: str) -> str:
        for step in range(1, self.max_steps + 1):
            raw = self.call_model(self.build_prompt(task), task)
            decision = self.parse_decision(raw)
            if decision["type"] == "finish":
                return decision["answer"]
            tool_name = decision.get("tool")
            if tool_name not in self.tools:
                return f"停止：模型请求了未注册工具 {tool_name}"
            observation = self.tools[tool_name].handler(decision.get("arguments", {}))
            self.memory.append(f"第 {step} 步：{tool_name} 返回 {observation}")
        return "停止：超过最大执行步数"


if __name__ == "__main__":
    agent = PocketAgent(TOOLS)
    print(agent.run("计算 128 * 3 + 42，把过程保存成学习笔记"))
    for item in agent.memory:
        print("-", item)