"""Tests for stream_structured()."""

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 pydantic import BaseModel

from koan.llm.config import LLMConfig
from koan.output.parser import OutputParseError
from koan.streaming.structured import stream_structured
from koan.types import Message, StreamDelta


class Answer(BaseModel):
    text: str
    score: float


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 TestStreamStructured:
    async def test_accumulates_and_parses(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(text='{"text":'),
            StreamDelta(text=' "hello",'),
            StreamDelta(text=' "score": 0.9}'),
            StreamDelta(finish_reason="stop"),
        ]
        client = _make_client_with_deltas(deltas)

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

        # Text deltas yielded progressively
        text_events = [e for e in events if e.text_delta]
        assert len(text_events) == 3

        # Final event has the parsed result
        final = events[-1]
        assert final.done is True
        assert final.result is not None
        assert final.result.text == "hello"
        assert final.result.score == 0.9

    async def test_accumulated_text_grows(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(text='{"text": "he'),
            StreamDelta(text='llo", "score": 0.5}'),
            StreamDelta(finish_reason="stop"),
        ]
        client = _make_client_with_deltas(deltas)

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

        text_events = [e for e in events if e.text_delta]
        assert text_events[0].accumulated_text == '{"text": "he'
        assert text_events[1].accumulated_text == '{"text": "hello", "score": 0.5}'

    async def test_malformed_json_raises(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(text="not json at all"),
            StreamDelta(finish_reason="stop"),
        ]
        client = _make_client_with_deltas(deltas)

        with pytest.raises(OutputParseError):
            async for _ in stream_structured(client, config, messages, Answer):
                pass

    async def test_json_in_code_fence(self, config: LLMConfig, messages: list[Message]) -> None:
        deltas = [
            StreamDelta(text='```json\n{"text": "fenced", "score": 1.0}\n```'),
            StreamDelta(finish_reason="stop"),
        ]
        client = _make_client_with_deltas(deltas)

        events = [e async for e in stream_structured(client, config, messages, Answer)]
        final = events[-1]
        assert final.done is True
        assert final.result is not None
        assert final.result.text == "fenced"
