from typing import Literal, Optional

from pydantic import BaseModel, Field


class FunctionCall(BaseModel):
    name: str = Field(..., description="The name of the function to call.")
    arguments: str = Field(
        ...,
        description=(
            "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."
        ),
    )


class ToolCall(BaseModel):
    id: str = Field(..., description="The ID of the tool call.")
    type: str = Field(
        Literal["function"],
        description="The type of the tool. Currently, only `function` is supported.",
    )
    function: FunctionCall = Field(
        ..., description="The function that the model called."
    )


class ContentPart(BaseModel):
    type: str = Field(..., description="The type of the content part.")


class TextContentPart(BaseModel):
    type: str = Field("text", description="The type of the content part.")
    text: str = Field(..., description="The text content.")


class ImageUrl(BaseModel):
    url: str = Field(
        ..., description="Either a URL of the image or the base64 encoded image data."
    )
    detail: Optional[str] = Field(
        "auto",
        description="Specifies the detail level of the image. Learn more in the Vision guide.",
    )


class ImageContentPart(BaseModel):
    type: str = Field("image", description="The type of the content part.")
    image_url: ImageUrl


class Message(BaseModel):
    role: str = Field(..., description="The role of the messages author.")
    name: Optional[str] = Field(
        None,
        description="An optional name for the participant. Provides the model information to differentiate between participants of the same role.",
    )


class SystemMessage(Message):
    role: str = Field(
        "system",
        description="The role of the messages author. Must be `system`.",
    )
    content: str = Field(..., description="The contents of the system message.")


class UserMessage(Message):
    role: str = Field(
        "user",
        description="The role of the messages author. Must be `user`.",
    )
    content: str | list[TextContentPart | ImageContentPart] = Field(
        ...,
        description="The contents of the user message.",
    )


class AssistantMessage(Message):
    role: str = Field(
        "assistant",
        description="The role of the messages author. Must be `assistant`.",
    )
    content: Optional[str] = Field(
        ...,
        description="The contents of the assistant message. Required unless tool_calls or function_call is specified.",
    )
    tool_calls: Optional[list[ToolCall]] = Field(
        None,
        description="The tool calls generated by the model, such as function calls.",
    )


class ToolMessage(Message):
    role: str = Field(
        "tool",
        description="The role of the messages author, in this case tool.",
    )
    content: str = Field(
        ...,
        description="The contents of the tool message.",
    )
    tool_call_id: str = Field(
        ...,
        description="Tool call that this message is responding to.",
    )


class Messages(BaseModel):
    messages: list[SystemMessage | UserMessage | AssistantMessage | ToolMessage] = (
        Field(
            ...,
            description="A list of messages comprising the conversation so far.",
        )
    )
