"""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
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,
)

if TYPE_CHECKING:
    from ..clients import ModelId, Params, Provider


@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
)

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: "Provider",
        model_id: "ModelId",
        params: "Params",
        toolkit: ToolkitT,
        format: Format[FormattableT] | None = None,
        input_messages: Sequence[Message],
        chunk_iterator: ChunkIteratorT,
    ) -> None:
        """Initialize the BaseStreamResponse.

        Args:
            provider: The provider name (e.g. "anthropic", "openai:completions").
            model_id: The model identifier that generated the response.
            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

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

        self.provider = provider
        self.model_id = model_id
        self.params = params
        self.toolkit = toolkit
        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=provider,
            model_id=model_id,
            raw_message=None,
        )

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

        self._chunk_iterator = chunk_iterator
        self._current_content: Text | Thought | ToolCall | None = None

        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:
                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:
                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"
                )
            self._current_content = ToolCall(
                id=chunk.id,
                name=chunk.name,
                args="",
            )

        elif chunk.type == "tool_call_chunk":
            if (
                self._current_content is None
                or self._current_content.type != "tool_call"
            ):
                raise RuntimeError(
                    "Received tool_call_chunk while not processing tool call."
                )
            self._current_content.args += chunk.delta

        elif chunk.type == "tool_call_end_chunk":
            if (
                self._current_content is None
                or self._current_content.type != "tool_call"
            ):
                raise RuntimeError(
                    "Received tool_call_end_chunk while not processing tool call."
                )
            if not self._current_content.args:
                self._current_content.args = "{}"
            self._content.append(self._current_content)
            self._tool_calls.append(self._current_content)
            self._current_content = None

    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
            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 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

        if not printed:
            yield "**[No Content]**"

    def structured_stream(
        self,
    ) -> Iterator[Partial[FormattableT]]:
        """Returns an iterator that yields partial structured objects as content streams.

        Returns:
            Iterator[Partial[FormatT]]: Synchronous iterator yielding partial structured objects

        This method yields Partial[FormatT] objects as the response content is streamed,
        allowing you to access partial structured data before the response is fully complete.
        Each yielded object represents the current state of the parsed structure with all
        fields optional.

        Fully iterating through this iterator will fully consume the underlying stream,
        updating the Response with all collected content.
        """
        raise NotImplementedError()


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
            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 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

        if not printed:
            yield "**[No Content]**"

    def structured_stream(
        self,
    ) -> AsyncIterator[Partial[FormattableT]]:
        """Returns an async iterator that yields partial structured objects as content streams.

        Returns:
            AsyncIterator[Partial[FormatT]]: Async iterator yielding partial structured objects

        This method yields Partial[FormatT] objects as the response content is streamed,
        allowing you to access partial structured data before the response is fully complete.
        Each yielded object represents the current state of the parsed structure with all
        fields optional.

        Fully iterating through this iterator will fully consume the underlying stream,
        updating the Response with all collected content.
        """
        raise NotImplementedError()
