"""ReACT (Reasoning + Acting) agent loop for LLM-based agents."""

import json
import re
from typing import Callable


def react_loop(
    llm_chat: Callable[[list[dict]], str],
    tools: dict[str, Callable[..., str]],
    system_prompt: str,
    user_prompt: str,
    max_iterations: int = 5,
    on_thought: Callable[[str], None] | None = None,
    on_action: Callable[[str, dict], None] | None = None,
    on_observation: Callable[[str], None] | None = None,
) -> str:
    """Execute a ReACT (Reasoning + Acting) agent loop.

    The agent reasons, selects tools, observes results, and repeats until
    producing a Final Answer or exhausting max_iterations.

    Args:
        llm_chat: Callable that takes a list of messages and returns a string response.
        tools: Mapping of tool name to callable. Each tool returns a string.
        system_prompt: System instruction for the agent.
        user_prompt: Initial user request.
        max_iterations: Maximum reasoning iterations before fallback.
        on_thought: Optional callback invoked with each Thought string.
        on_action: Optional callback invoked with (action_name, action_input_dict).
        on_observation: Optional callback invoked with each Observation string.

    Returns:
        The Final Answer string, or the last raw LLM response as fallback.
    """
    tools_description = _format_tools(tools)

    augmented_system = (
        f"{system_prompt}\n\n"
        "You have access to the following tools:\n"
        f"{tools_description}\n\n"
        "Use this format:\n"
        "Thought: <your reasoning>\n"
        "Action: <tool name>\n"
        "Action Input: <JSON object with tool arguments>\n"
        "Observation: <tool result will appear here>\n"
        "... (repeat as needed)\n"
        "Final Answer: <your final response to the user>"
    )

    messages: list[dict] = [
        {"role": "system", "content": augmented_system},
        {"role": "user", "content": user_prompt},
    ]

    last_response = ""

    for _ in range(max_iterations):
        response = llm_chat(messages)
        last_response = response

        final_answer = _parse_final_answer(response)
        if final_answer is not None:
            return final_answer

        thought = _parse_section(response, "Thought")
        action = _parse_section(response, "Action")
        action_input_raw = _parse_section(response, "Action Input")

        if thought and on_thought:
            on_thought(thought)

        if not action:
            # No action and no final answer — return as-is
            return response

        action_input: dict = {}
        if action_input_raw:
            try:
                action_input = json.loads(action_input_raw)
            except json.JSONDecodeError:
                action_input = {"input": action_input_raw}

        if on_action:
            on_action(action, action_input)

        tool_fn = tools.get(action)
        if tool_fn is None:
            observation = f"Error: tool '{action}' not found. Available tools: {list(tools.keys())}"
        else:
            try:
                observation = tool_fn(**action_input)
            except Exception as exc:
                observation = f"Error executing tool '{action}': {exc}"

        if on_observation:
            on_observation(observation)

        # Append assistant turn + observation to continue the conversation
        messages.append({"role": "assistant", "content": response})
        messages.append({"role": "user", "content": f"Observation: {observation}"})

    return last_response


# ---------------------------------------------------------------------------
# Private helpers
# ---------------------------------------------------------------------------

def _format_tools(tools: dict[str, Callable[..., str]]) -> str:
    """Build a human-readable tool list from the tools dict."""
    lines = []
    for name, fn in tools.items():
        doc = (fn.__doc__ or "").strip().split("\n")[0]
        lines.append(f"- {name}: {doc}" if doc else f"- {name}")
    return "\n".join(lines)


def _parse_section(text: str, label: str) -> str | None:
    """Extract the value after a 'Label:' marker, stopping at the next label."""
    pattern = rf"^{label}:\s*(.+?)(?=\n(?:Thought|Action Input|Action|Observation|Final Answer):|$)"
    match = re.search(pattern, text, re.IGNORECASE | re.MULTILINE | re.DOTALL)
    if match:
        return match.group(1).strip()
    return None


def _parse_final_answer(text: str) -> str | None:
    """Extract Final Answer from agent response."""
    match = re.search(r"Final Answer:\s*(.+)", text, re.IGNORECASE | re.DOTALL)
    if match:
        return match.group(1).strip()
    return None
