from typing import Annotated, Literal

from pydantic import BaseModel


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


class ChatCompletionMessageFunctionToolCall(BaseModel):
    """
    Function tool call from OpenAI's SDK.

    Example usage:

    {
        "id": "call_12345xyz",
        "type": "function",
        "function": {
            "name": "get_weather",
            "arguments": "{\"location\":\"Paris, France\"}"
        }
    }
    """

    id: str
    """The ID of the tool call."""

    function: Function
    """The function that the model called."""

    type: Literal["function"] = "function"
    """The type of the tool. Always `function`."""


class Custom(BaseModel):
    """Custom tool call parameters."""

    input: str
    """The input for the custom tool call generated by the model."""

    name: str
    """The name of the custom tool to call."""


class ChatCompletionMessageCustomToolCall(BaseModel):
    """
    Custom tool call from OpenAI's updated SDK.

    Example usage:

    {
        "id": "call_12345xyz",
        "type": "custom",
        "custom": {
            "name": "custom_tool",
            "input": "some input data"
        }
    }
    """

    id: str
    """The ID of the tool call."""

    custom: Custom
    """The custom tool that the model called."""

    type: Literal["custom"] = "custom"
    """The type of the tool. Always `custom`."""


class ChatCompetionMessageToolCallResult(BaseModel):
    """
    Example usage

    {
        "role": "tool",
        "tool_call_id": tool_call.id,
        "content": str(result)
    }
    """

    role: Literal["tool"] = "tool"
    """The role of the message. Always `tool`."""

    tool_call_id: str
    """The ID of the tool call that this message is responding to."""

    content: str
    """Result of the tool call in string format."""


# Type alias matching OpenAI's latest SDK structure
ChatCompletionMessageToolCall = Annotated[
    ChatCompletionMessageFunctionToolCall | ChatCompletionMessageCustomToolCall,
    "discriminator='type'",
]


class ChatCompletionMessage(BaseModel):
    content: str | None = None
    """The contents of the message."""

    role: Literal["assistant"] = "assistant"
    """The role of the author of this message."""

    tool_calls: list[ChatCompletionMessageToolCall] | None = None
    """The tool calls generated by the model, such as function calls or custom calls."""
