"""
llm_agents.react_agent
=======================
ReAct (Reason + Act) agent  – Yao et al. 2022 / AIMA §2.4 rational agent.

The agent interleaves Thought / Action / Observation traces inside the
system prompt to guide the LLM toward step-by-step tool use.

Usage
-----
    from ml_ai_library.llm_agents.react_agent import ReActAgent, Tool

    def search_web(query: str) -> str:
        "Search the web and return a short summary."
        ...

    def calculator(expression: str) -> str:
        "Evaluate a Python math expression safely."
        import math
        return str(eval(expression, {"__builtins__": {}}, vars(math)))

    agent = ReActAgent(
        model="gpt-4o-mini",
        tools=[Tool.from_fn(search_web), Tool.from_fn(calculator)],
        max_steps=8,
    )
    answer = agent.run("What is the square root of the year Columbus sailed to America?")
    print(answer)
"""

from __future__ import annotations
from typing import List, Optional

from ml_ai_library.llm_agents.base_agent import BaseAgent, Tool

_REACT_SYSTEM = """\
You are a ReAct (Reason + Act) agent.  For every task follow this loop:

Thought: <reasoning about what to do next>
Action: <tool_name>(<arguments as JSON>)
Observation: <tool result, filled in by the system>

Repeat Thought/Action/Observation until you have enough information, then:
Final Answer: <your final answer to the user>

RULES:
- Always begin with Thought.
- Only call one tool per Action step.
- When you know the answer write "Final Answer:" and stop.
"""


class ReActAgent:
    """
    ReAct agent using litellm tool-calling for Action steps.

    Parameters
    ----------
    model    : litellm model string
    tools    : list of Tool objects
    max_steps: maximum Thought/Action/Observation iterations
    verbose  : print the reasoning trace if True
    """

    def __init__(
        self,
        model: str = "gpt-4o-mini",
        tools: Optional[List[Tool]] = None,
        max_steps: int = 8,
        verbose: bool = False,
        temperature: float = 0.0,
    ) -> None:
        self.verbose = verbose
        self._agent  = BaseAgent(
            model=model,
            system=_REACT_SYSTEM,
            tools=tools or [],
            max_turns=max_steps * 2,      # allow tool calls + responses
            temperature=temperature,
        )

    def run(self, task: str) -> str:
        """Execute a task and return the Final Answer."""
        reply = self._agent.chat(task)
        if self.verbose:
            self._print_trace()
        # Extract Final Answer if present
        for line in reply.splitlines():
            if line.startswith("Final Answer:"):
                return line[len("Final Answer:"):].strip()
        return reply.strip()

    def reset(self):
        self._agent.reset()

    def _print_trace(self):
        for msg in self._agent.history:
            role = msg.get("role", "")
            content = msg.get("content") or ""
            if role in ("user", "assistant") and content:
                print(f"[{role.upper()}] {content}\n")
