from typing import List, Optional, Union, Dict, Any, Literal
from pydantic import BaseModel, Field


class ChatMessage(BaseModel):
    role: str = Field()
    content: str = Field()


class Logprob(BaseModel):
    token: str
    logprob: float
    top_logprobs: Optional[List[Dict[str, Any]]] = None


class LogprobsContent(BaseModel):
    content: Optional[List[Logprob]] = None
    refusal: Optional[List[Logprob]] = None


class FunctionCall(BaseModel):
    name: str
    arguments: str


class ChatCompletionMessage(BaseModel):
    role: Optional[str] = Field(
        None, description="The role of the author of this message"
    )
    content: Optional[str] = Field(None, description="The contents of the message")
    reasoning_content: Optional[str] = Field(
        None, description="The reasoning contents of the message"
    )
    tool_calls: Optional[List[Dict[str, Any]]] = Field(
        None, description="Tool calls generated by the model"
    )


class PromptTokensDetails(BaseModel):
    cached_tokens: int


class CompletionTokensDetails(BaseModel):
    reasoning_tokens: int
    accepted_prediction_tokens: int
    rejected_prediction_tokens: int


class Usage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    prompt_tokens_details: Optional[PromptTokensDetails]
    # completion_tokens_details: CompletionTokensDetails


class ChatCompletionChoice(BaseModel):
    index: int
    message: Optional[ChatCompletionMessage] = None
    delta: Optional[ChatCompletionMessage] = None
    logprobs: Optional[LogprobsContent] = None
    finish_reason: Optional[str] = Field(
        ..., description="Reason for stopping: stop, length, content_filter, tool_calls"
    )


class ChatCompletion(BaseModel):
    id: str = Field(..., description="Unique identifier for the chat completion")
    object: Literal["chat.completion"] = "chat.completion"
    created: int = Field(..., description="Unix timestamp of creation")
    model: str
    choices: List[ChatCompletionChoice]
    usage: Usage


class ChatCompletionChunk(BaseModel):
    id: str = Field(..., description="Unique identifier for the chat completion")
    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
    created: int = Field(..., description="Unix timestamp of creation")
    model: str
    choices: List[ChatCompletionChoice]
    usage: Optional[Usage]