"""Tests for sse_response()."""

from __future__ import annotations

from typing import TYPE_CHECKING
from unittest.mock import AsyncMock

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

import orjson
import pytest

from koan.llm.config import LLMConfig
from koan.streaming.sse import sse_response
from koan.types import Message, StreamDelta, Usage


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 TestSSEResponse:
    async def test_text_events(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(text="Hi"),
            StreamDelta(text=" there"),
            StreamDelta(finish_reason="stop"),
        ]
        client = _make_client_with_deltas(deltas)

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

        # Should have 3 data events + 1 [DONE]
        assert len(events) == 4
        assert events[-1] == "data: [DONE]\n\n"

        # Parse first event
        assert events[0].startswith("data: ")
        payload = orjson.loads(events[0].removeprefix("data: ").strip())
        assert payload["text"] == "Hi"

    async def test_done_sentinel(self, config: LLMConfig, messages: list[Message]) -> None:
        client = _make_client_with_deltas([])
        events = [e async for e in sse_response(client, config, messages)]
        assert events == ["data: [DONE]\n\n"]

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

        events = [e async for e in sse_response(client, config, messages)]
        # Each event ends with double newline
        for event in events:
            assert event.endswith("\n\n")
            assert event.startswith("data: ")

    async def test_usage_event(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(usage=Usage(input_tokens=10, output_tokens=5)),
        ]
        client = _make_client_with_deltas(deltas)

        events = [e async for e in sse_response(client, config, messages)]
        # Usage event + DONE
        assert len(events) == 2
        payload = orjson.loads(events[0].removeprefix("data: ").strip())
        assert payload["usage"]["input_tokens"] == 10
        assert payload["usage"]["output_tokens"] == 5

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

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

        # First event has tool_call_id and tool_name
        p1 = orjson.loads(events[0].removeprefix("data: ").strip())
        assert p1["tool_call_id"] == "tc1"
        assert p1["tool_name"] == "search"

        # Second event has tool_args_delta
        p2 = orjson.loads(events[1].removeprefix("data: ").strip())
        assert p2["tool_args_delta"] == '{"q": "test"}'

    async def test_skips_empty_deltas(self, config: LLMConfig, messages: list[Message]) -> None:
        """Deltas with no meaningful content should not emit SSE events."""
        deltas = [
            StreamDelta(),  # empty delta
            StreamDelta(text="real"),
        ]
        client = _make_client_with_deltas(deltas)

        events = [e async for e in sse_response(client, config, messages)]
        # Only 1 data event + DONE (empty delta skipped)
        assert len(events) == 2
