from typing import Any, Literal

from pydantic import BaseModel, Field


class ChatMessage(BaseModel):
    role: Literal["system", "user", "assistant", "tool", "function"]
    content: str | None = None
    name: str | None = None
    tool_call_id: str | None = None
    function_call: dict[str, Any] | None = None  # For function/tool messages


class FunctionParameters(BaseModel):
    # Accept any structure for function parameters
    RootModel: dict[str, Any]


class FunctionDefinition(BaseModel):
    name: str
    description: str | None = None
    parameters: FunctionParameters


class ToolFunction(BaseModel):
    type: Literal["function"]
    function: FunctionDefinition


class ResponseFormat(BaseModel):
    type: Literal["text", "json_object"]


class LogitBias(BaseModel):
    RootModel: dict[str, float]


class ChatCompletionRequest(BaseModel):
    model: str
    messages: list[ChatMessage]
    requirements: list[str | None] | None = Field(default_factory=list)
    functions: list[FunctionDefinition] | None = None
    function_call: Literal["none", "auto"] | dict[str, str] | None = None
    tools: list[ToolFunction] | None = None
    tool_choice: Literal["none", "auto"] | dict[str, Any] | None = None
    temperature: float | None = Field(default=1.0, ge=0, le=2)
    top_p: float | None = Field(default=1.0, ge=0, le=1)
    n: int | None = Field(default=1, ge=1)
    stream: bool | None = False
    stop: str | list[str] | None = None
    max_tokens: int | None = None
    presence_penalty: float | None = Field(default=0, ge=-2, le=2)
    frequency_penalty: float | None = Field(default=0, ge=-2, le=2)
    logit_bias: dict[str, float] | None = None
    user: str | None = None
    seed: int | None = None
    response_format: ResponseFormat | None = None

    # For future/undocumented fields
    extra: dict[str, Any] = Field(default_factory=dict)

    class Config:
        extra = "allow"


# Taking this from OpenAI types https://github.com/openai/openai-python/blob/main/src/openai/types/chat/chat_completion.py,
class ChatCompletionMessage(BaseModel):
    content: str | None = None
    """The contents of the message."""

    refusal: str | None = None
    """The refusal message generated by the model."""

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


class Choice(BaseModel):
    index: int
    """The index of the choice in the list of choices."""

    message: ChatCompletionMessage
    """A chat completion message generated by the model."""


class ChatCompletion(BaseModel):
    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."""

    object: Literal["chat.completion"]
    """The object type, which is always `chat.completion`."""
