import json
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Callable, Literal


Role = Literal["system", "user", "assistant", "tool"]


@dataclass
class Message:
    role: Role
    content: str


@dataclass
class Config:
    provider: str = "offline"
    model: str = "mini-agent-model"
    temperature: float = 0.1
    max_steps: int = 4


class OfflineModel:
    def complete(self, messages: list[Message]) -> str:
        text = "\n".join(message.content for message in messages)
        if "可用工具" in text and "calculator" in text:
            return json.dumps({"tool": "calculator", "arguments": {"expression": "9 * 9 + 1"}}, ensure_ascii=False)
        return "离线模型回答：已收到任务。"


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


class ToolBox:
    def __init__(self) -> None:
        self.tools: dict[str, Tool] = {}

    def register(self, tool: Tool) -> None:
        self.tools[tool.name] = tool

    def call(self, name: str, arguments: dict) -> str:
        if name not in self.tools:
            raise KeyError(f"未知工具：{name}")
        return self.tools[name].handler(arguments)

    def describe(self) -> str:
        return "\n".join(f"{tool.name}: {tool.description}" for tool in self.tools.values())


class Agent(ABC):
    def __init__(self, model: OfflineModel | None = None, config: Config | None = None) -> None:
        self.model = model or OfflineModel()
        self.config = config or Config()
        self.history: list[Message] = []

    @abstractmethod
    def run(self, task: str) -> str:
        raise NotImplementedError


class SimpleAgent(Agent):
    def run(self, task: str) -> str:
        self.history.append(Message("user", task))
        answer = self.model.complete(self.history)
        self.history.append(Message("assistant", answer))
        return answer


class ReActAgent(Agent):
    def __init__(self, toolbox: ToolBox, model: OfflineModel | None = None, config: Config | None = None) -> None:
        super().__init__(model, config)
        self.toolbox = toolbox

    def run(self, task: str) -> str:
        prompt = f"任务：{task}\n可用工具：\n{self.toolbox.describe()}\n只输出工具调用 JSON。"
        decision = json.loads(self.model.complete([Message("user", prompt)]))
        result = self.toolbox.call(decision["tool"], decision["arguments"])
        return f"Observation: {result}"


class ReflectionAgent(Agent):
    def run(self, task: str) -> str:
        draft = f"初稿：{task}"
        critique = "反馈：补充证据、边界和失败处理。"
        return f"{draft}\n{critique}\n修订：加入工具日志和人工接管。"


class PlanAndSolveAgent(Agent):
    def run(self, task: str) -> str:
        plan = ["理解问题", "调用必要工具", "汇总答案"]
        return "\n".join(f"{index + 1}. {step}" for index, step in enumerate(plan)) + f"\n任务：{task}"


class FunctionCallAgent(Agent):
    def __init__(self, toolbox: ToolBox, model: OfflineModel | None = None, config: Config | None = None) -> None:
        super().__init__(model, config)
        self.toolbox = toolbox

    def run(self, task: str) -> str:
        call = json.loads(self.model.complete([Message("user", f"{task}\n可用工具：\n{self.toolbox.describe()}")]))
        return self.toolbox.call(call["tool"], call["arguments"])


def calculator(arguments: dict) -> str:
    expression = str(arguments.get("expression", ""))
    return str(eval(expression, {"__builtins__": {}}, {}))


def default_toolbox() -> ToolBox:
    toolbox = ToolBox()
    toolbox.register(Tool("calculator", "计算数学表达式", calculator))
    return toolbox


if __name__ == "__main__":
    tools = default_toolbox()
    print(SimpleAgent().run("你好"))
    print(ReActAgent(tools).run("计算 9 * 9 + 1"))
    print(ReflectionAgent().run("写一段 Agent 介绍"))
    print(PlanAndSolveAgent().run("分析复杂问题"))
    print(FunctionCallAgent(tools).run("计算 9 * 9 + 1"))