"""Tests for stream_text()."""

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.text import stream_text
from koan.types import Message, StreamDelta


def _make_client_with_deltas(deltas: list[StreamDelta]) -> AsyncMock:
    """Create a mock LLMClient whose stream() yields the given deltas."""
    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 TestStreamText:
    async def test_yields_text_chunks(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(text="Hello"),
            StreamDelta(text=" world"),
            StreamDelta(text="!"),
            StreamDelta(finish_reason="stop"),
        ]
        client = _make_client_with_deltas(deltas)

        chunks = [chunk async for chunk in stream_text(client, config, messages)]
        assert chunks == ["Hello", " world", "!"]

    async def test_skips_empty_text(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(text=""),
            StreamDelta(text="hi"),
            StreamDelta(text=""),
            StreamDelta(finish_reason="stop"),
        ]
        client = _make_client_with_deltas(deltas)

        chunks = [chunk async for chunk in stream_text(client, config, messages)]
        assert chunks == ["hi"]

    async def test_empty_stream(self, config: LLMConfig, messages: list[Message]) -> None:
        client = _make_client_with_deltas([])
        chunks = [chunk async for chunk in stream_text(client, config, messages)]
        assert chunks == []

    async def test_ignores_tool_deltas(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(text="before"),
            StreamDelta(tool_call_id="tc1", tool_name="search"),
            StreamDelta(tool_args_delta='{"q":'),
            StreamDelta(tool_args_delta='"hi"}'),
            StreamDelta(text="after"),
        ]
        client = _make_client_with_deltas(deltas)

        chunks = [chunk async for chunk in stream_text(client, config, messages)]
        assert chunks == ["before", "after"]
