"""Tests for StreamingReactAgent."""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock

from koan.agents.events import (
    AgentComplete,
    AgentError,
    IterationStart,
    TextDelta,
    ToolCallEnd,
    ToolCallStart,
)
from koan.agents.stream import StreamingReactAgent
from koan.agents.types import AgentConfig, RunContext
from koan.tools.types import ToolResult
from koan.types import StreamDelta


async def _async_iter(items):
    """Helper to create an async iterator from a list."""
    for item in items:
        yield item


def _text_stream(text: str, chunk_size: int = 3) -> list[StreamDelta]:
    """Create StreamDelta chunks for a text response."""
    chunks = []
    for i in range(0, len(text), chunk_size):
        chunks.append(StreamDelta(text=text[i : i + chunk_size]))
    chunks.append(StreamDelta(finish_reason="end_turn"))
    return chunks


def _tool_stream(tool_name: str, arguments: str, text_before: str = "") -> list[StreamDelta]:
    """Create StreamDelta chunks for a tool call response."""
    chunks = []
    if text_before:
        chunks.append(StreamDelta(text=text_before))
    chunks.append(StreamDelta(tool_call_id="tc_1", tool_name=tool_name))
    chunks.append(StreamDelta(tool_args_delta=arguments))
    chunks.append(StreamDelta(finish_reason="tool_calls"))
    return chunks


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 TestStreamingReactAgent:
    async def test_single_turn_yields_text_deltas(self) -> None:
        """Single turn yields TextDelta events then AgentComplete."""
        client = AsyncMock()
        client.stream = MagicMock(return_value=_async_iter(_text_stream("Hello world!")))

        agent = StreamingReactAgent("test", client=client)
        events = [e async for e in agent.run_stream("Hi")]

        # Should start with IterationStart
        assert isinstance(events[0], IterationStart)
        assert events[0].iteration == 1

        # TextDelta events in the middle
        text_events = [e for e in events if isinstance(e, TextDelta)]
        assert len(text_events) > 0
        full_text = "".join(e.text for e in text_events)
        assert full_text == "Hello world!"

        # Should end with AgentComplete
        assert isinstance(events[-1], AgentComplete)
        assert events[-1].output == "Hello world!"
        assert events[-1].iterations == 1
        assert events[-1].tool_calls_made == 0

    async def test_tool_call_yields_tool_events(self) -> None:
        """Tool call yields ToolCallStart, ToolCallEnd, then more text."""
        client = AsyncMock()

        # First stream: tool call
        tool_chunks = _tool_stream("search", '{"query": "weather"}')
        # Second stream: final text
        text_chunks = _text_stream("It's sunny.")

        client.stream = MagicMock(side_effect=[
            _async_iter(tool_chunks),
            _async_iter(text_chunks),
        ])

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

        agent = StreamingReactAgent("researcher", client=client, tools=registry)
        events = [e async for e in agent.run_stream("Weather?")]

        # Check event sequence
        event_types = [type(e).__name__ for e in events]

        assert "IterationStart" in event_types
        assert "ToolCallStart" in event_types
        assert "ToolCallEnd" in event_types
        assert "AgentComplete" in event_types

        # Verify tool events
        tool_starts = [e for e in events if isinstance(e, ToolCallStart)]
        assert len(tool_starts) == 1
        assert tool_starts[0].tool_name == "search"

        tool_ends = [e for e in events if isinstance(e, ToolCallEnd)]
        assert len(tool_ends) == 1
        assert tool_ends[0].tool_name == "search"

        complete = [e for e in events if isinstance(e, AgentComplete)]
        assert len(complete) == 1
        assert complete[0].tool_calls_made == 1

    async def test_max_iterations_yields_error(self) -> None:
        """Max iterations reached yields AgentError."""
        client = AsyncMock()

        # Always return tool calls — use side_effect to create fresh iterator each call
        client.stream = MagicMock(side_effect=lambda *a, **kw: _async_iter(_tool_stream("loop_tool", "{}")))

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

        config = AgentConfig(max_iterations=2)
        agent = StreamingReactAgent("looper", client=client, config=config, tools=registry)
        events = [e async for e in agent.run_stream("Loop")]

        error_events = [e for e in events if isinstance(e, AgentError)]
        assert len(error_events) == 1
        assert error_events[0].error == "max_iterations_reached"
        assert error_events[0].iterations == 2

    async def test_monotonic_timestamps(self) -> None:
        """All events have monotonically increasing timestamps."""
        client = AsyncMock()
        client.stream = MagicMock(return_value=_async_iter(_text_stream("Hello")))

        agent = StreamingReactAgent("test", client=client)
        events = [e async for e in agent.run_stream("Hi")]

        for i in range(1, len(events)):
            assert events[i].timestamp >= events[i - 1].timestamp

    async def test_event_types_distinguishable(self) -> None:
        """isinstance checks work correctly on all event types."""
        client = AsyncMock()

        tool_chunks = _tool_stream("mytool", "{}")
        text_chunks = _text_stream("Done.")

        client.stream = MagicMock(side_effect=[
            _async_iter(tool_chunks),
            _async_iter(text_chunks),
        ])

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

        agent = StreamingReactAgent("test", client=client, tools=registry)
        events = [e async for e in agent.run_stream("Go")]

        # All events should be AgentEvent instances
        from koan.agents.events import AgentEvent

        for e in events:
            assert isinstance(e, AgentEvent)

        # Check specific subtypes exist
        has_iteration = any(isinstance(e, IterationStart) for e in events)
        has_text = any(isinstance(e, TextDelta) for e in events)
        has_tool_start = any(isinstance(e, ToolCallStart) for e in events)
        has_tool_end = any(isinstance(e, ToolCallEnd) for e in events)
        has_complete = any(isinstance(e, AgentComplete) for e in events)

        assert has_iteration
        assert has_text
        assert has_tool_start
        assert has_tool_end
        assert has_complete

    async def test_run_returns_agent_result(self) -> None:
        """run() collects events and returns AgentResult."""
        client = AsyncMock()
        client.stream = MagicMock(return_value=_async_iter(_text_stream("Result text.")))

        agent = StreamingReactAgent("test", client=client)
        result = await agent.run("Hi")

        assert result.output == "Result text."
        assert result.agent_name == "test"
        assert result.iterations == 1

    async def test_tool_error_yields_error_in_tool_call_end(self) -> None:
        """Tool execution error is reported in ToolCallEnd event."""
        client = AsyncMock()

        tool_chunks = _tool_stream("broken", "{}")
        text_chunks = _text_stream("Sorry.")

        client.stream = MagicMock(side_effect=[
            _async_iter(tool_chunks),
            _async_iter(text_chunks),
        ])

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

        agent = StreamingReactAgent("test", client=client, tools=registry)
        events = [e async for e in agent.run_stream("Break")]

        tool_ends = [e for e in events if isinstance(e, ToolCallEnd)]
        assert len(tool_ends) == 1
        assert tool_ends[0].error == "boom"
        assert tool_ends[0].tool_name == "broken"

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

        tool_chunks = _tool_stream("ctx_tool", "{}")
        text_chunks = _text_stream("Done.")

        client.stream = MagicMock(side_effect=[
            _async_iter(tool_chunks),
            _async_iter(text_chunks),
        ])

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

        ctx = RunContext(user_id="u1")
        agent = StreamingReactAgent("test", client=client, tools=registry)
        _ = [e async for e in agent.run_stream("Test", context=ctx)]

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