"""StreamingReactAgent — ReAct loop with streaming events."""

from __future__ import annotations

from typing import TYPE_CHECKING

import orjson
import structlog

from koan.agents.base import BaseAgent
from koan.agents.events import (
    AgentComplete,
    AgentError,
    AgentEvent,
    IterationStart,
    TextDelta,
    ToolCallEnd,
    ToolCallStart,
)
from koan.streaming.tools import stream_with_tools
from koan.types import Message

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from koan.agents.types import AgentConfig, RunContext
    from koan.llm.client import LLMClient
    from koan.tools.registry import ToolRegistry

log = structlog.get_logger()


class StreamingReactAgent(BaseAgent):
    """Agent that streams events during a ReAct loop.

    Instead of returning a single ``AgentResult``, yields ``AgentEvent``
    objects as the agent thinks, calls tools, and produces output.

    Usage::

        agent = StreamingReactAgent("researcher", client=llm_client, tools=registry)
        async for event in agent.run_stream("What's the weather?"):
            if isinstance(event, TextDelta):
                print(event.text, end="", flush=True)
            elif isinstance(event, AgentComplete):
                print(f"\\nDone in {event.iterations} iterations")
    """

    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):  # type: ignore[override]
        """Execute the ReAct loop, collecting all events into a final result."""
        from koan.agents.types import AgentResult

        last_event: AgentComplete | AgentError | None = None
        async for event in self.run_stream(input_text, context=context):
            if isinstance(event, (AgentComplete, AgentError)):
                last_event = event

        if isinstance(last_event, AgentComplete):
            return AgentResult(
                output=last_event.output,
                agent_name=self.name,
                iterations=last_event.iterations,
                tool_calls_made=last_event.tool_calls_made,
            )
        if isinstance(last_event, AgentError):
            return AgentResult(
                output=last_event.output,
                agent_name=self.name,
                iterations=last_event.iterations,
                tool_calls_made=last_event.tool_calls_made,
                metadata={"error": last_event.error},
            )

        return AgentResult(output="", agent_name=self.name)

    async def run_stream(
        self,
        input_text: str,
        *,
        context: RunContext | None = None,
    ) -> AsyncIterator[AgentEvent]:
        """Execute the ReAct loop, yielding events as they occur."""
        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
        accumulated_text = ""

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

            accumulated_text = ""
            has_tool_calls = False
            tool_call_list: list[dict[str, str]] = []

            async for event in stream_with_tools(self.client, llm_config, messages):
                if event.is_text:
                    accumulated_text += event.text
                    yield TextDelta(text=event.text)

                if event.has_tool_calls and self.tools:
                    has_tool_calls = True
                    for tc in event.tool_calls:
                        tool_call_list.append({"id": tc.id, "name": tc.name, "arguments": tc.arguments})

            if has_tool_calls and self.tools:
                # Add assistant message with tool_calls for provider compatibility
                raw_tool_calls = [
                    {"id": tc.get("id", ""), "name": tc["name"], "arguments": tc["arguments"]}
                    for tc in tool_call_list
                ]
                messages.append(Message(
                    role="assistant",
                    content=accumulated_text,
                    tool_calls=raw_tool_calls,
                ))

                # Execute each tool call
                for tc in tool_call_list:
                    tc_id = tc.get("id", "")
                    tc_name = tc["name"]
                    tc_args = tc["arguments"]

                    yield ToolCallStart(tool_name=tc_name, arguments=tc_args)

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

                    messages.append(Message(
                        role="tool",
                        content=result_text,
                        name=tc_name,
                        tool_call_id=tc_id,
                    ))
                    tool_calls_made += 1

                continue

            # Final response — no tool calls
            yield AgentComplete(
                output=accumulated_text,
                iterations=iterations,
                tool_calls_made=tool_calls_made,
            )
            return

        # Max iterations reached
        yield AgentError(
            error="max_iterations_reached",
            output=accumulated_text,
            iterations=iterations,
            tool_calls_made=tool_calls_made,
        )
