"""Tests for ReactAgent."""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock

from koan.agents.react import ReactAgent
from koan.agents.types import AgentConfig, RunContext
from koan.tools.types import ToolResult
from koan.types import LLMResponse, Usage


def _make_response(
    text: str = "Hello!",
    stop_reason: str | None = "end_turn",
    tool_calls: list[dict] | None = None,
) -> LLMResponse:
    raw = {}
    if tool_calls:
        raw["tool_calls"] = tool_calls
    return LLMResponse(
        text=text,
        model="test-model",
        provider="test",
        usage=Usage(input_tokens=10, output_tokens=5),
        stop_reason=stop_reason,
        raw=raw,
    )


def _mock_registry(**execute_kwargs) -> MagicMock:
    """Create a mock ToolRegistry that is truthy."""
    registry = MagicMock()
    registry.__bool__ = lambda self: True
    registry.execute = AsyncMock(**execute_kwargs)
    return registry


class TestReactAgent:
    async def test_single_turn_no_tools(self) -> None:
        """Agent returns text when LLM doesn't call tools."""
        client = AsyncMock()
        client.complete = AsyncMock(return_value=_make_response("The answer is 42."))

        agent = ReactAgent("test", client=client)
        result = await agent.run("What is the meaning of life?")

        assert result.output == "The answer is 42."
        assert result.agent_name == "test"
        assert result.iterations == 1
        assert result.tool_calls_made == 0
        assert not result.is_error

    async def test_tool_call_then_response(self) -> None:
        """Agent calls a tool, feeds result back, then returns final text."""
        client = AsyncMock()

        tool_response = _make_response(
            text="Let me look that up.",
            stop_reason="tool_calls",
            tool_calls=[{"name": "search", "arguments": '{"query": "weather"}'}],
        )
        final_response = _make_response("It's sunny in NYC.")
        client.complete = AsyncMock(side_effect=[tool_response, final_response])

        registry = _mock_registry(return_value=ToolResult(tool_name="search", output="sunny, 72F"))

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

        assert result.output == "It's sunny in NYC."
        assert result.iterations == 2
        assert result.tool_calls_made == 1

    async def test_max_iterations_guard(self) -> None:
        """Agent stops after max_iterations even if tools keep being called."""
        client = AsyncMock()

        tool_response = _make_response(
            text="Calling tool...",
            stop_reason="tool_calls",
            tool_calls=[{"name": "loop_tool", "arguments": "{}"}],
        )
        client.complete = AsyncMock(return_value=tool_response)

        registry = _mock_registry(return_value=ToolResult(tool_name="loop_tool", output="ok"))

        config = AgentConfig(max_iterations=3)
        agent = ReactAgent("looper", client=client, config=config, tools=registry)
        result = await agent.run("Do something")

        assert result.iterations == 3
        assert result.is_error
        assert result.metadata["error"] == "max_iterations_reached"

    async def test_tool_error_handled_gracefully(self) -> None:
        """When a tool raises an error, the agent reports it back to the LLM."""
        client = AsyncMock()

        tool_response = _make_response(
            text="Let me try.",
            stop_reason="tool_calls",
            tool_calls=[{"name": "broken", "arguments": "{}"}],
        )
        final_response = _make_response("Sorry, I couldn't do that.")
        client.complete = AsyncMock(side_effect=[tool_response, final_response])

        registry = _mock_registry(side_effect=RuntimeError("boom"))

        agent = ReactAgent("handler", client=client, tools=registry)
        result = await agent.run("Break things")

        assert result.output == "Sorry, I couldn't do that."
        assert result.iterations == 2

    async def test_system_prompt_included(self) -> None:
        """System prompt from config is prepended to messages."""
        client = AsyncMock()
        client.complete = AsyncMock(return_value=_make_response("Done."))

        config = AgentConfig(system_prompt="You are a helpful assistant.")
        agent = ReactAgent("sys", client=client, config=config)
        await agent.run("Hi")

        call_args = client.complete.call_args
        messages = call_args[0][1]
        assert messages[0].role == "system"
        assert messages[0].content == "You are a helpful assistant."
        assert messages[1].role == "user"

    async def test_context_passed_to_tools(self) -> None:
        """RunContext is forwarded to tool execution."""
        client = AsyncMock()

        tool_response = _make_response(
            text="Calling tool.",
            stop_reason="tool_calls",
            tool_calls=[{"name": "ctx_tool", "arguments": "{}"}],
        )
        final_response = _make_response("Done.")
        client.complete = AsyncMock(side_effect=[tool_response, final_response])

        registry = _mock_registry(return_value=ToolResult(tool_name="ctx_tool", output="ok"))

        ctx = RunContext(user_id="u1")
        agent = ReactAgent("ctx_agent", client=client, tools=registry)
        await agent.run("Test", context=ctx)

        registry.execute.assert_called_once_with("ctx_tool", "{}", context=ctx)
