"""Streaming with structured output — accumulate text, validate at end."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Generic, TypeVar

from koan.output.parser import parse_response

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from pydantic import BaseModel

    from koan.llm.client import LLMClient
    from koan.llm.config import LLMConfig
    from koan.types import Message

T = TypeVar("T", bound="BaseModel")


@dataclass
class StructuredStreamEvent(Generic[T]):
    """An event from stream_structured."""

    text_delta: str = ""
    accumulated_text: str = ""
    result: T | None = field(default=None)
    done: bool = False


async def stream_structured(
    client: LLMClient,
    config: LLMConfig,
    messages: list[Message],
    model_class: type[T],
) -> AsyncIterator[StructuredStreamEvent[T]]:
    """Stream text deltas and parse the final accumulated text into a Pydantic model.

    Yields ``StructuredStreamEvent`` objects with partial text as it arrives.
    The final event has ``done=True`` and ``result`` set to the parsed model.

    Raises ``OutputParseError`` if the accumulated text cannot be parsed.

    Usage::

        async for event in stream_structured(client, config, messages, MyModel):
            if event.text_delta:
                print(event.text_delta, end="")
            if event.done:
                print(f"\\nParsed: {event.result}")
    """
    accumulated = ""

    async for delta in client.stream(config, messages):
        if delta.text:
            accumulated += delta.text
            yield StructuredStreamEvent(
                text_delta=delta.text,
                accumulated_text=accumulated,
            )

        if delta.finish_reason:
            result = parse_response(accumulated, model_class)
            yield StructuredStreamEvent(
                accumulated_text=accumulated,
                result=result,
                done=True,
            )
