from enum import Enum, IntEnum, StrEnum
from typing import Annotated, Dict, List, Literal, Optional, Union

from pydantic import BaseModel, ConfigDict, Field, RootModel


class LLMBaseModel(BaseModel):
    model_config = ConfigDict(
        use_enum_values=True,
        validate_default=True,
    )


class Models(StrEnum):
    """
    ID of the model to use. See the [model endpoint compatibility](/docs/models/model-endpoint-compatibility) table for
    details on which models work with the Chat API.
    """

    gpt_4o = "gpt-4o"
    gpt_4o_mini = "gpt-4o-mini"

    gpt_4_turbo = "gpt-4-turbo"
    gpt_4_turbo_2024_04_09 = "gpt-4-turbo-2024-04-09"
    gpt_4_0125_preview = "gpt-4-0125-preview"
    gpt_4_turbo_preview = "gpt-4-turbo-preview"
    gpt_4_1106_preview = "gpt-4-1106-preview"
    gpt_4_vision_preview = "gpt-4-vision-preview"
    gpt_4 = "gpt-4"
    gpt_4_0314 = "gpt-4-0314"
    gpt_4_0613 = "gpt-4-0613"
    gpt_4_32k = "gpt-4-32k"
    gpt_4_32k_0314 = "gpt-4-32k-0314"
    gpt_4_32k_0613 = "gpt-4-32k-0613"
    gpt_3_5_turbo = "gpt-3.5-turbo"
    gpt_3_5_turbo_16k = "gpt-3.5-turbo-16k"
    gpt_3_5_turbo_0301 = "gpt-3.5-turbo-0301"
    gpt_3_5_turbo_0613 = "gpt-3.5-turbo-0613"
    gpt_3_5_turbo_1106 = "gpt-3.5-turbo-1106"
    gpt_3_5_turbo_0125 = "gpt-3.5-turbo-0125"
    gpt_3_5_turbo_16k_0613 = "gpt-3.5-turbo-16k-0613"
    gpt_4_1 = "gpt-4.1"
    gpt_4_1_mini = "gpt-4.1-mini"
    gpt_4_1_nano = "gpt-4.1-nano"

    o1 = "o1"
    o1_pro = "o1-pro"
    o3 = "o3"
    o3_mini = "o3-mini"
    o4_mini = "o4-mini"

    gemini_pro_vision = "gemini/gemini-pro-vision"
    gemini_2_0_flash = "gemini/gemini-2.0-flash"
    gemini_2_5_flash = "gemini/gemini-2.5-flash-preview-05-20"
    gemini_2_5_pro = "gemini/gemini-2.5-pro-preview-05-06"


class Role(StrEnum):
    """
    The role of the messages author, in this case `system`.
    """

    system = "system"
    user = "user"
    assistant = "assistant"
    tool = "tool"
    function = "function"


class ToolType(StrEnum):
    """
    The type of the tool. Currently, only `function` is supported.
    """

    function = "function"


class FunctionCall(LLMBaseModel):
    """
    The function that the model called.
    """

    name: str
    """
    The name of the function to call.
    """
    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.
    """


class ChatCompletionMessageToolCall(LLMBaseModel):

    id: str
    """
    The ID of the tool call.
    """
    type: ToolType
    """
    The type of the tool. Currently, only `function` is supported.
    """
    function: FunctionCall
    """
    The function that the model called.
    """


class SystemMessage(LLMBaseModel):

    content: str
    """
    The contents of the system message.
    """
    role: Literal[Role.system] = Field(default=Role.system)
    """
    The role of the messages author, in this case `system`.
    """
    name: Optional[str] = None
    """
    An optional name for the participant. Provides the model information to differentiate between participants
    of the same role.
    """


class AssistantMessage(LLMBaseModel):

    content: Optional[str] = None
    """
    The contents of the assistant message. Required unless `tool_calls` or `function_call` is specified.

    """
    role: Literal[Role.assistant] = Field(default=Role.assistant)
    """
    The role of the messages author, in this case `assistant`.
    """
    name: Optional[str] = None
    """
    An optional name for the participant. Provides the model information to differentiate between participants of the same role.
    """
    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
    """
    The tool calls generated by the model, such as function calls.
    """


class ContentPartType(StrEnum):
    text = "text"
    """
    The text content type.
    """
    image_url = "image_url"
    """
    The image content type.
    """
    file_url = "file_url"
    """
    The file content type
    """


class TextContentPart(LLMBaseModel):

    type: Literal[ContentPartType.text] = ContentPartType.text
    """
    The type of the content part.
    """
    text: str
    """
    The text content.
    """


class ImageDetail(Enum):
    """
    Specifies the detail level of the image. Learn more in the [Vision guide]
    (/docs/guides/vision/low-or-high-fidelity-image-understanding).
    """

    auto = "auto"
    low = "low"
    high = "high"


class ImageUrl(LLMBaseModel):

    url: str
    """
    Either a URL of the image or the base64 encoded image data.
    """
    detail: Optional[ImageDetail] = ImageDetail.auto
    """
    Specifies the detail level of the image. Learn more in the [Vision guide]
    (/docs/guides/vision/low-or-high-fidelity-image-understanding).
    """


class FileUrl(LLMBaseModel):
    url: str


class ImageContentPart(LLMBaseModel):

    type: Literal[ContentPartType.image_url] = ContentPartType.image_url
    """
    The type of the content part.
    """
    image_url: ImageUrl


class FileContentPart(LLMBaseModel):

    type: Literal[ContentPartType.file_url] = ContentPartType.file_url
    """
    The type of the content part
    """
    file_url: FileUrl


class ArrayOfContentParts(RootModel[List[Union[TextContentPart, ImageContentPart, FileContentPart]]]):
    root: Annotated[
        List[Union[TextContentPart, ImageContentPart, FileContentPart]],
        Field(min_length=1, title="Array of content parts"),
    ]
    """
    An array of content parts with a defined type, each can be of type `text` or `image_url` when passing in images.
    You can pass multiple images by adding multiple `image_url` content parts. Image input is only supported when using
    the `gpt-4-visual-preview` model.
    """


class UserMessage(LLMBaseModel):

    content: Union[str, ArrayOfContentParts]
    """
    The contents of the user message.

    """
    role: Literal[Role.user] = Field(default=Role.user)
    """
    The role of the messages author, in this case `user`.
    """
    name: Optional[str] = None
    """
    An optional name for the participant. Provides the model information to differentiate between participants of the same role.
    """


class ToolMessage(LLMBaseModel):

    role: Literal[Role.tool] = Field(default=Role.tool)
    """
    The role of the messages author, in this case `tool`.
    """
    content: str
    """
    The contents of the tool message.
    """
    tool_call_id: str
    """
    Tool call that this message is responding to.
    """


class FunctionMessage(LLMBaseModel):

    role: Literal[Role.function] = Field(default=Role.function)
    """
    The role of the messages author, in this case `function`.
    """
    content: str
    """
    The contents of the function message.
    """
    name: str
    """
    The name of the function to call.
    """


class ResponseType(StrEnum):
    """
    Must be one of `text` or `json_object`.
    """

    text = "text"
    json_object = "json_object"


class Function(LLMBaseModel):
    name: str
    """
    The name of the function to call.
    """


class ChatCompletionNamedToolChoice(LLMBaseModel):
    """
    Specifies a tool the model should use. Use to force the model to call a specific function.
    """

    type: ToolType
    """
    The type of the tool. Currently, only `function` is supported.
    """
    function: Function


class ResponseFormat(LLMBaseModel):
    """
    An object specifying the format that the model must output. Compatible with [GPT-4 Turbo](/docs/models/gpt-4-and-gpt-4-turbo)
    and all GPT-3.5 Turbo models newer than `gpt-3.5-turbo-1106`.

    Setting to `{ "type": "json_object" }` enables JSON mode, which guarantees the message the model generates is valid JSON.

    **Important:** when using JSON mode, you **must** also instruct the model to produce JSON yourself via a system or user message.
    Without this, the model may generate an unending stream of whitespace until the generation reaches the token limit,
    resulting in a long-running and seemingly "stuck" request. Also note that the message content may be partially cut off if
    `finish_reason="length"`, which indicates the generation exceeded `max_tokens` or the conversation exceeded the max context length.

    """

    type: Annotated[Optional[ResponseType], Field("text", examples=["json_object"])]
    """
    Must be one of `text` or `json_object`.
    """


class Stop(RootModel[List[str]]):
    root: Annotated[List[str], Field(max_length=4, min_length=1)]
    """
    Up to 4 sequences where the API will stop generating further tokens.

    """


class FunctionParameters(LLMBaseModel):
    """
    The parameters the functions accepts, described as a JSON Schema object.
    See the [guide](/docs/guides/text-generation/function-calling) for examples, and the
    [JSON Schema reference](https://json-schema.org/understanding-json-schema/) for documentation about the format.

    Omitting `parameters` defines a function with an empty parameter list.
    """

    model_config = ConfigDict(
        extra="allow",
    )


class FunctionObject(LLMBaseModel):
    description: Optional[str] = None
    """
    A description of what the function does, used by the model to choose when and how to call the function.
    """
    name: str
    """
    The name of the function to be called. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length of 64.
    """
    parameters: Optional[FunctionParameters] = None


class ChatCompletionToolChoiceOption(Enum):
    """
    `none` means the model will not call a function and instead generates a message. `auto` means
    the model can pick between generating a message or calling a function.

    """

    none = "none"
    auto = "auto"


class ChatCompletionTool(LLMBaseModel):
    type: ToolType
    """
    The type of the tool. Currently, only `function` is supported.
    """
    function: FunctionObject


class ChatCompletionRequest(LLMBaseModel):
    """
    A request to generate chat completions.
    """

    messages: Annotated[
        List[
            Union[
                SystemMessage,
                UserMessage,
                AssistantMessage,
                ToolMessage,
                FunctionMessage,
            ]
        ],
        Field(min_length=1),
    ]
    """
    A list of messages comprising the conversation so far.
    [Example Python code](https://cookbook.openai.com/examples/how_to_format_inputs_to_chatgpt_models).
    """
    frequency_penalty: Annotated[Optional[float], Field(0, ge=-2.0, le=2.0)] = None
    """
    Number between -2.0 and 2.0. 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.

    [See more information about frequency and presence penalties.](/docs/guides/text-generation/parameter-details)

    """
    logit_bias: Optional[Dict[str, int]] = None
    """
    Modify the likelihood of specified tokens appearing in the completion.

    Accepts a JSON object that maps tokens (specified by their token ID in the tokenizer) to an associated bias value from
    -100 to 100. Mathematically, the bias is added to the logits generated by the model prior to sampling.
    The exact effect will vary per model, but values between -1 and 1 should decrease or increase likelihood of selection;
    values like -100 or 100 should result in a ban or exclusive selection of the relevant token.

    """
    logprobs: Optional[bool] = False
    """
    Whether to return log probabilities of the output tokens or not. If true, returns the log probabilities of each
    output token returned in the `content` of `message`.
    """
    top_logprobs: Annotated[Optional[int], Field(None, ge=0, le=20)] = None
    """
    An integer between 0 and 20 specifying the number of most likely tokens to return at each token position,
    each with an associated log probability. `logprobs` must be set to `true` if this parameter is used.
    """
    max_tokens: Optional[int] = None
    """
    The maximum number of [tokens](/tokenizer) that can be generated in the chat completion.

    The total length of input tokens and generated tokens is limited by the model's context length.
    [Example Python code](https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken) for counting tokens.

    """
    n: Annotated[Optional[int], Field(1, examples=[1], ge=1, le=128)] = None
    """
    How many chat completion choices to generate for each input message. Note that you will be charged based on the
    number of generated tokens across all of the choices. Keep `n` as `1` to minimize costs.
    """
    presence_penalty: Annotated[Optional[float], Field(0, ge=-2.0, le=2.0)] = None
    """
    Number between -2.0 and 2.0. 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.

    [See more information about frequency and presence penalties.](/docs/guides/text-generation/parameter-details)

    """
    response_format: Optional[ResponseFormat] = None
    """
    An object specifying the format that the model must output. Compatible with [GPT-4 Turbo](/docs/models/gpt-4-and-gpt-4-turbo)
    and all GPT-3.5 Turbo models newer than `gpt-3.5-turbo-1106`.

    Setting to `{ "type": "json_object" }` enables JSON mode, which guarantees the message the model generates is valid JSON.

    **Important:** when using JSON mode, you **must** also instruct the model to produce JSON yourself via a system or user message.
    Without this, the model may generate an unending stream of whitespace until the generation reaches the token limit,
    resulting in a long-running and seemingly "stuck" request. Also note that the message content may be partially cut off if
    `finish_reason="length"`, which indicates the generation exceeded `max_tokens` or the conversation exceeded the max context
    length.

    """
    seed: Annotated[Optional[int], Field(None, ge=-9223372036854775808, le=9223372036854775807)] = None
    """
    This feature is in Beta.
    If specified, our system will make a best effort to sample deterministically, such that repeated requests with the same
    `seed` and parameters should return the same result.
    Determinism is not guaranteed, and you should refer to the `system_fingerprint` response parameter to monitor changes in
    the backend.

    """
    stop: Optional[Union[str, Stop]] = None
    """
    Up to 4 sequences where the API will stop generating further tokens.

    """
    temperature: Annotated[Optional[float], Field(1, examples=[1], ge=0.0, le=2.0)] = None
    """
    What sampling temperature to use, between 0 and 2. Higher values like 0.8 will make the output more random,
    while lower values like 0.2 will make it more focused and deterministic.

    We generally recommend altering this or `top_p` but not both.

    """
    top_p: Annotated[Optional[float], Field(1, examples=[1], ge=0.0, le=1.0)] = None
    """
    An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of
    the tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered.

    We generally recommend altering this or `temperature` but not both.

    """
    tools: Optional[List[ChatCompletionTool]] = None
    """
    A list of tools the model may call. Currently, only functions are supported as a tool. Use this to provide a
    list of functions the model may generate JSON inputs for. A max of 128 functions are supported.

    """
    tool_choice: Optional[Union[ChatCompletionToolChoiceOption, ChatCompletionNamedToolChoice]] = None
    """
    Controls which (if any) function is called by the model.
    `none` means the model will not call a function and instead generates a message.
    `auto` means the model can pick between generating a message or calling a function.
    Specifying a particular function via `{"type": "function", "function": {"name": "my_function"}}` forces the model
    to call that function.

    `none` is the default when no functions are present. `auto` is the default if functions are present.

    """
    user: Annotated[Optional[str], Field(None, examples=["user-1234"])] = None
    """
    A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.
    [Learn more](/docs/guides/safety-best-practices/end-user-ids).

    """

    # Other configuration parameters
    stream_options: Optional[dict] = None
    """
    A dictionary containing options for the streaming response. Only set this when you set stream: true.
    """

    timeout: Optional[float] = None
    """
    The maximum time in seconds to wait for the completion.
    If the completion is not returned within this time, the request will time out.
    """

    # openai v1.0+ new params
    deployment_id: Optional[str] = None
    """
    The deployment ID for generating completions. Defaults to None.
    """

    max_retries: Optional[int] = None
    """
    The number of retries to attempt (default is 0).
    """

    mock_response: Optional[str] = None
    """
    If provided, return a mock completion response for testing or debugging purposes (default is None).
    """


class FinishReason(StrEnum):
    """
    The reason the model stopped generating tokens. This will be `stop` if the model hit a natural stop point or a provided
    stop sequence,
    `length` if the maximum number of tokens specified in the request was reached,
    `content_filter` if content was omitted due to a flag from our content filters,
    `tool_calls` if the model called a tool, or `function_call` (deprecated) if the model called a function.

    """

    stop = "stop"
    length = "length"
    tool_calls = "tool_calls"
    content_filter = "content_filter"
    function_call = "function_call"


class TopLogprob(LLMBaseModel):
    token: str
    """
    The token.
    """
    logprob: float
    """
    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: List[int]
    """
    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.
    """


class ChatCompletionTokenLogprob(LLMBaseModel):
    token: str
    """
    The token.
    """
    logprob: float
    """
    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: List[int]
    """
    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[TopLogprob]
    """
    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.
    """


class Choice(LLMBaseModel):

    finish_reason: FinishReason
    """
    The reason the model stopped generating tokens. This will be `stop` if the model hit a natural stop point or a
    provided stop sequence,
    `length` if the maximum number of tokens specified in the request was reached,
    `content_filter` if content was omitted due to a flag from our content filters,
    `tool_calls` if the model called a tool, or `function_call` (deprecated) if the model called a function.

    """
    index: int
    """
    The index of the choice in the list of choices.
    """
    message: Optional[AssistantMessage] = None


class CompletionUsage(LLMBaseModel):
    """
    Usage statistics for the completion request.
    """

    completion_tokens: int
    """
    Number of tokens in the generated completion.
    """
    prompt_tokens: int
    """
    Number of tokens in the prompt.
    """
    total_tokens: int
    """
    Total number of tokens used in the request (prompt + completion).
    """


class ChatCompletionResponse(LLMBaseModel):
    """
    Represents a chat completion response returned by model, based on the provided input.
    """

    id: str
    """
    A unique identifier for the chat completion.
    """
    choices: List[Choice]
    """
    A list of chat completion choices. Can be more than one if `n` is greater than 1.
    """
    created: int
    """
    The Unix timestamp (in seconds) of when the chat completion was created.
    """
    model: str
    """
    The model used for the chat completion.
    """
    system_fingerprint: Optional[str] = None
    """
    This fingerprint represents the backend configuration that the model runs with.

    Can be used in conjunction with the `seed` request parameter to understand when backend changes have been made
    that might impact determinism.

    """

    usage: Optional[CompletionUsage] = None
    """
    Usage statistics for the completion request.
    """


class ChatCompletionResponseEmpty(LLMBaseModel):
    """
    Represents an empty chat completion response.
    """

    id: str
    """
    A unique identifier for the chat completion.
    """
    created: int
    """
    The Unix timestamp (in seconds) of when the chat completion was created.
    """


class ResponseCodes(IntEnum):
    BAD_REQUEST_ERROR = 400
    AUTHENTICATION_ERROR = 401
    PERMISSION_DENIED_ERROR = 403
    UNPROCESSABLE_ENTITY_ERROR = 422
    NOT_FOUND_ERROR = 404
    RATE_LIMIT_ERROR = 429
    INTERNAL_SERVER_ERROR = 500
    CONFLICT_ERROR = 409
    CONTEXT_WINDOW_EXCEEDED_ERROR = 1001
    CONTENT_POLICY_VIOLATION_ERROR = 1002
    API_CONNECTION_ERROR = 1003
    API_TIMEOUT_ERROR = 1004


class ChatCompletionError(LLMBaseModel):
    """
    Represents an error response returned by the API.
    """

    status_code: ResponseCodes = ResponseCodes.BAD_REQUEST_ERROR
    """
    The status code of the response.
    """
    message: str
    """
    The error message.
    """
    response: Optional[str] = None
    """
    The response object.
    """
    model: Optional[str] = None
    """
    The model used for the chat completion.
    """
    request_messages: List[
        Union[
            SystemMessage,
            UserMessage,
            AssistantMessage,
            ToolMessage,
            FunctionMessage,
        ]
    ]
    """
    A list of messages comprising the conversation so far.
    """


class APIError(Exception):
    status_code: int
    message: str
