"""
Wrapper for OpenAI types to make them compatible with Pydantic.

Pydantic doesn't support Iterable[T] well, so we use List[T] instead for tool_calls,
https://github.com/pydantic/pydantic/issues/9541

Otherwise we are using OpenAI SDK types directly.
"""

from typing import (
    Iterable,
    List,
    Literal,
    Optional,
    TypeAlias,
    Union,
)

from openai.types.chat import (
    ChatCompletionContentPartTextParam,
    ChatCompletionDeveloperMessageParam,
    ChatCompletionFunctionMessageParam,
    ChatCompletionMessageToolCallParam,
    ChatCompletionSystemMessageParam,
    ChatCompletionUserMessageParam,
)
from openai.types.chat.chat_completion_assistant_message_param import (
    Audio,
    ContentArrayOfContentPart,
    FunctionCall,
)
from typing_extensions import Required, TypedDict


class ChatCompletionAssistantMessageParamWrapper(TypedDict, total=False):
    """
    Almost exact copy of ChatCompletionAssistantMessageParam, but two changes.

    First change: List[T] instead of Iterable[T] for tool_calls. Addresses pydantic issue.
    https://github.com/pydantic/pydantic/issues/9541

    Second change: Add reasoning_content to the message. A LiteLLM property for reasoning data.
    """

    role: Required[Literal["assistant"]]
    """The role of the messages author, in this case `assistant`."""

    audio: Optional[Audio]
    """Data about a previous audio response from the model.

    [Learn more](https://platform.openai.com/docs/guides/audio).
    """

    content: Union[str, Iterable[ContentArrayOfContentPart], None]
    """The contents of the assistant message.

    Required unless `tool_calls` or `function_call` is specified.
    """

    reasoning_content: Optional[str]
    """The reasoning content of the assistant message. 
    
    A LiteLLM property for reasoning data: https://docs.litellm.ai/docs/reasoning_content
    """

    function_call: Optional[FunctionCall]
    """Deprecated and replaced by `tool_calls`.

    The name and arguments of a function that should be called, as generated by the
    model.
    """

    name: str
    """An optional name for the participant.

    Provides the model information to differentiate between participants of the same
    role.
    """

    refusal: Optional[str]
    """The refusal message by the assistant."""

    tool_calls: List[ChatCompletionMessageToolCallParam]
    """The tool calls generated by the model, such as function calls."""


class ChatCompletionToolMessageParamWrapper(TypedDict, total=False):
    content: Required[Union[str, Iterable[ChatCompletionContentPartTextParam]]]
    """The contents of the tool message."""

    role: Required[Literal["tool"]]
    """The role of the messages author, in this case `tool`."""

    tool_call_id: Required[str]
    """Tool call that this message is responding to."""

    kiln_task_tool_data: Optional[str]
    """The data for the Kiln task tool that this message is responding to.
    
    Formatted as `<project_id>:::<tool_id>:::<task_id>:::<run_id>`
    """


ChatCompletionMessageParam: TypeAlias = Union[
    ChatCompletionDeveloperMessageParam,
    ChatCompletionSystemMessageParam,
    ChatCompletionUserMessageParam,
    ChatCompletionAssistantMessageParamWrapper,
    ChatCompletionToolMessageParamWrapper,
    ChatCompletionFunctionMessageParam,
]
