"""End-to-end streaming tests — requires OPENAI_API_KEY."""

from __future__ import annotations

from typing import TYPE_CHECKING

import pytest

from koan.agents import (
    AgentComplete,
    AgentConfig,
    AgentStart,
    IterationStart,
    Orchestrator,
    ReactAgent,
    RuleRouter,
    StreamingReactAgent,
    ToolCallEnd,
    ToolCallStart,
)
from koan.streaming.tools import stream_with_tools
from koan.tools import ToolRegistry, tool
from koan.types import Message

from .conftest import skip_no_openai

if TYPE_CHECKING:
    from koan.agents.events import AgentEvent
    from koan.llm.client import LLMClient

pytestmark = [pytest.mark.integration, skip_no_openai]


# --- Test tools ---


@tool
def add_numbers(a: int, b: int) -> int:
    """Add two numbers together."""
    return a + b


@tool
def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is 22C and sunny."


@tool
def translate(text: str, language: str) -> str:
    """Translate text to a given language."""
    return f"[{language}] {text}"


# --- Fixtures ---


@pytest.fixture
def openai_config() -> AgentConfig:
    return AgentConfig(
        provider="openai",
        model="gpt-4o-mini",
        system_prompt="Use the provided tools to answer questions. Be concise.",
    )


@pytest.fixture
def tool_registry() -> ToolRegistry:
    registry = ToolRegistry()
    registry.add(add_numbers)
    registry.add(get_weather)
    return registry


# --- Streaming tool calls ---


class TestStreamingToolCalls:
    """Test streaming with tool call accumulation at the LLM client level."""

    async def test_stream_with_tools_calls_tool(self, llm_client: LLMClient) -> None:
        """stream_with_tools should accumulate tool call deltas and yield them."""
        from koan.llm.config import LLMConfig

        config = LLMConfig(
            provider="openai",
            model="gpt-4o-mini",
            tools=[
                {
                    "type": "function",
                    "function": {
                        "name": "get_weather",
                        "description": "Get the current weather for a city.",
                        "parameters": {
                            "type": "object",
                            "properties": {"city": {"type": "string"}},
                            "required": ["city"],
                        },
                    },
                }
            ],
        )
        messages = [
            Message(role="system", content="Use the get_weather tool to answer."),
            Message(role="user", content="What's the weather in Berlin?"),
        ]

        events = []
        async for event in stream_with_tools(llm_client, config, messages):
            events.append(event)

        # Should have at least one event with tool calls
        tool_events = [e for e in events if e.has_tool_calls]
        assert len(tool_events) >= 1, f"Expected tool call events, got: {events}"

        tc = tool_events[0].tool_calls[0]
        assert tc.name == "get_weather"
        assert tc.id  # Should have an ID
        assert "berlin" in tc.arguments.lower() or "Berlin" in tc.arguments


class TestStreamingReactAgent:
    """Test StreamingReactAgent with real tool calling."""

    async def test_streaming_agent_with_tools(
        self, llm_client: LLMClient, openai_config: AgentConfig, tool_registry: ToolRegistry
    ) -> None:
        """StreamingReactAgent should stream events including tool call/result events."""
        agent = StreamingReactAgent(
            "streamer",
            client=llm_client,
            config=openai_config,
            tools=tool_registry,
        )

        events: list[AgentEvent] = []
        async for event in agent.run_stream("What is 17 + 28? Use the add_numbers tool."):
            events.append(event)

        # Should have iteration starts
        iteration_events = [e for e in events if isinstance(e, IterationStart)]
        assert len(iteration_events) >= 1

        # Should have tool call start/end
        tool_starts = [e for e in events if isinstance(e, ToolCallStart)]
        tool_ends = [e for e in events if isinstance(e, ToolCallEnd)]
        event_types = [type(e).__name__ for e in events]
        assert len(tool_starts) >= 1, f"Expected ToolCallStart, got: {event_types}"
        assert len(tool_ends) >= 1
        assert tool_starts[0].tool_name == "add_numbers"
        assert "45" in tool_ends[0].result

        # Should complete with final text
        completes = [e for e in events if isinstance(e, AgentComplete)]
        assert len(completes) == 1
        assert "45" in completes[0].output

    async def test_streaming_agent_multi_turn_tools(
        self, llm_client: LLMClient, tool_registry: ToolRegistry
    ) -> None:
        """StreamingReactAgent should handle multi-turn tool calling."""
        config = AgentConfig(
            provider="openai",
            model="gpt-4o-mini",
            system_prompt=(
                "You have access to get_weather and add_numbers tools. "
                "Use get_weather to answer weather questions. Be concise."
            ),
        )

        agent = StreamingReactAgent(
            "weather-streamer",
            client=llm_client,
            config=config,
            tools=tool_registry,
        )

        events: list[AgentEvent] = []
        async for event in agent.run_stream("What's the weather in Tokyo?"):
            events.append(event)

        tool_starts = [e for e in events if isinstance(e, ToolCallStart)]
        assert len(tool_starts) >= 1
        assert tool_starts[0].tool_name == "get_weather"

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

    async def test_streaming_agent_run_collects_result(
        self, llm_client: LLMClient, openai_config: AgentConfig, tool_registry: ToolRegistry
    ) -> None:
        """StreamingReactAgent.run() should collect events into an AgentResult."""
        agent = StreamingReactAgent(
            "collector",
            client=llm_client,
            config=openai_config,
            tools=tool_registry,
        )

        result = await agent.run("What is 100 + 200? Use the add_numbers tool.")

        assert result.agent_name == "collector"
        assert "300" in result.output
        assert result.tool_calls_made >= 1


class TestOrchestratorStreaming:
    """Test Orchestrator with streaming sub-agents."""

    async def test_orchestrator_routes_and_streams(
        self, llm_client: LLMClient, tool_registry: ToolRegistry
    ) -> None:
        """Orchestrator.run_stream should route and stream events from sub-agent."""
        weather_agent = StreamingReactAgent(
            "weather",
            client=llm_client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="Use the get_weather tool to answer weather questions. Be concise.",
            ),
            tools=tool_registry,
        )

        math_agent = StreamingReactAgent(
            "math",
            client=llm_client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="Use the add_numbers tool to answer math questions. Be concise.",
            ),
            tools=tool_registry,
        )

        orchestrator = Orchestrator(
            agents={"weather": weather_agent, "math": math_agent},
            router=RuleRouter({"weather": "weather", "add": "math", "sum": "math", "plus": "math"}),
        )

        events: list[AgentEvent] = []
        async for event in orchestrator.run_stream("What's the weather in Paris?"):
            events.append(event)

        # Should start with AgentStart
        starts = [e for e in events if isinstance(e, AgentStart)]
        assert len(starts) >= 1
        assert starts[0].agent_name == "weather"

        # Should have tool usage
        tool_starts = [e for e in events if isinstance(e, ToolCallStart)]
        assert len(tool_starts) >= 1

        # Should complete
        completes = [e for e in events if isinstance(e, AgentComplete)]
        assert len(completes) == 1

    async def test_orchestrator_run_with_tools(
        self, llm_client: LLMClient, tool_registry: ToolRegistry
    ) -> None:
        """Orchestrator.run (non-streaming) should work with tool-calling agents."""
        calc_agent = ReactAgent(
            "calc",
            client=llm_client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="Use the add_numbers tool to answer. Be concise.",
            ),
            tools=tool_registry,
        )

        general_agent = ReactAgent(
            "general",
            client=llm_client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="Be concise.",
            ),
        )

        orchestrator = Orchestrator(
            agents={"calc": calc_agent, "general": general_agent},
            router=RuleRouter({"add": "calc", "sum": "calc", "plus": "calc"}, default="general"),
        )

        result = await orchestrator.run("What is 42 plus 58?")

        assert "100" in result.output
        assert "calc" in result.agent_name
        assert result.tool_calls_made >= 1


class TestMultiAgentStreaming:
    """Test multi-agent patterns with streaming."""

    async def test_orchestrator_routes_different_queries(
        self, llm_client: LLMClient, tool_registry: ToolRegistry
    ) -> None:
        """Orchestrator should route different queries to different agents."""
        weather_agent = ReactAgent(
            "weather",
            client=llm_client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="Use the get_weather tool. Be concise.",
            ),
            tools=tool_registry,
        )

        math_agent = ReactAgent(
            "math",
            client=llm_client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="Use the add_numbers tool. Be concise.",
            ),
            tools=tool_registry,
        )

        orchestrator = Orchestrator(
            agents={"weather": weather_agent, "math": math_agent},
            router=RuleRouter({"weather": "weather", "add": "math", "sum": "math"}),
        )

        # Weather query
        weather_result = await orchestrator.run("What's the weather in London?")
        assert weather_result.metadata.get("routed_to") == "weather"
        assert weather_result.tool_calls_made >= 1

        # Math query
        math_result = await orchestrator.run("What is the sum of 10 and 20?")
        assert math_result.metadata.get("routed_to") == "math"
        assert "30" in math_result.output
