"""Base class for StreamResponse and AsyncStreamResponse."""

from collections.abc import AsyncIterator, Iterator, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, TypeVar

from ..content import (
    AssistantContentChunk,
    AssistantContentPart,
    Text,
    TextChunk,
    TextEndChunk,
    TextStartChunk,
    Thought,
    ThoughtChunk,
    ThoughtEndChunk,
    ThoughtStartChunk,
    ToolCall,
    ToolCallChunk,
    ToolCallEndChunk,
    ToolCallStartChunk,
)
from ..formatting import (
    Format,
    FormattableT,
    Partial,
    is_output_parser,
)
from ..messages import AssistantMessage, Message
from ..tools import FORMAT_TOOL_NAME, ToolkitT
from ..types import Jsonable
from .finish_reason import FinishReasonChunk
from .root_response import RootResponse
from .streams import (
    AsyncStream,
    AsyncTextStream,
    AsyncThoughtStream,
    AsyncToolCallStream,
    Stream,
    TextStream,
    ThoughtStream,
    ToolCallStream,
)
from .usage import Usage, UsageDeltaChunk

if TYPE_CHECKING:
    from ..models import Params
    from ..providers import ModelId, ProviderId


StreamResponseT = TypeVar("StreamResponseT", bound="BaseStreamResponse[Any, Any, Any]")


@dataclass(kw_only=True)
class RawStreamEventChunk:
    """A chunk containing a raw stream event from the underlying provider.

    Will be accumulated on `StreamResponse.raw` for debugging purposes.
    """

    type: Literal["raw_stream_event_chunk"] = "raw_stream_event_chunk"

    raw_stream_event: Any
    """The raw stream event from the underlying provider."""


@dataclass(kw_only=True)
class RawMessageChunk:
    """A chunk containing provider-specific raw message content that will be added to the `AssistantMessage`.

    This chunk contains a provider-specific representation of a piece of content that
    will be added to the `AssistantMessage` reconstructed by the containing stream.
    This content should be a Jsonable Python object for serialization purposes.

    The intention is that this content may be passed as-is back to the provider when the
    generated `AssistantMessage` is being reused in conversation.
    """

    type: Literal["raw_message_chunk"] = "raw_message_chunk"

    raw_message: Jsonable
    """The provider-specific raw content.
    
    Should be a Jsonable object.
    """


StreamResponseChunk: TypeAlias = (
    AssistantContentChunk
    | FinishReasonChunk
    | RawStreamEventChunk
    | RawMessageChunk
    | UsageDeltaChunk
)

ChunkIterator: TypeAlias = Iterator[StreamResponseChunk]
"""Synchronous iterator yielding chunks with raw data."""

AsyncChunkIterator: TypeAlias = AsyncIterator[StreamResponseChunk]
"""Asynchronous iterator yielding chunks with raw data."""

ChunkIteratorT = TypeVar("ChunkIteratorT", bound=ChunkIterator | AsyncChunkIterator)


class BaseStreamResponse(
    RootResponse[ToolkitT, FormattableT],
    Generic[ChunkIteratorT, ToolkitT, FormattableT],
):
    """Base class underpinning StreamResponse and AsyncStreamResponse.

    Manages chunk handling logic for both.
    """

    raw_stream_events: Sequence[Any]
    """The raw stream event chunks from the LLM. Provider-specific."""

    chunks: Sequence[AssistantContentChunk]
    """All of the Mirascope chunks consumed from the stream."""

    content: Sequence[AssistantContentPart]
    """The content generated by the LLM.
    
    Content is updated in this array as it is consumed by the stream. Text content will 
    update with each text chunk (this will mutate the Text object that is returned 
    rather than creating a new one). Other content will be added once each part
    is fully streamed.
    """

    messages: list[Message]
    """The message history, including the most recent assistant message.

    The most recent assistant message will have all of the completed content that has 
    already been consumed from the stream. Text content will be included as each chunk 
    is processed; other content will be included only when its corresponding part is 
    completed (to avoid partial tool calls and the like). If no content has been 
    streamed, then the final assistant message will be present (to maintain turn order 
    expectations), but will be empty.
    """

    texts: Sequence[Text]
    """The text content in the generated response, if any.
    
    Text content updates with each text chunk as it streams. The `Text` objects are 
    mutated in place rather than creating new ones for each chunk.
    """

    tool_calls: Sequence[ToolCall]
    """The tools the LLM wants called on its behalf, if any.
    
    Tool calls are only added to this sequence once they have been fully streamed 
    to avoid partial tool calls in the response.
    """

    thoughts: Sequence[Thought]
    """The readable thoughts from the model's thinking process, if any.
    
    The thoughts may be direct output from the model thinking process, or may be a
    generated summary. (This depends on the provider; newer models tend to summarize.)

    Thoughts are added to the sequence as they are streamed. The `Thought` objects are
    mutated in place rather than creating new ones for each chunk.
    """

    consumed: bool = False
    """Whether the stream has been fully consumed.
    
    This is True after all chunks have been processed from the underlying iterator.
    When False, more content may be available by calling the stream methods.
    """

    def __init__(
        self,
        *,
        provider_id: "ProviderId",
        model_id: "ModelId",
        provider_model_name: str,
        params: "Params",
        toolkit: ToolkitT,
        format: Format[FormattableT] | None = None,
        input_messages: Sequence[Message],
        chunk_iterator: ChunkIteratorT,
        usage: Usage | None = None,
    ) -> None:
        """Initialize the BaseStreamResponse.

        Args:
            provider: The provider name (e.g. "anthropic", "openai").
            model_id: The model identifier that generated the response.
            provider_model_name: Optional provider-specific model name. May include
                provider-specific additional info (like api mode in "gpt-5:responses").
            params: The params used to generate the response (or None).
            toolkit: Toolkit containing all the tools used to generate the response.
            format: The `Format` for the expected structured output format (or None).
            input_messages: The input messages that were sent to the LLM
            usage: Token usage statistics for the response.

        The BaseStreamResponse will process the tuples to build the chunks and raw lists
        as the stream is consumed.
        """

        self.provider_id = provider_id
        self.model_id = model_id
        self.provider_model_name = provider_model_name
        self.params = params
        self.toolkit = toolkit
        self.usage = usage
        self.format = format

        # Internal-only lists which we mutate (append) during chunk processing
        self._chunks: list[AssistantContentChunk] = []
        self._content: list[AssistantContentPart] = []
        self._texts: list[Text] = []
        self._thoughts: list[Thought] = []
        self._tool_calls: list[ToolCall] = []
        self._raw_stream_events: list[Any] = []
        self._last_raw_stream_event_chunk: Any | None = None

        # Externally-facing references typed as immutable Sequences
        self.chunks = self._chunks
        self.content = self._content
        self.texts = self._texts
        self.thoughts = self._thoughts
        self.tool_calls = self._tool_calls
        self.raw_stream_events = self._raw_stream_events

        self.finish_reason = None

        self._assistant_message = AssistantMessage(
            content=self._content,
            provider_id=provider_id,
            model_id=model_id,
            provider_model_name=provider_model_name,
            raw_message=None,
        )

        self.messages = list(input_messages) + [self._assistant_message]

        self._chunk_iterator = chunk_iterator
        self._current_content: Text | Thought | None = None
        self._current_tool_calls: dict[str, ToolCall] = {}

        self._processing_format_tool: bool = False

    def _transform_format_tool_chunks(
        self, chunk: AssistantContentChunk
    ) -> AssistantContentChunk:
        if chunk.type == "tool_call_start_chunk" and chunk.name.startswith(
            FORMAT_TOOL_NAME
        ):
            self._processing_format_tool = True
            return TextStartChunk()
        if self._processing_format_tool and chunk.type == "tool_call_chunk":
            return TextChunk(delta=chunk.delta)
        if self._processing_format_tool and chunk.type == "tool_call_end_chunk":
            self._processing_format_tool = False
            return TextEndChunk()
        return chunk

    def _handle_chunk(self, chunk: AssistantContentChunk) -> AssistantContentChunk:
        if self.finish_reason:
            raise RuntimeError(
                f"Stream already finished with reason: {self.finish_reason}"
            )
        chunk = self._transform_format_tool_chunks(chunk)

        if chunk.content_type == "text":
            self._handle_text_chunk(chunk)
        elif chunk.content_type == "tool_call":
            self._handle_tool_call_chunk(chunk)
        elif chunk.content_type == "thought":
            self._handle_thought_chunk(chunk)
        else:
            raise NotImplementedError

        self._chunks.append(chunk)
        return chunk

    def _handle_text_chunk(
        self, chunk: TextStartChunk | TextChunk | TextEndChunk
    ) -> None:
        if chunk.type == "text_start_chunk":
            if self._current_content or self._current_tool_calls:
                raise RuntimeError(
                    "Received text_start_chunk while processing another chunk"
                )
            self._current_content = Text(text="")
            # Text gets included in content even when unfinished.
            self._content.append(self._current_content)
            self._texts.append(self._current_content)

        elif chunk.type == "text_chunk":
            if self._current_content is None or self._current_content.type != "text":
                raise RuntimeError("Received text_chunk while not processing text.")
            self._current_content.text += chunk.delta

        elif chunk.type == "text_end_chunk":
            if self._current_content is None or self._current_content.type != "text":
                raise RuntimeError("Received text_end_chunk while not processing text.")
            self._current_content = None

    def _handle_thought_chunk(
        self, chunk: ThoughtStartChunk | ThoughtChunk | ThoughtEndChunk
    ) -> None:
        if chunk.type == "thought_start_chunk":
            if self._current_content or self._current_tool_calls:
                raise RuntimeError(
                    "Received thought_start_chunk while processing another chunk"
                )
            self._current_content = Thought(thought="")
            # Thoughts get included even when unfinished.
            self._content.append(self._current_content)
            self._thoughts.append(self._current_content)

        elif chunk.type == "thought_chunk":
            if self._current_content is None or self._current_content.type != "thought":
                raise RuntimeError(
                    "Received thought_chunk while not processing thought."
                )
            self._current_content.thought += chunk.delta

        elif chunk.type == "thought_end_chunk":
            if self._current_content is None or self._current_content.type != "thought":
                raise RuntimeError(
                    "Received thought_end_chunk while not processing thought."
                )
            self._current_content = None

    def _handle_tool_call_chunk(
        self, chunk: ToolCallStartChunk | ToolCallChunk | ToolCallEndChunk
    ) -> None:
        if chunk.type == "tool_call_start_chunk":
            if self._current_content:
                raise RuntimeError(
                    "Received tool_call_start_chunk while processing another chunk"
                )
            if chunk.id in self._current_tool_calls:
                raise RuntimeError("Got tool_call_start_chunk with conflicting id")
            # Create a new tool call and track it by ID
            # Multiple tool calls can be in progress simultaneously (interleaved)
            tool_call = ToolCall(
                id=chunk.id,
                name=chunk.name,
                args="",
            )
            self._current_tool_calls[chunk.id] = tool_call

        elif chunk.type == "tool_call_chunk":
            # Look up the tool call by ID
            tool_call = self._current_tool_calls.get(chunk.id)
            if tool_call is None:
                raise RuntimeError(
                    f"Received tool_call_chunk for unknown tool call ID: {chunk.id}"
                )
            tool_call.args += chunk.delta

        elif chunk.type == "tool_call_end_chunk":
            # Finalize the tool call
            tool_call = self._current_tool_calls.get(chunk.id)
            if tool_call is None:
                raise RuntimeError(
                    f"Received tool_call_end_chunk for unknown tool call ID: {chunk.id}"
                )
            if not tool_call.args:
                tool_call.args = "{}"
            self._content.append(tool_call)
            self._tool_calls.append(tool_call)
            del self._current_tool_calls[chunk.id]

    def _pretty_chunk(self, chunk: AssistantContentChunk, spacer: str) -> str:
        match chunk.type:
            case "text_start_chunk":
                return spacer
            case "text_chunk":
                return chunk.delta
            case "tool_call_start_chunk":
                return spacer + f"**ToolCall ({chunk.name}):** "
            case "tool_call_chunk":
                return chunk.delta
            case "thought_start_chunk":
                return spacer + "**Thinking:**\n  "
            case "thought_chunk":
                return chunk.delta.replace("\n", "\n  ")  # Indent every line
            case _:
                return ""


class BaseSyncStreamResponse(BaseStreamResponse[ChunkIterator, ToolkitT, FormattableT]):
    """A base class for synchronous Stream Responses."""

    def streams(self) -> Iterator[Stream]:
        """Returns an iterator that yields streams for each content part in the response.

        Returns:
            Iterator[Stream]: Synchronous iterator yielding Stream objects

        Each content part in the response will correspond to one stream, which will yield
        chunks of content as they come in from the underlying LLM.

        Fully iterating through this iterator will fully consume the underlying stream,
        updating the Response with all collected content.

        As content is consumed, it is cached on the StreamResponse. If a new iterator
        is constructed via calling `streams()`, it will start by replaying the cached
        content from the response, and (if there is still more content to consume from
        the LLM), it will proceed to consume it once it has iterated through all the
        cached chunks.
        """
        chunk_iter = self.chunk_stream()

        for start_chunk in chunk_iter:
            # At the start of this loop, we always expect to find a start chunk. Then,
            # before proceeding, we will collect from the stream we create (in case the
            # user did not exhaust it), which ensures we will be expecting a start chunk
            # again on the next iteration
            match start_chunk.type:
                case "text_start_chunk":

                    def text_stream_iterator() -> Iterator[TextChunk]:
                        for chunk in chunk_iter:
                            if chunk.type == "text_chunk":
                                yield chunk
                            else:
                                return  # Stream finished

                    stream = TextStream(chunk_iterator=text_stream_iterator())
                    yield stream

                case "thought_start_chunk":

                    def thought_stream_iterator() -> Iterator[ThoughtChunk]:
                        for chunk in chunk_iter:
                            if chunk.type == "thought_chunk":
                                yield chunk
                            else:
                                return  # Stream finished

                    stream = ThoughtStream(chunk_iterator=thought_stream_iterator())
                    yield stream

                case "tool_call_start_chunk":
                    tool_id = start_chunk.id
                    tool_name = start_chunk.name

                    def tool_call_stream_iterator() -> Iterator[ToolCallChunk]:
                        for chunk in chunk_iter:
                            if chunk.type == "tool_call_chunk":
                                yield chunk
                            else:
                                return  # Stream finished

                    stream = ToolCallStream(
                        tool_id=tool_id,
                        tool_name=tool_name,
                        chunk_iterator=tool_call_stream_iterator(),
                    )
                    yield stream

                case _:  # pragma: no cover
                    raise RuntimeError(f"Expected start chunk, got: {start_chunk.type}")

            # Before continuing to the next stream, make sure the last stream is consumed
            # (If the user did not do so when we yielded it)
            stream.collect()

    def chunk_stream(
        self,
    ) -> Iterator[AssistantContentChunk]:
        """Returns an iterator that yields content chunks as they are received.

        Returns:
            Iterator[AssistantContentChunk]: Synchronous iterator yielding chunks

        This provides access to the Mirascope chunk data including start, delta, and end chunks
        for each content type (text, thought, tool_call). Unlike the streams() method
        that groups chunks by content part, this yields individual chunks as they arrive.

        Fully iterating through this iterator will fully consume the underlying stream,
        updating the Response with all collected content.

        As chunks are consumed, they are cached on the StreamResponse. If a new iterator
        is constructed via calling `chunk_stream()`, it will start by replaying the cached
        chunks from the response, and (if there is still more content to consume from
        the LLM), it will proceed to consume it once it has iterated through all the
        cached chunks.
        """
        for chunk in self.chunks:
            yield chunk

        if self.consumed:
            return

        for chunk in self._chunk_iterator:
            if chunk.type == "raw_stream_event_chunk":
                self._raw_stream_events.append(chunk.raw_stream_event)
            elif chunk.type == "raw_message_chunk":
                self._assistant_message.raw_message = chunk.raw_message
            elif chunk.type == "finish_reason_chunk":
                self.finish_reason = chunk.finish_reason
            elif chunk.type == "usage_delta_chunk":
                if self.usage is None:
                    self.usage = Usage()
                self.usage.input_tokens += chunk.input_tokens
                self.usage.output_tokens += chunk.output_tokens
                self.usage.cache_read_tokens += chunk.cache_read_tokens
                self.usage.cache_write_tokens += chunk.cache_write_tokens
                self.usage.reasoning_tokens += chunk.reasoning_tokens
                if chunk.provider_tool_usage:
                    self.usage.provider_tool_usage = chunk.provider_tool_usage
            else:
                yield self._handle_chunk(chunk)

        self.consumed = True

    def finish(self) -> None:
        """Finish streaming all of this response's content."""
        for _ in self.chunk_stream():
            pass

    def text_stream(self, sep: str = "\n") -> Iterator[str]:
        """Stream only the text content from the response.

        Args:
            sep: Separator to yield between text parts. Defaults to newline.

        Returns:
            Iterator[str]: Iterator yielding text delta strings

        Yields text deltas as they arrive, ignoring thoughts, tool calls, and other
        content types. Ideal for displaying text to users in real-time.

        If you concatenate the strings from `.text_stream()`, it will be equivalent to
        calling `.text(sep=sep)` on the fully consumed response.
        """
        for stream in self.streams():
            if stream.content_type == "text":
                yield from stream
                yield sep

    def pretty_stream(self) -> Iterator[str]:
        """Stream a readable representation of the stream_response as text.

        Returns:
            Iterator[str]: Iterator yielding string chunks depicting the content

        Iterating through the pretty stream will populate the stream response by consuming
        the underlying iterator (if it hasn't been consumed already). Calling `.pretty_stream()`
        will always return a fresh iterator that begins from the start of the stream.

        If you concatenate the text from `.pretty_stream()`, it will be equivalent to the
        text generated by calling `.pretty()` (assuming the response was fully consumed
        at the time when you call `.pretty()`).
        """
        printed = False

        for chunk in self.chunk_stream():
            pretty = self._pretty_chunk(chunk, "\n\n" if printed else "")
            if pretty != "":
                printed = True
            yield pretty

    def structured_stream(
        self,
    ) -> Iterator[Partial[FormattableT]]:
        """Drive the stream forward, yielding partial formatted outputs as they arrive.

        This method consumes the underlying stream chunk by chunk, yielding parsed
        partial outputs each time new content arrives. Each yielded value is a
        Partial[FormattableT] with optional fields that may be None until fully received.

        Example:
            >>> response = recommend_book.stream("fantasy")
            >>> for partial in response.structured_stream():
            >>>     print(f"Title so far: {partial.title}")
            >>> book = response.parse()  # Get final complete result

        Fully iterating through this iterator will fully consume the underlying stream,
        updating the Response with all collected content.

        Yields:
            Partial[FormattableT]: Partial objects with fields populated as they arrive.
                Fields not yet received will be None.

        Raises:
            ValueError: If format parameter not set.
            NotImplementedError: If format uses OutputParser (not supported).
        """
        if self.format is None:
            raise ValueError("structured_stream() requires format parameter")

        if is_output_parser(self.format.formattable):
            raise NotImplementedError(
                "structured_stream() not supported for OutputParser. "
                "Use BaseModel or primitive types."
            )

        for chunk in self.chunk_stream():
            if chunk.type == "text_chunk":
                partial = self.parse(partial=True)
                if partial:
                    yield partial


class BaseAsyncStreamResponse(
    BaseStreamResponse[AsyncChunkIterator, ToolkitT, FormattableT]
):
    """A base class for asynchronous Stream Responses."""

    async def streams(self) -> AsyncIterator[AsyncStream]:
        """Returns an async iterator that yields streams for each content part in the response.

        Returns:
            AsyncIterator[AsyncStream]: Async iterator yielding AsyncStream objects

        Each content part in the response will correspond to one stream, which will yield
        chunks of content as they come in from the underlying LLM.

        Fully iterating through this iterator will fully consume the underlying stream,
        updating the Response with all collected content.

        As content is consumed, it is cached on the AsyncStreamResponse. If a new iterator
        is constructed via calling `streams()`, it will start by replaying the cached
        content from the response, and (if there is still more content to consume from
        the LLM), it will proceed to consume it once it has iterated through all the
        cached chunks.
        """
        chunk_iter = self.chunk_stream()

        async for start_chunk in chunk_iter:
            # At the start of this loop, we always expect to find a start chunk. Then,
            # before proceeding, we will collect from the stream we create (in case the
            # user did not exhaust it), which ensures we will be expecting a start chunk
            # again on the next iteration
            match start_chunk.type:
                case "text_start_chunk":

                    async def text_stream_iterator() -> AsyncIterator[TextChunk]:
                        async for chunk in chunk_iter:
                            if chunk.type == "text_chunk":
                                yield chunk
                            else:
                                return  # Stream finished

                    stream = AsyncTextStream(chunk_iterator=text_stream_iterator())
                    yield stream

                case "thought_start_chunk":

                    async def thought_stream_iterator() -> AsyncIterator[ThoughtChunk]:
                        async for chunk in chunk_iter:
                            if chunk.type == "thought_chunk":
                                yield chunk
                            else:
                                return  # Stream finished

                    stream = AsyncThoughtStream(
                        chunk_iterator=thought_stream_iterator()
                    )
                    yield stream

                case "tool_call_start_chunk":
                    tool_id = start_chunk.id
                    tool_name = start_chunk.name

                    async def tool_call_stream_iterator() -> AsyncIterator[
                        ToolCallChunk
                    ]:
                        async for chunk in chunk_iter:
                            if chunk.type == "tool_call_chunk":
                                yield chunk
                            else:
                                return  # Stream finished

                    stream = AsyncToolCallStream(
                        tool_id=tool_id,
                        tool_name=tool_name,
                        chunk_iterator=tool_call_stream_iterator(),
                    )
                    yield stream

                case _:  # pragma: no cover
                    raise RuntimeError(f"Expected start chunk, got: {start_chunk.type}")

            # Before continuing to the next stream, make sure the last stream is consumed
            # (If the user did not do so when we yielded it)
            await stream.collect()

    async def chunk_stream(
        self,
    ) -> AsyncIterator[AssistantContentChunk]:
        """Returns an async iterator that yields content chunks as they are received.

        Returns:
            AsyncIterator[AssistantContentChunk]: Async iterator yielding chunks

        This provides access to the Mirascope chunk data including start, delta, and end chunks
        for each content type (text, thinking, tool_call). Unlike the streams() method
        that groups chunks by content part, this yields individual chunks as they arrive.

        Fully iterating through this iterator will fully consume the underlying stream,
        updating the Response with all collected content.

        As chunks are consumed, they are cached on the AsyncStreamResponse. If a new iterator
        is constructed via calling `chunk_stream()`, it will start by replaying the cached
        chunks from the response, and (if there is still more content to consume from
        the LLM), it will proceed to consume it once it has iterated through all the
        cached chunks.
        """

        for chunk in self.chunks:
            yield chunk

        if self.consumed:
            return

        async for chunk in self._chunk_iterator:
            if chunk.type == "raw_stream_event_chunk":
                self._raw_stream_events.append(chunk.raw_stream_event)
            elif chunk.type == "raw_message_chunk":
                self._assistant_message.raw_message = chunk.raw_message
            elif chunk.type == "finish_reason_chunk":
                self.finish_reason = chunk.finish_reason
            elif chunk.type == "usage_delta_chunk":
                if self.usage is None:
                    self.usage = Usage()
                self.usage.input_tokens += chunk.input_tokens
                self.usage.output_tokens += chunk.output_tokens
                self.usage.cache_read_tokens += chunk.cache_read_tokens
                self.usage.cache_write_tokens += chunk.cache_write_tokens
                self.usage.reasoning_tokens += chunk.reasoning_tokens
                if chunk.provider_tool_usage:
                    self.usage.provider_tool_usage = chunk.provider_tool_usage
            else:
                yield self._handle_chunk(chunk)

        self.consumed = True

    async def finish(self) -> None:
        """Finish streaming all of this response's content."""
        async for _ in self.chunk_stream():
            pass

    async def text_stream(self, sep: str = "\n") -> AsyncIterator[str]:
        """Stream only the text content from the response.

        Args:
            sep: Separator to yield between text parts. Defaults to newline.

        Returns:
            AsyncIterator[str]: Async iterator yielding text delta strings

        Yields text deltas as they arrive, ignoring thoughts, tool calls, and other
        content types. Ideal for displaying text to users in real-time.

        If you concatenate the strings from `.text_stream()`, it will be equivalent to
        calling `.text(sep=sep)` on the fully consumed response.
        """
        async for stream in self.streams():
            if stream.content_type == "text":
                async for delta in stream:
                    yield delta
                yield sep

    async def pretty_stream(self) -> AsyncIterator[str]:
        """Stream a readable representation of the stream_response as text.

        Returns:
            AsyncIterator[str]: Async iterator yielding string chunks depicting the content

        Iterating through the pretty stream will populate the stream response by consuming
        the underlying iterator (if it hasn't been consumed already). Calling `.pretty_stream()`
        will always return a fresh iterator that begins from the start of the stream.

        If you concatenate the text from `.pretty_stream()`, it will be equivalent to the
        text generated by calling `.pretty()` (assuming the response was fully consumed
        at the time when you call `.pretty()`).
        """
        printed = False

        async for chunk in self.chunk_stream():
            pretty = self._pretty_chunk(chunk, "\n\n" if printed else "")
            if pretty != "":
                printed = True
            yield pretty

    async def structured_stream(
        self,
    ) -> AsyncIterator[Partial[FormattableT]]:
        """Drive the stream forward, yielding partial formatted outputs as they arrive.

        This method consumes the underlying stream chunk by chunk, yielding parsed
        partial outputs each time new content arrives. Each yielded value is a
        Partial[FormattableT] with optional fields that may be None until fully received.

        Example:
            >>> response = await recommend_book.stream("fantasy")
            >>> async for partial in response.structured_stream():
            >>>     print(f"Title so far: {partial.title}")
            >>> book = response.parse()  # Get final complete result

        Fully iterating through this iterator will fully consume the underlying stream,
        updating the Response with all collected content.

        Yields:
            Partial[FormattableT]: Partial objects with fields populated as they arrive.
                Fields not yet received will be None.

        Raises:
            ValueError: If format parameter not set.
            NotImplementedError: If format uses OutputParser (not supported).
        """
        if self.format is None:
            raise ValueError("structured_stream() requires format parameter")

        if is_output_parser(self.format.formattable):
            raise NotImplementedError(
                "structured_stream() not supported for OutputParser. "
                "Use BaseModel or primitive types."
            )

        async for chunk in self.chunk_stream():
            if chunk.type == "text_chunk":
                partial = self.parse(partial=True)
                if partial:
                    yield partial
