"""ReactAgent — think → act → observe loop with tool calling."""

from __future__ import annotations

from typing import TYPE_CHECKING

import orjson
import structlog

from koan.agents.base import BaseAgent
from koan.agents.types import AgentConfig, AgentResult, RunContext
from koan.types import Message

if TYPE_CHECKING:
    from koan.llm.client import LLMClient
    from koan.tools.registry import ToolRegistry

log = structlog.get_logger()


class ReactAgent(BaseAgent):
    """Agent that follows the ReAct pattern: Reason + Act.

    Sends messages to the LLM. If the LLM responds with tool calls
    (indicated by ``stop_reason``), the tools are executed and results
    fed back. Repeats until the LLM returns a final text response
    or ``max_iterations`` is reached.

    Usage::

        agent = ReactAgent("researcher", client=llm_client, tools=registry)
        result = await agent.run("What's the weather in NYC?")
        print(result.output)
    """

    def __init__(
        self,
        name: str,
        *,
        client: LLMClient,
        config: AgentConfig | None = None,
        tools: ToolRegistry | None = None,
    ) -> None:
        super().__init__(name, client=client, config=config, tools=tools)

    async def run(self, input_text: str, *, context: RunContext | None = None) -> AgentResult:
        """Execute the ReAct loop."""
        llm_config = self._build_llm_config()
        messages: list[Message] = []

        if self.config.system_prompt:
            messages.append(Message(role="system", content=self.config.system_prompt))

        messages.append(Message(role="user", content=input_text))

        iterations = 0
        tool_calls_made = 0
        last_text = ""

        for _ in range(self.config.max_iterations):
            iterations += 1

            response = await self.client.complete(llm_config, messages)
            last_text = response.text

            # Check if the model wants to call tools
            if response.stop_reason in ("tool_calls", "tool_use") and self.tools:
                tool_results = await self._handle_tool_calls(response, context)
                tool_calls_made += len(tool_results)

                # Add assistant message with tool_calls so providers can reconstruct the turn
                raw = response.raw or {}
                messages.append(Message(
                    role="assistant",
                    content=response.text or "",
                    tool_calls=raw.get("tool_calls"),
                ))

                # Add tool results with call IDs for provider compatibility
                for call_id, name, result_text in tool_results:
                    messages.append(Message(
                        role="tool",
                        content=result_text,
                        name=name,
                        tool_call_id=call_id,
                    ))

                continue

            # Final response — no tool calls
            return AgentResult(
                output=response.text,
                agent_name=self.name,
                iterations=iterations,
                tool_calls_made=tool_calls_made,
            )

        # Max iterations reached
        return AgentResult(
            output=last_text,
            agent_name=self.name,
            iterations=iterations,
            tool_calls_made=tool_calls_made,
            metadata={"error": "max_iterations_reached"},
        )

    async def _handle_tool_calls(
        self,
        response: object,
        context: RunContext | None,
    ) -> list[tuple[str, str, str]]:
        """Extract and execute tool calls from the response.

        Returns a list of (call_id, tool_name, result_text) tuples.
        """
        results: list[tuple[str, str, str]] = []
        if not self.tools:
            return results

        # Tool calls are expected in response.raw or parsed from text
        raw = getattr(response, "raw", None) or {}
        tool_calls = raw.get("tool_calls", [])

        for tc in tool_calls:
            call_id = tc.get("id", "")
            name = tc.get("name", "")
            arguments = tc.get("arguments", "{}")

            log.debug("agent.tool_call", agent=self.name, tool=name)

            try:
                result = await self.tools.execute(name, arguments, context=context)
                result_text = orjson.dumps(result.output).decode() if result.output is not None else ""
            except Exception as e:
                log.warning("agent.tool_error", agent=self.name, tool=name, error=str(e))
                result_text = f"Error: {e}"

            results.append((call_id, name, result_text))

        return results
