import re
import json
from typing import Callable, Dict, Any, List


class Tool:
    def __init__(self, name: str, description: str, fn: Callable[[str], str]):
        self.name = name
        self.description = description
        self.fn = fn

    def __call__(self, arg: str) -> str:
        return self.fn(arg)


def calculator_tool(expression: str) -> str:
    # a restricted eval: only digits, operators, parentheses, spaces
    if not re.fullmatch(r"[0-9+\-*/(). ]+", expression):
        return "Error: invalid characters in expression."
    try:
        return str(eval(expression, {"__builtins__": {}}, {}))
    except Exception as e:
        return f"Error: {e}"


def lookup_tool(query: str) -> str:
    knowledge_base = {
        "capital of france": "Paris",
        "capital of japan": "Tokyo",
        "speed of light": "299792458 m/s",
    }
    return knowledge_base.get(query.strip().lower(), "Not found.")


class ReActAgent:
    def __init__(self, tools: List[Tool], model_fn: Callable[[str], str]):
        self.tools = {t.name: t for t in tools}
        self.model_fn = model_fn

    def run(self, task: str, max_steps: int = 5) -> str:
        trace = f"Task: {task}\n"
        for step in range(max_steps):
            response = self.model_fn(trace)
            trace += response + "\n"

            final_match = re.search(r"Final Answer:\s*(.+)", response)
            if final_match:
                return final_match.group(1).strip()

            action_match = re.search(
                r"Action:\s*(\w+)\((.*)\)", response
            )
            if action_match:
                tool_name, arg = action_match.groups()
                arg = arg.strip().strip('"').strip("'")
                if tool_name in self.tools:
                    observation = self.tools[tool_name](arg)
                else:
                    observation = f"Error: unknown tool '{tool_name}'."
                trace += f"Observation: {observation}\n"

        return "Agent did not reach a final answer within max_steps."
