from typing import Iterator, Optional

from openai.types.chat import (
    ChatCompletionChunk,
)
from openai.types.chat import ChatCompletion


class OpenAIStream:
    """
    A class that reconstructs a ChatCompletion from an OpenAI stream.

    Usage: initialize this class, iterate through all the chunks, then
    call final_response to get the final ChatCompletion.

    Args:
        chunk_iter: The output of an openai chat completion streaming=True call.
    """

    def __init__(self, chunk_iter: Iterator[ChatCompletionChunk]) -> None:
        self.chunk_iter = chunk_iter
        self.first_chunk: Optional[ChatCompletionChunk] = None
        self.output_choices: list[dict] = []

    def __iter__(self) -> Iterator[ChatCompletionChunk]:
        for chunk in self.chunk_iter:
            yield chunk
            if self.first_chunk is None:
                self.first_chunk = chunk
            for chunk_choice in chunk.choices:
                for i in range(chunk_choice.index + 1 - len(self.output_choices)):
                    self.output_choices.append(
                        {
                            "index": len(self.output_choices),
                            "message": {
                                "content": None,
                                "tool_calls": None,
                            },
                            "finish_reason": None,
                            "logprobs": None,
                        }
                    )

                # choice fields
                choice = self.output_choices[chunk_choice.index]
                if chunk_choice.finish_reason:
                    choice["finish_reason"] = chunk_choice.finish_reason
                if chunk_choice.logprobs:
                    choice["logprobs"] = chunk_choice.logprobs

                # message
                if chunk_choice.delta.content:
                    if choice["message"]["content"] is None:
                        choice["message"]["content"] = ""
                    choice["message"]["content"] += chunk_choice.delta.content
                if chunk_choice.delta.role:
                    choice["message"]["role"] = chunk_choice.delta.role

                # function call
                if chunk_choice.delta.function_call:
                    raise NotImplementedError("Function calls not supported")

                # tool calls
                if chunk_choice.delta.tool_calls:
                    if choice["message"]["tool_calls"] is None:
                        choice["message"]["tool_calls"] = []
                    for tool_call_delta in chunk_choice.delta.tool_calls:
                        for i in range(
                            tool_call_delta.index
                            + 1
                            - len(choice["message"]["tool_calls"])  # type: ignore
                        ):
                            choice["message"]["tool_calls"].append(  # type: ignore
                                {
                                    "function": {"name": None, "arguments": ""},
                                }
                            )
                        tool_call = choice["message"]["tool_calls"][  # type: ignore
                            tool_call_delta.index
                        ]
                        if tool_call_delta.id is not None:
                            tool_call["id"] = tool_call_delta.id
                        if tool_call_delta.type is not None:
                            tool_call["type"] = tool_call_delta.type
                        if tool_call_delta.function is not None:
                            if tool_call_delta.function.name is not None:
                                tool_call["function"][
                                    "name"
                                ] = tool_call_delta.function.name
                            if tool_call_delta.function.arguments is not None:
                                tool_call["function"][
                                    "arguments"
                                ] += tool_call_delta.function.arguments

    def final_response(self) -> ChatCompletion:
        if self.first_chunk is None:
            raise ValueError("No chunks received")
        return ChatCompletion(
            id=self.first_chunk.id,
            choices=self.output_choices,  # type: ignore
            created=self.first_chunk.created,
            model=self.first_chunk.model,
            object="chat.completion",
        )
