#
#    Copyright 2024 OpenAI
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

from dataclasses import dataclass
from typing import List, Optional

from qwak.llmops.generation.base import ModelResponse
from .chat_completion_token_logprob import ChatCompletionTokenLogprob

__all__ = [
    "ChatCompletionChunk",
    "Choice",
    "ChoiceDelta",
    "ChoiceDeltaFunctionCall",
    "ChoiceDeltaToolCall",
    "ChoiceDeltaToolCallFunction",
    "ChoiceLogprobs",
]


@dataclass
class ChoiceDeltaFunctionCall:
    arguments: Optional[str] = None
    """
    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.
    """

    name: Optional[str] = None
    """The name of the function to call."""


@dataclass
class ChoiceDeltaToolCallFunction:
    arguments: Optional[str] = None
    """
    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.
    """

    name: Optional[str] = None
    """The name of the function to call."""


@dataclass
class ChoiceDeltaToolCall:
    index: int

    id: Optional[str] = None
    """The ID of the tool call."""

    function: Optional[ChoiceDeltaToolCallFunction] = None

    type: Optional[str] = None
    """The type of the tool. Currently, only `function` is supported."""

    def __post_init__(self):
        """Validates that type is 'function' if present."""
        if self.type is not None and self.type != "function":
            raise ValueError(
                f"Invalid type: '{self.type}'. Must be 'function' or None."
            )


@dataclass
class ChoiceDelta:
    content: Optional[str] = None
    """The contents of the chunk message."""

    function_call: Optional[ChoiceDeltaFunctionCall] = None
    """Deprecated and replaced by `tool_calls`.

    The name and arguments of a function that should be called, as generated by the
    model.
    """

    role: Optional[str] = None
    """The role of the author of this message."""

    tool_calls: Optional[List[ChoiceDeltaToolCall]] = None

    def __post_init__(self):
        """Validates that role is one of the allowed values if present."""
        if self.role is not None:
            allowed_roles = {"system", "user", "assistant", "tool"}
            if self.role not in allowed_roles:
                raise ValueError(
                    f"Invalid role: '{self.role}'. "
                    f"Must be one of {allowed_roles} or None."
                )


@dataclass
class ChoiceLogprobs:
    content: Optional[List[ChatCompletionTokenLogprob]] = None
    """A list of message content tokens with log probability information."""


@dataclass
class Choice:
    delta: ChoiceDelta
    """A chat completion delta generated by streamed model responses."""

    index: int
    """The index of the choice in the list of choices."""

    finish_reason: Optional[str] = None
    """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.
    """

    logprobs: Optional[ChoiceLogprobs] = None
    """Log probability information for the choice."""

    def __post_init__(self):
        """Validates that finish_reason is one of the allowed values if present."""
        if self.finish_reason is not None:
            allowed_reasons = {
                "stop",
                "length",
                "tool_calls",
                "content_filter",
                "function_call",
            }
            if self.finish_reason not in allowed_reasons:
                raise ValueError(
                    f"Invalid finish_reason: '{self.finish_reason}'. "
                    f"Must be one of {allowed_reasons} or None."
                )


@dataclass
class ChatCompletionChunk(ModelResponse):
    id: str
    """A unique identifier for the chat completion. Each chunk has the same ID."""

    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.

    Each chunk has the same timestamp.
    """

    model: str
    """The model to generate the completion."""

    object: str
    """The object type, which is always `chat.completion.chunk`."""

    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.
    """

    def __post_init__(self):
        """Validates that the object type is correct."""
        if self.object != "chat.completion.chunk":
            raise ValueError(
                f"Invalid object type: '{self.object}'. Must be 'chat.completion.chunk'"
            )
