"""Tests for stream_with_tools()."""

from __future__ import annotations

from typing import TYPE_CHECKING
from unittest.mock import AsyncMock

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

import pytest

from koan.llm.config import LLMConfig
from koan.streaming.tools import AccumulatedToolCall, StreamToolEvent, stream_with_tools
from koan.types import Message, StreamDelta


def _make_client_with_deltas(deltas: list[StreamDelta]) -> AsyncMock:
    client = AsyncMock()

    async def _stream(config: LLMConfig, messages: list[Message]) -> AsyncIterator[StreamDelta]:
        for d in deltas:
            yield d

    client.stream = _stream
    return client


@pytest.fixture
def config() -> LLMConfig:
    return LLMConfig(provider="openai", model="gpt-4o-mini")


@pytest.fixture
def messages() -> list[Message]:
    return [Message(role="user", content="Hello")]


class TestStreamWithTools:
    async def test_text_only(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(text="Hello"),
            StreamDelta(text=" there"),
            StreamDelta(finish_reason="stop"),
        ]
        client = _make_client_with_deltas(deltas)

        events = [e async for e in stream_with_tools(client, config, messages)]

        text_events = [e for e in events if e.is_text]
        assert len(text_events) == 2
        assert text_events[0].text == "Hello"
        assert text_events[1].text == " there"

        finish_events = [e for e in events if e.finish_reason]
        assert len(finish_events) == 1
        assert finish_events[0].finish_reason == "stop"

    async def test_single_tool_call(self, config: LLMConfig, messages: list[Message]) -> None:
        """Simulate OpenAI-style tool call streaming."""
        deltas = [
            StreamDelta(tool_call_id="call_1", tool_name="get_weather"),
            StreamDelta(tool_args_delta='{"city":'),
            StreamDelta(tool_args_delta=' "NYC"}'),
            StreamDelta(finish_reason="tool_calls"),
        ]
        client = _make_client_with_deltas(deltas)

        events = [e async for e in stream_with_tools(client, config, messages)]

        tool_events = [e for e in events if e.has_tool_calls]
        assert len(tool_events) == 1
        assert len(tool_events[0].tool_calls) == 1

        tc = tool_events[0].tool_calls[0]
        assert tc.id == "call_1"
        assert tc.name == "get_weather"
        assert tc.arguments == '{"city": "NYC"}'

    async def test_multiple_tool_calls(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(tool_call_id="call_1", tool_name="search"),
            StreamDelta(tool_args_delta='{"q": "a"}'),
            StreamDelta(tool_call_id="call_2", tool_name="fetch"),
            StreamDelta(tool_args_delta='{"url": "http://x"}'),
            StreamDelta(finish_reason="tool_calls"),
        ]
        client = _make_client_with_deltas(deltas)

        events = [e async for e in stream_with_tools(client, config, messages)]

        tool_events = [e for e in events if e.has_tool_calls]
        assert len(tool_events) == 1
        assert len(tool_events[0].tool_calls) == 2

        names = {tc.name for tc in tool_events[0].tool_calls}
        assert names == {"search", "fetch"}

    async def test_text_then_tool(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(text="Let me check..."),
            StreamDelta(tool_call_id="call_1", tool_name="lookup"),
            StreamDelta(tool_args_delta='{"id": 42}'),
            StreamDelta(finish_reason="tool_calls"),
        ]
        client = _make_client_with_deltas(deltas)

        events = [e async for e in stream_with_tools(client, config, messages)]

        text_events = [e for e in events if e.is_text]
        tool_events = [e for e in events if e.has_tool_calls]
        assert len(text_events) == 1
        assert text_events[0].text == "Let me check..."
        assert len(tool_events) == 1

    async def test_anthropic_tool_use_finish(self, config: LLMConfig, messages: list[Message]) -> None:
        """Anthropic uses 'tool_use' as finish_reason instead of 'tool_calls'."""
        deltas = [
            StreamDelta(tool_call_id="toolu_1", tool_name="calc"),
            StreamDelta(tool_args_delta='{"x": 1}'),
            StreamDelta(finish_reason="tool_use"),
        ]
        client = _make_client_with_deltas(deltas)

        events = [e async for e in stream_with_tools(client, config, messages)]

        tool_events = [e for e in events if e.has_tool_calls]
        assert len(tool_events) == 1
        assert tool_events[0].tool_calls[0].name == "calc"


class TestStreamToolEvent:
    def test_is_text(self) -> None:
        assert StreamToolEvent(text="hi").is_text
        assert not StreamToolEvent(text="").is_text

    def test_has_tool_calls(self) -> None:
        tc = AccumulatedToolCall(id="1", name="f", arguments="{}")
        assert StreamToolEvent(tool_calls=[tc]).has_tool_calls
        assert not StreamToolEvent().has_tool_calls
