import time
import uuid
from dataclasses import asdict, dataclass, field, fields
from typing import Any, Literal

from mlflow.types.schema import AnyType, Array, ColSpec, DataType, Map, Object, Property, Schema

# TODO: Switch to pydantic in a future version of MLflow.
#       For now, to prevent adding pydantic as a core dependency,
#       we use dataclasses instead.
#
#       Unfortunately, validation for generic types is not that
#       straightforward. For example, `isinstance(thing, List[T])``
#       is not supported, so the code here is a little ugly.


JSON_SCHEMA_TYPES = ["string", "number", "integer", "object", "array", "boolean", "null"]


class _BaseDataclass:
    def _validate_field(self, key, val_type, required):
        value = getattr(self, key, None)
        if required and value is None:
            raise ValueError(f"`{key}` is required")
        if value is not None and not isinstance(value, val_type):
            raise ValueError(
                f"`{key}` must be of type {val_type.__name__}, got {type(value).__name__}"
            )

    def _validate_literal(self, key, allowed_values, required):
        value = getattr(self, key, None)
        if required and value is None:
            raise ValueError(f"`{key}` is required")
        if value is not None and value not in allowed_values:
            raise ValueError(f"`{key}` must be one of {allowed_values}, got {value}")

    def _validate_list(self, key, val_type, required):
        values = getattr(self, key, None)
        if required and values is None:
            raise ValueError(f"`{key}` is required")

        if values is not None:
            if isinstance(values, list) and not all(isinstance(v, val_type) for v in values):
                raise ValueError(f"All items in `{key}` must be of type {val_type.__name__}")
            elif not isinstance(values, list):
                raise ValueError(f"`{key}` must be a list, got {type(values).__name__}")

    def _convert_dataclass(self, key: str, cls: "_BaseDataclass", required=True):
        value = getattr(self, key)
        if value is None:
            if required:
                raise ValueError(f"`{key}` is required")
            return

        if isinstance(value, cls):
            return

        if not isinstance(value, dict):
            raise ValueError(
                f"Expected `{key}` to be either an instance of `{cls.__name__}` or "
                f"a dict matching the schema. Received `{type(value).__name__}`"
            )

        try:
            setattr(self, key, cls.from_dict(value))
        except TypeError as e:
            raise ValueError(f"Error when coercing {value} to {cls.__name__}: {e}")

    def _convert_dataclass_list(self, key: str, cls: "_BaseDataclass", required=True):
        values = getattr(self, key)
        if values is None:
            if required:
                raise ValueError(f"`{key}` is required")
            return
        if not isinstance(values, list):
            raise ValueError(f"`{key}` must be a list")

        if len(values) > 0:
            # if the items are all dicts, try to convert them to the desired class
            if all(isinstance(v, dict) for v in values):
                try:
                    setattr(self, key, [cls.from_dict(v) for v in values])
                except TypeError as e:
                    raise ValueError(f"Error when coercing {values} to {cls.__name__}: {e}")
            elif any(not isinstance(v, cls) for v in values):
                raise ValueError(
                    f"Items in `{key}` must all have the same type: {cls.__name__} or dict"
                )

    def _convert_dataclass_map(self, key, cls, required=True):
        mapping = getattr(self, key)
        if mapping is None:
            if required:
                raise ValueError(f"`{key}` is required")
            return

        if not isinstance(mapping, dict):
            raise ValueError(f"`{key}` must be a dict")

        # create a new map to avoid mutating the original
        new_mapping = {}
        for k, v in mapping.items():
            if isinstance(v, cls):
                new_mapping[k] = v
            elif isinstance(v, dict):
                try:
                    new_mapping[k] = cls.from_dict(v)
                except TypeError as e:
                    raise ValueError(f"Error when coercing {v} to {cls.__name__}: {e}")
            else:
                raise ValueError(
                    f"Items in `{key}` must be either an instance of `{cls.__name__}` "
                    f"or a dict matching the schema. Received `{type(v).__name__}`"
                )
        setattr(self, key, new_mapping)

    def to_dict(self):
        return asdict(self, dict_factory=lambda obj: {k: v for (k, v) in obj if v is not None})

    @classmethod
    def from_dict(cls, data):
        """
        Create an instance of the class from a dict, ignoring any undefined fields.
        This is useful when the dict contains extra fields, causing cls(**data) to fail.
        """
        field_names = [field.name for field in fields(cls)]
        filtered_data = {k: v for k, v in data.items() if k in field_names}
        return cls(**filtered_data)


@dataclass
class FunctionToolCallArguments(_BaseDataclass):
    """
    The arguments of a function tool call made by the model.

    Args:
        arguments (str): A JSON string of arguments that should be passed to the tool.
        name (str): The name of the tool that is being called.
    """

    name: str
    arguments: str

    def __post_init__(self):
        self._validate_field("name", str, True)
        self._validate_field("arguments", str, True)

    def to_tool_call(self, id=None):
        if id is None:
            id = str(uuid.uuid4())
        return ToolCall(id=id, function=self)


@dataclass
class ToolCall(_BaseDataclass):
    """
    A tool call made by the model.

    Args:
        function (:py:class:`FunctionToolCallArguments`): The arguments of the function tool call.
        id (str): The ID of the tool call. Defaults to a random UUID.
        type (str): The type of the object. Defaults to "function".
    """

    function: FunctionToolCallArguments
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    type: str = "function"

    def __post_init__(self):
        self._validate_field("id", str, True)
        self._convert_dataclass("function", FunctionToolCallArguments, True)
        self._validate_field("type", str, True)


@dataclass
class ChatMessage(_BaseDataclass):
    """
    A message in a chat request or response.

    Args:
        role (str): The role of the entity that sent the message (e.g. ``"user"``,
            ``"system"``, ``"assistant"``, ``"tool"``).
        content (str): The content of the message.
            **Optional** Can be ``None`` if refusal or tool_calls are provided.
        refusal (str): The refusal message content.
            **Optional** Supplied if a refusal response is provided.
        name (str): The name of the entity that sent the message. **Optional**.
        tool_calls (List[:py:class:`ToolCall`]): A list of tool calls made by the model.
            **Optional** defaults to ``None``
        tool_call_id (str): The ID of the tool call that this message is a response to.
            **Optional** defaults to ``None``
    """

    role: str
    content: str | None = None
    refusal: str | None = None
    name: str | None = None
    tool_calls: list[ToolCall] | None = None
    tool_call_id: str | None = None

    def __post_init__(self):
        self._validate_field("role", str, True)

        if self.refusal:
            self._validate_field("refusal", str, True)
            if self.content:
                raise ValueError("Both `content` and `refusal` cannot be set")
        elif self.tool_calls:
            self._validate_field("content", str, False)
        else:
            self._validate_field("content", str, True)

        self._validate_field("name", str, False)
        self._convert_dataclass_list("tool_calls", ToolCall, False)
        self._validate_field("tool_call_id", str, False)


@dataclass
class ChatChoiceDelta(_BaseDataclass):
    """
    A streaming message delta in a chat response.

    Args:
        role (str): The role of the entity that sent the message (e.g. ``"user"``,
            ``"system"``, ``"assistant"``, ``"tool"``).
            **Optional** defaults to ``"assistant"``
            This is optional because OpenAI clients can explicitly return None for
            the role
        content (str): The content of the new token being streamed
            **Optional** Can be ``None`` on the last delta chunk or if refusal or
            tool_calls are provided
        refusal (str): The refusal message content.
            **Optional** Supplied if a refusal response is provided.
        name (str): The name of the entity that sent the message. **Optional**.
        tool_calls (List[:py:class:`ToolCall`]): A list of tool calls made by the model.
            **Optional** defaults to ``None``
    """

    role: str | None = "assistant"
    content: str | None = None
    refusal: str | None = None
    name: str | None = None
    tool_calls: list[ToolCall] | None = None

    def __post_init__(self):
        self._validate_field("role", str, False)

        if self.refusal:
            self._validate_field("refusal", str, True)
            if self.content:
                raise ValueError("Both `content` and `refusal` cannot be set")
        self._validate_field("content", str, False)
        self._validate_field("name", str, False)
        self._convert_dataclass_list("tool_calls", ToolCall, False)


@dataclass
class ParamType(_BaseDataclass):
    type: Literal["string", "number", "integer", "object", "array", "boolean", "null"]

    def __post_init__(self):
        self._validate_literal("type", JSON_SCHEMA_TYPES, True)


@dataclass
class ParamProperty(ParamType):
    """
    A single parameter within a function definition.

    Args:
        type (str): The type of the parameter. Possible values are "string", "number", "integer",
            "object", "array", "boolean", or "null", conforming to the JSON Schema specification.
        description (str): A description of the parameter.
            **Optional**, defaults to ``None``
        enum (List[str]): Used to constrain the possible values for the parameter.
            **Optional**, defaults to ``None``
        items (:py:class:`ParamProperty`): If the param is of ``array`` type, this field can be
            used to specify the type of its items. **Optional**, defaults to ``None``
    """

    description: str | None = None
    enum: list[str] | None = None
    items: ParamType | None = None

    def __post_init__(self):
        self._validate_field("description", str, False)
        self._validate_list("enum", str, False)
        self._convert_dataclass("items", ParamType, False)
        super().__post_init__()


@dataclass
class ToolParamsSchema(_BaseDataclass):
    """
    A tool parameter definition.

    Args:
        properties (Dict[str, :py:class:`ParamProperty`]): A mapping of parameter names to
            their definitions.
        type (str): The type of the parameter. Currently only "object" is supported.
        required (List[str]): A list of required parameter names. **Optional**, defaults to ``None``
        additionalProperties (bool): Whether additional properties are allowed in the object.
            **Optional**, defaults to ``None``
    """

    properties: dict[str, ParamProperty]
    type: Literal["object"] = "object"
    required: list[str] | None = None
    additionalProperties: bool | None = None

    def __post_init__(self):
        self._convert_dataclass_map("properties", ParamProperty, True)
        self._validate_literal("type", ["object"], True)
        self._validate_list("required", str, False)
        self._validate_field("additionalProperties", bool, False)


@dataclass
class FunctionToolDefinition(_BaseDataclass):
    """
    Definition for function tools (currently the only supported type of tool).

    Args:
        name (str): The name of the tool.
        description (str): A description of what the tool does, and how it should be used.
            **Optional**, defaults to ``None``
        parameters: A mapping of parameter names to their
            definitions. If not provided, this defines a function without parameters.
            **Optional**, defaults to ``None``
        strict (bool): A flag that represents whether or not the model should
            strictly follow the schema provided.
    """

    name: str
    description: str | None = None
    parameters: ToolParamsSchema | None = None
    strict: bool = False

    def __post_init__(self):
        self._validate_field("name", str, True)
        self._validate_field("description", str, False)
        self._convert_dataclass("parameters", ToolParamsSchema, False)
        self._validate_field("strict", bool, True)

    def to_tool_definition(self):
        """
        Convenience function for wrapping this in a ToolDefinition
        """
        return ToolDefinition(type="function", function=self)


@dataclass
class ToolDefinition(_BaseDataclass):
    """
    Definition for tools that can be called by the model.

    Args:
        function (:py:class:`FunctionToolDefinition`): The definition of a function tool.
        type (str): The type of the tool. Currently only "function" is supported.
    """

    function: FunctionToolDefinition
    type: Literal["function"] = "function"

    def __post_init__(self):
        self._validate_literal("type", ["function"], True)
        self._convert_dataclass("function", FunctionToolDefinition, True)


@dataclass
class ChatParams(_BaseDataclass):
    """
    Common parameters used for chat inference

    Args:
        temperature (float): A param used to control randomness and creativity during inference.
            **Optional**, defaults to ``1.0``
        max_tokens (int): The maximum number of new tokens to generate.
            **Optional**, defaults to ``None`` (unlimited)
        stop (List[str]): A list of tokens at which to stop generation.
            **Optional**, defaults to ``None``
        n (int): The number of responses to generate.
            **Optional**, defaults to ``1``
        stream (bool): Whether to stream back responses as they are generated.
            **Optional**, defaults to ``False``
        top_p (float): An optional param to control sampling with temperature, the model considers
            the results of the tokens with top_p probability mass. E.g., 0.1 means only the tokens
            comprising the top 10% probability mass are considered.
        top_k (int): An optional param for reducing the vocabulary size to top k tokens
            (sorted in descending order by their probabilities).
        frequency_penalty: (float): An optional param of positive or negative value,
            positive values penalize new tokens based on
            their existing frequency in the text so far, decreasing the model's likelihood to repeat
            the same line verbatim.
        presence_penalty: (float): An optional param of positive or negative value,
            positive values penalize new tokens based on whether they appear in the text so far,
            increasing the model's likelihood to talk about new topics.
        custom_inputs (Dict[str, Any]): An optional param to provide arbitrary additional context
            to the model. The dictionary values must be JSON-serializable.
        tools (List[:py:class:`ToolDefinition`]): An optional list of tools that can be called by
            the model.

    .. warning::

        In an upcoming MLflow release, default values for `temperature`, `n` and `stream` will be
        removed. Please provide these values explicitly in your code if needed.
    """

    temperature: float = 1.0
    max_tokens: int | None = None
    stop: list[str] | None = None
    n: int = 1
    stream: bool = False

    top_p: float | None = None
    top_k: int | None = None
    frequency_penalty: float | None = None
    presence_penalty: float | None = None

    custom_inputs: dict[str, Any] | None = None
    tools: list[ToolDefinition] | None = None

    def __post_init__(self):
        self._validate_field("temperature", float, True)
        self._validate_field("max_tokens", int, False)
        self._validate_list("stop", str, False)
        self._validate_field("n", int, True)
        self._validate_field("stream", bool, True)

        self._validate_field("top_p", float, False)
        self._validate_field("top_k", int, False)
        self._validate_field("frequency_penalty", float, False)
        self._validate_field("presence_penalty", float, False)
        self._convert_dataclass_list("tools", ToolDefinition, False)

        # validate that the custom_inputs field is a map from string to string
        if self.custom_inputs is not None:
            if not isinstance(self.custom_inputs, dict):
                raise ValueError(
                    "Expected `custom_inputs` to be a dictionary, "
                    f"received `{type(self.custom_inputs).__name__}`"
                )
            for key, value in self.custom_inputs.items():
                if not isinstance(key, str):
                    raise ValueError(
                        "Expected `custom_inputs` to be of type `Dict[str, Any]`, "
                        f"received key of type `{type(key).__name__}` (key: {key})"
                    )

    @classmethod
    def keys(cls) -> set[str]:
        """
        Return the keys of the dataclass
        """
        return {field.name for field in fields(cls)}


@dataclass()
class ChatCompletionRequest(ChatParams):
    """
    Format of the request object expected by the chat endpoint.

    Args:
        messages (List[:py:class:`ChatMessage`]): A list of :py:class:`ChatMessage`
            that will be passed to the model. **Optional**, defaults to empty list (``[]``)
        temperature (float): A param used to control randomness and creativity during inference.
            **Optional**, defaults to ``1.0``
        max_tokens (int): The maximum number of new tokens to generate.
            **Optional**, defaults to ``None`` (unlimited)
        stop (List[str]): A list of tokens at which to stop generation.
            **Optional**, defaults to ``None``
        n (int): The number of responses to generate.
            **Optional**, defaults to ``1``
        stream (bool): Whether to stream back responses as they are generated.
            **Optional**, defaults to ``False``
        top_p (float): An optional param to control sampling with temperature, the model considers
            the results of the tokens with top_p probability mass. E.g., 0.1 means only the tokens
            comprising the top 10% probability mass are considered.
        top_k (int): An optional param for reducing the vocabulary size to top k tokens
            (sorted in descending order by their probabilities).
        frequency_penalty: (float): An optional param of positive or negative value,
            positive values penalize new tokens based on
            their existing frequency in the text so far, decreasing the model's likelihood to repeat
            the same line verbatim.
        presence_penalty: (float): An optional param of positive or negative value,
            positive values penalize new tokens based on whether they appear in the text so far,
            increasing the model's likelihood to talk about new topics.
        custom_inputs (Dict[str, Any]): An optional param to provide arbitrary additional context
            to the model. The dictionary values must be JSON-serializable.
        tools (List[:py:class:`ToolDefinition`]): An optional list of tools that can be called by
            the model.

    .. warning::

        In an upcoming MLflow release, default values for `temperature`, `n` and `stream` will be
        removed. Please provide these values explicitly in your code if needed.
    """

    messages: list[ChatMessage] = field(default_factory=list)

    def __post_init__(self):
        self._convert_dataclass_list("messages", ChatMessage)
        super().__post_init__()


@dataclass
class TopTokenLogProb(_BaseDataclass):
    """
    Token and its log probability.

    Args:
        token: The token.
        logprob: The log probability of this token, if it is within the top
            20 most likely tokens. Otherwise, the value -9999.0 is used to
            signify that the token is very unlikely.
        bytes: A list of integers representing the UTF-8 bytes representation
            of the token. Useful in instances where characters are represented
            by multiple tokens and their byte representations must be combined
            to generate the correct text representation. Can be null if there
            is no bytes representation for the token.
    """

    token: str
    logprob: float
    bytes: list[int] | None = None

    def __post_init__(self):
        self._validate_field("token", str, True)
        self._validate_field("logprob", float, True)
        self._validate_list("bytes", int, False)


@dataclass
class TokenLogProb(_BaseDataclass):
    """
    Message content token with log probability information.

    Args:
        token: The token.
        logprob: The log probability of this token, if it is within the top
            20 most likely tokens. Otherwise, the value -9999.0 is used to
            signify that the token is very unlikely.
        bytes: A list of integers representing the UTF-8 bytes representation
            of the token. Useful in instances where characters are represented
            by multiple tokens and their byte representations must be combined
            to generate the correct text representation. Can be null if there
            is no bytes representation for the token.
        top_logprobs: List of the most likely tokens and their log probability,
            at this token position. In rare cases, there may be fewer than the
            number of requested top_logprobs returned.
    """

    token: str
    logprob: float
    top_logprobs: list[TopTokenLogProb]
    bytes: list[int] | None = None

    def __post_init__(self):
        self._validate_field("token", str, True)
        self._validate_field("logprob", float, True)
        self._convert_dataclass_list("top_logprobs", TopTokenLogProb)
        self._validate_list("bytes", int, False)


@dataclass
class ChatChoiceLogProbs(_BaseDataclass):
    """
    Log probability information for the choice.

    Args:
        content: A list of message content tokens with log probability information.
    """

    content: list[TokenLogProb] | None = None

    def __post_init__(self):
        self._convert_dataclass_list("content", TokenLogProb, False)


@dataclass
class ChatChoice(_BaseDataclass):
    """
    A single chat response generated by the model.
    ref: https://platform.openai.com/docs/api-reference/chat/object

    Args:
        message (:py:class:`ChatMessage`): The message that was generated.
        index (int): The index of the response in the list of responses.
            Defaults to ``0``
        finish_reason (str): The reason why generation stopped.
            **Optional**, defaults to ``"stop"``
        logprobs (:py:class:`ChatChoiceLogProbs`): Log probability information for the choice.
            **Optional**, defaults to ``None``
    """

    message: ChatMessage
    index: int = 0
    finish_reason: str = "stop"
    logprobs: ChatChoiceLogProbs | None = None

    def __post_init__(self):
        self._validate_field("index", int, True)
        self._validate_field("finish_reason", str, True)
        self._convert_dataclass("message", ChatMessage, True)
        self._convert_dataclass("logprobs", ChatChoiceLogProbs, False)


@dataclass
class ChatChunkChoice(_BaseDataclass):
    """
    A single chat response chunk generated by the model.
    ref: https://platform.openai.com/docs/api-reference/chat/streaming

    Args:
        index (int): The index of the response in the list of responses.
            defaults to ``0``
        delta (:py:class:`ChatChoiceDelta`): The streaming chunk message that was generated.
        finish_reason (str): The reason why generation stopped.
            **Optional**, defaults to ``None``
        logprobs (:py:class:`ChatChoiceLogProbs`): Log probability information for the choice.
            **Optional**, defaults to ``None``
    """

    delta: ChatChoiceDelta
    index: int = 0
    finish_reason: str | None = None
    logprobs: ChatChoiceLogProbs | None = None

    def __post_init__(self):
        self._validate_field("index", int, True)
        self._validate_field("finish_reason", str, False)
        self._convert_dataclass("delta", ChatChoiceDelta, True)
        self._convert_dataclass("logprobs", ChatChoiceLogProbs, False)


@dataclass
class TokenUsageStats(_BaseDataclass):
    """
    Stats about the number of tokens used during inference.

    Args:
        prompt_tokens (int): The number of tokens in the prompt.
            **Optional**, defaults to ``None``
        completion_tokens (int): The number of tokens in the generated completion.
            **Optional**, defaults to ``None``
        total_tokens (int): The total number of tokens used.
            **Optional**, defaults to ``None``
    """

    prompt_tokens: int | None = None
    completion_tokens: int | None = None
    total_tokens: int | None = None

    def __post_init__(self):
        self._validate_field("prompt_tokens", int, False)
        self._validate_field("completion_tokens", int, False)
        self._validate_field("total_tokens", int, False)


@dataclass
class ChatCompletionResponse(_BaseDataclass):
    """
    The full response object returned by the chat endpoint.

    Args:
        choices (List[:py:class:`ChatChoice`]): A list of :py:class:`ChatChoice` objects
            containing the generated responses
        usage (:py:class:`TokenUsageStats`): An object describing the tokens used by the request.
            **Optional**, defaults to ``None``.
        id (str): The ID of the response. **Optional**, defaults to ``None``
        model (str): The name of the model used. **Optional**, defaults to ``None``
        object (str): The object type. Defaults to 'chat.completion'
        created (int): The time the response was created.
            **Optional**, defaults to the current time.
        custom_outputs (Dict[str, Any]): An field that can contain arbitrary additional context.
            The dictionary values must be JSON-serializable.
            **Optional**, defaults to ``None``
    """

    choices: list[ChatChoice]
    usage: TokenUsageStats | None = None
    id: str | None = None
    model: str | None = None
    object: str = "chat.completion"
    created: int = field(default_factory=lambda: int(time.time()))
    custom_outputs: dict[str, Any] | None = None

    def __post_init__(self):
        self._validate_field("id", str, False)
        self._validate_field("object", str, True)
        self._validate_field("created", int, True)
        self._validate_field("model", str, False)
        self._convert_dataclass_list("choices", ChatChoice)
        self._convert_dataclass("usage", TokenUsageStats, False)


@dataclass
class ChatCompletionChunk(_BaseDataclass):
    """
    The streaming chunk returned by the chat endpoint.
    ref: https://platform.openai.com/docs/api-reference/chat/streaming

    Args:
        choices (List[:py:class:`ChatChunkChoice`]): A list of :py:class:`ChatChunkChoice` objects
            containing the generated chunk of a streaming response
        usage (:py:class:`TokenUsageStats`): An object describing the tokens used by the request.
            **Optional**, defaults to ``None``.
        id (str): The ID of the response. **Optional**, defaults to ``None``
        model (str): The name of the model used. **Optional**, defaults to ``None``
        object (str): The object type. Defaults to 'chat.completion.chunk'
        created (int): The time the response was created.
            **Optional**, defaults to the current time.
        custom_outputs (Dict[str, Any]): An field that can contain arbitrary additional context.
            The dictionary values must be JSON-serializable.
            **Optional**, defaults to ``None``
    """

    choices: list[ChatChunkChoice]
    usage: TokenUsageStats | None = None
    id: str | None = None
    model: str | None = None
    object: str = "chat.completion.chunk"
    created: int = field(default_factory=lambda: int(time.time()))
    custom_outputs: dict[str, Any] | None = None

    def __post_init__(self):
        self._validate_field("id", str, False)
        self._validate_field("object", str, True)
        self._validate_field("created", int, True)
        self._validate_field("model", str, False)
        self._convert_dataclass_list("choices", ChatChunkChoice)
        self._convert_dataclass("usage", TokenUsageStats, False)


# turn off formatting for the model signatures to preserve readability
# fmt: off

_token_usage_stats_col_spec = ColSpec(
    name="usage",
    type=Object(
        [
            Property("prompt_tokens", DataType.long),
            Property("completion_tokens", DataType.long),
            Property("total_tokens", DataType.long),
        ]
    ),
    required=False,
)
_custom_inputs_col_spec = ColSpec(name="custom_inputs", type=Map(AnyType()), required=False)
_custom_outputs_col_spec = ColSpec(name="custom_outputs", type=Map(AnyType()), required=False)

CHAT_MODEL_INPUT_SCHEMA = Schema(
    [
        ColSpec(
            name="messages",
            type=Array(
                Object(
                    [
                        Property("role", DataType.string),
                        Property("content", DataType.string, False),
                        Property("name", DataType.string, False),
                        Property("refusal", DataType.string, False),
                        Property("tool_calls", Array(Object([
                            Property("id", DataType.string),
                            Property("function", Object([
                                Property("name", DataType.string),
                                Property("arguments", DataType.string),
                            ])),
                            Property("type", DataType.string),
                        ])), False),
                        Property("tool_call_id", DataType.string, False),
                    ]
                )
            ),
        ),
        ColSpec(name="temperature", type=DataType.double, required=False),
        ColSpec(name="max_tokens", type=DataType.long, required=False),
        ColSpec(name="stop", type=Array(DataType.string), required=False),
        ColSpec(name="n", type=DataType.long, required=False),
        ColSpec(name="stream", type=DataType.boolean, required=False),
        ColSpec(name="top_p", type=DataType.double, required=False),
        ColSpec(name="top_k", type=DataType.long, required=False),
        ColSpec(name="frequency_penalty", type=DataType.double, required=False),
        ColSpec(name="presence_penalty", type=DataType.double, required=False),
        ColSpec(
            name="tools",
            type=Array(
                Object([
                    Property("type", DataType.string),
                    Property("function", Object([
                        Property("name", DataType.string),
                        Property("description", DataType.string, False),
                        Property("parameters", Object([
                            Property("properties", Map(Object([
                                Property("type", DataType.string),
                                Property("description", DataType.string, False),
                                Property("enum", Array(DataType.string), False),
                                Property("items", Object([Property("type", DataType.string)]), False), # noqa
                            ]))),
                            Property("type", DataType.string, False),
                            Property("required", Array(DataType.string), False),
                            Property("additionalProperties", DataType.boolean, False),
                        ])),
                        Property("strict", DataType.boolean, False),
                    ]), False),
                ]),
            ),
            required=False,
        ),
        _custom_inputs_col_spec,
    ]
)

CHAT_MODEL_OUTPUT_SCHEMA = Schema(
    [
        ColSpec(name="id", type=DataType.string),
        ColSpec(name="object", type=DataType.string),
        ColSpec(name="created", type=DataType.long),
        ColSpec(name="model", type=DataType.string),
        ColSpec(
            name="choices",
            type=Array(Object([
                Property("index", DataType.long),
                Property("message", Object([
                    Property("role", DataType.string),
                    Property("content", DataType.string, False),
                    Property("name", DataType.string, False),
                    Property("refusal", DataType.string, False),
                    Property("tool_calls", Array(Object([
                        Property("id", DataType.string),
                        Property("function", Object([
                            Property("name", DataType.string),
                            Property("arguments", DataType.string),
                        ])),
                        Property("type", DataType.string),
                    ])), False),
                    Property("tool_call_id", DataType.string, False),
                ])),
                Property("finish_reason", DataType.string),
            ])),
        ),
        _token_usage_stats_col_spec,
        _custom_outputs_col_spec
    ]
)

CHAT_MODEL_INPUT_EXAMPLE = {
    "messages": [
        {"role": "user", "content": "Hello!"},
    ],
    "temperature": 1.0,
    "max_tokens": 10,
    "stop": ["\n"],
    "n": 1,
    "stream": False,
}

COMPLETIONS_MODEL_INPUT_SCHEMA = Schema(
    [
        ColSpec(name="prompt", type=DataType.string),
        ColSpec(name="temperature", type=DataType.double, required=False),
        ColSpec(name="max_tokens", type=DataType.long, required=False),
        ColSpec(name="stop", type=Array(DataType.string), required=False),
        ColSpec(name="n", type=DataType.long, required=False),
        ColSpec(name="stream", type=DataType.boolean, required=False),
    ]
)

COMPLETIONS_MODEL_OUTPUT_SCHEMA = Schema(
    [
        ColSpec(name="id", type=DataType.string),
        ColSpec(name="object", type=DataType.string),
        ColSpec(name="created", type=DataType.long),
        ColSpec(name="model", type=DataType.string),
        ColSpec(
            name="choices",
            type=Array(
                Object(
                    [
                        Property("index", DataType.long),
                        Property(
                            "text",
                            DataType.string,
                        ),
                        Property("finish_reason", DataType.string),
                    ]
                )
            ),
        ),
        ColSpec(
            name="usage",
            type=Object(
                [
                    Property("prompt_tokens", DataType.long),
                    Property("completion_tokens", DataType.long),
                    Property("total_tokens", DataType.long),
                ]
            ),
        ),
    ]
)

EMBEDDING_MODEL_INPUT_SCHEMA = Schema(
    [
        ColSpec(name="input", type=DataType.string),
    ]
)

EMBEDDING_MODEL_OUTPUT_SCHEMA = Schema(
    [
        ColSpec(name="object", type=DataType.string),
        ColSpec(
            name="data",
            type=Array(
                Object(
                    [
                        Property("index", DataType.long),
                        Property("object", DataType.string),
                        Property("embedding", Array(DataType.double)),
                    ]
                )
            ),
        ),
        ColSpec(
            name="usage",
            type=Object(
                [
                    Property("prompt_tokens", DataType.long),
                    Property("total_tokens", DataType.long),
                ]
            ),
        ),
    ]
)
# fmt: on
