import datetime
from typing import List, Literal, Optional, Union

from pydantic import BaseModel

# class ToolCallFunction(BaseModel):
#     name: str
#     arguments: str


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


class ToolCall(BaseModel):
    id: str
    # "Currently, only function is supported"
    type: Literal["function"] = "function"
    # function: ToolCallFunction
    function: FunctionCall


class ToolCallDenial(ToolCall):
    reason: Optional[str] = None


class LogProbToken(BaseModel):
    token: str
    logprob: float
    bytes: Optional[List[int]]


# Legacy?
class MessageContentLogProb(BaseModel):
    token: str
    logprob: float
    bytes: Optional[List[int]]
    top_logprobs: Optional[List[LogProbToken]]


class TopLogprob(BaseModel):
    token: str
    bytes: Optional[List[int]] = None
    logprob: float


class ChatCompletionTokenLogprob(BaseModel):
    token: str
    bytes: Optional[List[int]] = None
    logprob: float
    top_logprobs: List[TopLogprob]


class ChoiceLogprobs(BaseModel):
    content: Optional[List[ChatCompletionTokenLogprob]] = None

    refusal: Optional[List[ChatCompletionTokenLogprob]] = None


class Message(BaseModel):
    content: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None
    role: str
    function_call: Optional[FunctionCall] = None  # Deprecated
    reasoning_content: Optional[str] = None  # Used in newer reasoning APIs, e.g. DeepSeek
    reasoning_content_signature: Optional[str] = None  # NOTE: for Anthropic
    redacted_reasoning_content: Optional[str] = None  # NOTE: for Anthropic
    omitted_reasoning_content: bool = False  # NOTE: for OpenAI o1/o3


class Choice(BaseModel):
    finish_reason: str
    index: int
    message: Message
    logprobs: Optional[ChoiceLogprobs] = None
    seed: Optional[int] = None  # found in TogetherAI


class UsageStatisticsPromptTokenDetails(BaseModel):
    # None means provider didn't report this field, 0 means provider reported 0
    cached_tokens: Optional[int] = None  # OpenAI/Gemini: tokens served from cache
    cache_read_tokens: Optional[int] = None  # Anthropic: tokens read from cache
    cache_creation_tokens: Optional[int] = None  # Anthropic: tokens written to cache
    # NOTE: OAI specific
    # audio_tokens: int = 0

    def __add__(self, other: "UsageStatisticsPromptTokenDetails") -> "UsageStatisticsPromptTokenDetails":
        # Helper to add optional ints (None + None = None, None + N = N, N + M = N+M)
        def add_optional(a: Optional[int], b: Optional[int]) -> Optional[int]:
            if a is None and b is None:
                return None
            return (a or 0) + (b or 0)

        return UsageStatisticsPromptTokenDetails(
            cached_tokens=add_optional(self.cached_tokens, other.cached_tokens),
            cache_read_tokens=add_optional(self.cache_read_tokens, other.cache_read_tokens),
            cache_creation_tokens=add_optional(self.cache_creation_tokens, other.cache_creation_tokens),
        )


class UsageStatisticsCompletionTokenDetails(BaseModel):
    # None means provider didn't report this field, 0 means provider reported 0
    reasoning_tokens: Optional[int] = None
    # NOTE: OAI specific
    # audio_tokens: int = 0
    # accepted_prediction_tokens: int = 0
    # rejected_prediction_tokens: int = 0

    def __add__(self, other: "UsageStatisticsCompletionTokenDetails") -> "UsageStatisticsCompletionTokenDetails":
        # Helper to add optional ints (None + None = None, None + N = N, N + M = N+M)
        def add_optional(a: Optional[int], b: Optional[int]) -> Optional[int]:
            if a is None and b is None:
                return None
            return (a or 0) + (b or 0)

        return UsageStatisticsCompletionTokenDetails(
            reasoning_tokens=add_optional(self.reasoning_tokens, other.reasoning_tokens),
        )


class UsageStatistics(BaseModel):
    completion_tokens: int = 0
    prompt_tokens: int = 0
    total_tokens: int = 0

    prompt_tokens_details: Optional[UsageStatisticsPromptTokenDetails] = None
    completion_tokens_details: Optional[UsageStatisticsCompletionTokenDetails] = None

    def __add__(self, other: "UsageStatistics") -> "UsageStatistics":
        if self.prompt_tokens_details is None and other.prompt_tokens_details is None:
            total_prompt_tokens_details = None
        elif self.prompt_tokens_details is None:
            total_prompt_tokens_details = other.prompt_tokens_details
        elif other.prompt_tokens_details is None:
            total_prompt_tokens_details = self.prompt_tokens_details
        else:
            total_prompt_tokens_details = self.prompt_tokens_details + other.prompt_tokens_details

        if self.completion_tokens_details is None and other.completion_tokens_details is None:
            total_completion_tokens_details = None
        elif self.completion_tokens_details is None:
            total_completion_tokens_details = other.completion_tokens_details
        elif other.completion_tokens_details is None:
            total_completion_tokens_details = self.completion_tokens_details
        else:
            total_completion_tokens_details = self.completion_tokens_details + other.completion_tokens_details

        return UsageStatistics(
            completion_tokens=self.completion_tokens + other.completion_tokens,
            prompt_tokens=self.prompt_tokens + other.prompt_tokens,
            total_tokens=self.total_tokens + other.total_tokens,
            prompt_tokens_details=total_prompt_tokens_details,
            completion_tokens_details=total_completion_tokens_details,
        )


class ChatCompletionResponse(BaseModel):
    """https://platform.openai.com/docs/api-reference/chat/object"""

    id: str
    choices: List[Choice]
    created: Union[datetime.datetime, int]
    model: Optional[str] = None  # NOTE: this is not consistent with OpenAI API standard, however is necessary to support local LLMs
    # system_fingerprint: str  # docs say this is mandatory, but in reality API returns None
    system_fingerprint: Optional[str] = None
    # object: str = Field(default="chat.completion")
    object: Literal["chat.completion"] = "chat.completion"
    usage: UsageStatistics

    def __str__(self):
        return self.model_dump_json(indent=4)


class FunctionCallDelta(BaseModel):
    # arguments: Optional[str] = None
    name: Optional[str] = None
    arguments: Optional[str] = None
    # name: str


class ToolCallDelta(BaseModel):
    index: int
    id: Optional[str] = None
    # "Currently, only function is supported"
    type: Literal["function"] = "function"
    # function: ToolCallFunction
    function: Optional[FunctionCallDelta] = None


class MessageDelta(BaseModel):
    """Partial delta stream of a Message

    Example ChunkResponse:
    {
        'id': 'chatcmpl-9EOCkKdicNo1tiL1956kPvCnL2lLS',
        'object': 'chat.completion.chunk',
        'created': 1713216662,
        'model': 'gpt-4-0613',
        'system_fingerprint': None,
        'choices': [{
            'index': 0,
            'delta': {'content': 'User'},
            'logprobs': None,
            'finish_reason': None
        }]
    }
    """

    content: Optional[str] = None
    reasoning_content: Optional[str] = None
    reasoning_content_signature: Optional[str] = None  # NOTE: for Anthropic
    redacted_reasoning_content: Optional[str] = None  # NOTE: for Anthropic
    tool_calls: Optional[List[ToolCallDelta]] = None
    role: Optional[str] = None
    function_call: Optional[FunctionCallDelta] = None  # Deprecated


class ChunkChoice(BaseModel):
    finish_reason: Optional[str] = None  # NOTE: when streaming will be null
    index: int
    delta: MessageDelta
    logprobs: Optional[ChoiceLogprobs] = None


class ChatCompletionChunkResponse(BaseModel):
    """https://platform.openai.com/docs/api-reference/chat/streaming"""

    id: str
    choices: List[ChunkChoice]
    created: Union[datetime.datetime, int]
    model: str
    # system_fingerprint: str  # docs say this is mandatory, but in reality API returns None
    system_fingerprint: Optional[str] = None
    # object: str = Field(default="chat.completion")
    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
    output_tokens: int = 0
