"""A minimal ReAct agent in plain Python.

Companion code for "How Agents Actually Reason"
https://allthingsagentic.org/how-agents-actually-reason

Usage:
    export OPENAI_API_KEY=sk-...
    uv sync
    uv run python agent.py "What's the boiling point of water in Fahrenheit?"
"""

from __future__ import annotations

import ast
import operator
import os
import re
import sys
from dataclasses import dataclass

from openai import OpenAI


SYSTEM_PROMPT = """You are a helpful assistant that solves problems step by step.

You have access to these tools:
- calc(expression: str): evaluate an arithmetic expression
- lookup(key: str): look up a fact in the knowledge base

For each turn, respond in this exact format:

Thought: <your reasoning about what to do next>
Action: <tool name>
Action Input: <input to the tool>

When you have the final answer, respond with:

Thought: <final reasoning>
Final Answer: <the answer>

Use exactly one Action per turn and then stop. Do not write the Observation
or a Final Answer in the same response as an Action; wait for the system to
give you the Observation in the next turn.
"""


KB: dict[str, str] = {
    "speed of light": "299,792,458 metres per second",
    "boiling point of water": "100 degrees Celsius at 1 atm",
    "earth radius": "approximately 6,371 km",
    "pi": "3.14159265358979",
}


_OPS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.Mod: operator.mod,
    ast.Pow: operator.pow,
    ast.USub: operator.neg,
    ast.UAdd: operator.pos,
}


def eval_arith(expr: str) -> float:
    """Safely evaluate an arithmetic expression. No names, no calls."""
    node = ast.parse(expr, mode="eval").body
    return _eval(node)


def _eval(node: ast.AST) -> float:
    if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
        return node.value
    if isinstance(node, ast.BinOp) and type(node.op) in _OPS:
        return _OPS[type(node.op)](_eval(node.left), _eval(node.right))
    if isinstance(node, ast.UnaryOp) and type(node.op) in _OPS:
        return _OPS[type(node.op)](_eval(node.operand))
    raise ValueError(f"unsupported expression node: {ast.dump(node)}")


TOOLS = {
    "calc": lambda expr: f"{eval_arith(expr):g}",
    "lookup": lambda key: KB.get(key.strip().lower(), "unknown"),
}


@dataclass
class Step:
    kind: str  # "action" or "final"
    thought: str
    tool: str | None = None
    tool_input: str | None = None
    answer: str | None = None


def parse_step(text: str) -> Step:
    # If the model hallucinated its own Observation:, slice it off so we
    # never see the fake observation or anything it generated after it.
    # Real observations come from our tool runner, not the model.
    cutoff = re.search(r"\n\s*Observation:", text)
    if cutoff:
        text = text[: cutoff.start()]

    thought = re.search(r"Thought:\s*(.+?)(?=\n[A-Z]|\Z)", text, re.S)
    action = re.search(r"Action:\s*(.+)", text)
    inp = re.search(r"Action Input:\s*(.+)", text)
    # Prefer Action over Final Answer when the model wrote both in one turn.
    # That pattern means the model skipped past the Observation it should have
    # waited for; run the tool, get the real result, and let the next turn decide.
    if action and inp:
        return Step(
            kind="action",
            thought=thought.group(1).strip() if thought else "",
            tool=action.group(1).strip(),
            tool_input=inp.group(1).strip(),
        )
    final = re.search(r"Final Answer:\s*(.+)", text, re.S)
    if final:
        return Step(
            kind="final",
            thought=thought.group(1).strip() if thought else "",
            answer=final.group(1).strip(),
        )
    return Step(
        kind="action",
        thought=thought.group(1).strip() if thought else "",
        tool=action.group(1).strip() if action else "",
        tool_input=inp.group(1).strip() if inp else "",
    )


MAX_STEPS = 6
MODEL = os.environ.get("OPENAI_MODEL", "gpt-5")


def run(goal: str) -> str:
    client = OpenAI()
    messages: list[dict] = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": goal},
    ]

    for step_num in range(1, MAX_STEPS + 1):
        resp = client.chat.completions.create(
            model=MODEL,
            messages=messages,
        )
        text = resp.choices[0].message.content or ""
        step = parse_step(text)

        print(f"\n--- step {step_num} ---")
        print(text.strip())
        print(f"parsed: {step}")

        if step.kind == "final":
            return step.answer or ""

        if step.tool not in TOOLS:
            observation = f"unknown tool: {step.tool!r}. valid tools: {list(TOOLS)}"
        else:
            try:
                observation = TOOLS[step.tool](step.tool_input or "")
            except Exception as exc:
                observation = f"error running {step.tool}: {exc}"

        print(f"Observation: {observation}")
        messages.append({"role": "assistant", "content": text})
        messages.append({"role": "user", "content": f"Observation: {observation}"})

    raise RuntimeError("agent exceeded MAX_STEPS without producing a final answer")


def main() -> None:
    if len(sys.argv) < 2:
        print('usage: uv run python agent.py "<your question>"', file=sys.stderr)
        sys.exit(2)
    goal = " ".join(sys.argv[1:])
    answer = run(goal)
    print("\n=== final answer ===")
    print(answer)


if __name__ == "__main__":
    main()