from typing import Optional, Literal, TypedDict, Iterable, Union
from typing_extensions import Required, TypeAlias

##############################################################
############ Type Definitions Part One: Tool Calls ############
##############################################################

class Function(TypedDict, total=True):
    """The function that the model called."""

    arguments: Required[str]
    """
    The arguments to call the function with, as generated by the model in JSON
    format. Note that the model does not always generate valid JSON, and may
    hallucinate parameters not defined by your function schema. Validate the
    arguments in your code before calling your function.
    """
    name: Required[str]
    """The name of the function to call."""

class FunctionToolCall(TypedDict, total=True):
    """A call to a function tool created by the model."""

    id: Required[str]
    """The ID of the tool call."""
    function: Required[Function]
    """The function that the model called."""
    type: Required[Literal["function"]]
    """The type of the tool. Currently, only `function` is supported."""

##############################################################
############ Type Definitions Part Two: Messages #############
##############################################################

class SystemMessage(TypedDict, total=False):
    """Developer-provided instructions that the model should follow, regardless of messages sent by the user."""
    
    role: Required[Literal["system"]]
    """The role of the messages author, in this case `system`."""
    content: Required[str]
    """The contents of the system message, which is a text."""

class UserTextMessage(TypedDict, total=False):
    """Messages sent by an end user, containing prompts."""
    
    role: Required[Literal["user"]]
    """The role of the messages author, in this case `user`."""
    content: Required[str]
    """The content of the user message, which is a text."""

class AssistantTextMessage(TypedDict, total=False):
    """Messages sent by the model in response to user messages."""
    
    role: Required[Literal["assistant"]]
    """The role of the messages author, in this case `assistant`."""
    content: Optional[str]
    """The content of the assistant message, which is a text. Required unless `tool_calls` is specified."""
    tool_calls: Optional[Iterable[FunctionToolCall]]
    """The tool calls generated by the model, such as function calls."""

class ToolMessage(TypedDict, total=False):
    """Messages generated by tools."""
    
    role: Required[Literal["tool"]]
    """The role of the messages author, in this case `tool`."""
    content: Required[str]
    """The contents of the tool message."""
    tool_call_id: Required[str]
    """Tool call that this message is responding to."""

##############################################################
########### Type Definitions Part Three: TypeAlias ###########
##############################################################

ChatMessage: TypeAlias = Union[
    SystemMessage,
    UserTextMessage,
    AssistantTextMessage,
    ToolMessage,
]
