"""
💭 prompted.types.chat_completions

Contains base types used within the `prompted` package. These types
are based on the Chat Completions specification.
"""

from typing import (
    Any,
    Dict,
    Iterable,
    Literal,
    List,
    Optional,
    Union,
)
from typing_extensions import (
    Required,
    NotRequired,
    TypeAlias,
    TypedDict,
)
from pydantic import BaseModel, Field

__all__ = (
    "FunctionParameters",
    "Function",
    "Tool",
    "FunctionCall",
    "ToolCall",
    "MessageContentImageURL",
    "MessageContentImagePart",
    "MessageContentAudioPart",
    "MessageContentTextPart",
    "MessageContentPart",
    "MessageContent",
    "MessageTextContent",
    "MessageRole",
    "Message",
    "Subscriptable",
    "TopLogprob",
    "TokenLogprob",
    "ChoiceLogprobs",
    "CompletionFunction",
    "CompletionToolCall",
    "CompletionMessage",
    "CompletionUsage",
    "Completion",
    "CompletionChunk",
    "Embedding",
    "EmbeddingsUsage", 
)


# ----------------------------------------------------------------------------
# Tool & Function Calling
#
# This block contains both types for objects sent *AS A RESPONSE* (tool calls)
# as well as the schema for the user-defined functions that can be called.
# ----------------------------------------------------------------------------


FunctionParameters: TypeAlias = Dict[str, Any]
"""
A type alias for the parameters of a function.

This type is used to describe the parameters of a function that can be called.
See the OpenAI guide: https://platform.openai.com/docs/guides/function-calling
"""


class Function(TypedDict):
    """
    A dictionary representing a function that can be called.

    Example:
    ```python
    {
        "name": "function_name",
        "description": "function_description",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The city and state, e.g. San Francisco, CA"
                },
                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}
            },
            "required": ["location"]
        },
        "strict": True, # Optional
    }
    ```
    """

    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.
    """
    description: NotRequired[str]
    """
    A description of what the function does, used by the model to choose when and
    how to call the function.
    """
    parameters: FunctionParameters
    """
    The parameters the functions accepts, described as a JSON Schema object.
    See the OpenAI guide for examples: https://platform.openai.com/docs/guides/function-calling
    and the JSON Schema reference for documentation: https://json-schema.org/understanding-json-schema/
    """
    strict: NotRequired[bool]
    """
    Whether to enable strict schema adherence when generating the function call.
    If set to true, the model will follow the exact schema defined in the parameters field.
    Only a subset of JSON Schema is supported when strict is true.
    """


class Tool(TypedDict):
    """
    A dictionary representing a tool that can be called.

    Example:
    ```python
    {
        'type': 'function',
        'function': {
            'name': 'my_web_tool',
            'parameters': {
                'type': 'object',
                'properties': {'url': {'type': 'string', 'description': 'The URL of the website to get the title of.'}},
                'required': ['url'],
                'additionalProperties': False
            },
            'description': 'This is a tool that can be used to get the title of a website.\n'
        }
    }
    """

    type: Literal["function"]
    """
    The type of the tool. Currently, only `function` is supported.
    """
    function: Function
    """
    The function that the tool calls.
    """


# response types
class FunctionCall(TypedDict):
    """
    A dictionary representing a function call.

    (deprecated in favor of tool_calls)
    """

    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 ToolCall(TypedDict):
    """
    A dictionary representing a tool call.
    """

    id: str
    """
    The ID of the tool call. This ID must be referenced when you submit a tool
    message with the results of the tool call.
    """
    function: FunctionCall # OpenAI documentation implies FunctionCall structure for 'function' field
    """
    The function that the tool calls.
    """
    type: Literal["function"]
    """
    The type of the tool. Currently, only `function` is supported.
    """


# ----------------------------------------------------------------------------
# Messages Content Parts
#
# The 'content' field in a message supports various content part types,
# including text, images, and audio.
# ----------------------------------------------------------------------------

class MessageContentImageURL(TypedDict):
    """
    An object specifying the URL of an image.
    """
    url: str
    """
    Either a URL of the image or a base64 encoded image data.
    """
    detail: NotRequired[Literal["auto", "low", "high"]]
    """
    Specifies the detail level of the image. `low` uses fewer tokens,
    `high` uses more tokens. `auto` lets the model decide. Defaults to `auto`.
    """

class MessageContentImagePart(TypedDict):
    """
    A dictionary within the 'content' key of a message, representing parameters
    referencing an image for multi-modal chat completion.

    Example :
    ```python
    {
        "type": "image_url",
        "image_url": {
            "url": "https://example.com/image.png",
            "detail": "auto"
        }
    }
    ```
    """

    image_url: MessageContentImageURL
    """
    An object containing the URL and an optional detail field.
    """
    type: Literal["image_url"]
    """
    The type of the content part. (Always "image_url" for this type)
    """


class MessageContentAudioPart(TypedDict):
    """
    A dictionary within the 'content' key of a message, representing parameters
    referencing an audio file for multi-modal chat completion.
    This is a conceptual representation; official OpenAI support for audio in
    message content parts for models like GPT-4o might evolve.
    As of the last update, direct audio file input in 'content' is generally
    handled by specific audio models or transcription services first.

    Example (Conceptual):
    ```python
    {
        "type": "audio_url", # Or a similar type if officially supported
        "audio_url": {
            "url": "data:audio/wav;base64,..." # Or a direct URL
        }
    }
    # Or, for models supporting direct audio input in messages:
    # {
    #     "type": "input_audio",
    #     "input_audio": {
    #         "data": "base64_encoded_audio_data",
    #         "format": "wav"
    #     }
    # }
    ```
    """
    # Note: The structure for audio content parts may vary based on specific model capabilities.
    # The provided 'input_audio' structure seems plausible but should be verified against
    # the latest documentation for the specific model being targeted.
    # For now, retaining a structure similar to image_url for consistency if a URL based approach is used.
    # If direct base64 audio input is common in message parts:
    input_audio: NotRequired[Dict[str, Any]] = TypedDict( # Using NotRequired if not universally supported
        "InputAudio",
        {
            "data": str, # Base64 encoded audio
            "format": NotRequired[Literal["wav", "mp3", "opus", "flac", "aac", "m4a"]], # Common audio formats
        },
    )
    audio_url: NotRequired[Dict[str, Any]] = TypedDict( # Alternative structure if URL based
        "AudioURL",
        {
            "url": str,
        }
    )
    type: Required[Literal["audio_url", "input_audio"]] # Type might depend on how audio is provided
    """
    The type of the content part.
    """


class MessageContentTextPart(TypedDict):
    """
    A dictionary within the 'content' key of a message, representing a text
    message.

    Example :
    ```python
    {
        "type": "text",
        "text": "Hello, world!"
    }
    """

    text: Required[str]
    """
    The text of the message.
    """
    type: Required[Literal["text"]]
    """
    The type of the content part. (Always "text" for this type)
    """


MessageContentPart: TypeAlias = Union[
    MessageContentImagePart,
    MessageContentTextPart,
    MessageContentAudioPart, # Added, assuming potential future or specialized use
]
"""
A type alias for all possible content parts of a message.
"""


MessageContent: TypeAlias = Union[
    str,
    Iterable[MessageContentPart],
]
"""
A type alias for the content of a message.

This specific type is generally valid for the `user` role.
Assistant messages might also contain complex content, but often primarily text or tool_calls.
System messages typically contain only strings.
"""


MessageTextContent: TypeAlias = Union[
    str,
    Iterable[MessageContentTextPart], # Technically, a string content is just `[{"type": "text", "text": "string"}]`
]
"""
A type alias for the content of a message, supporting either a string
or an iterable of `MessageContentTextPart` objects. Primarily for scenarios
where only text content is expected.
"""


# ----------------------------------------------------------------------------
# Message Types
# ----------------------------------------------------------------------------


MessageRole: TypeAlias = Literal[
    "system", "user", "assistant", "tool"
    # "developer" role is not standard in the public OpenAI API.
    # "function" role is deprecated in favor of "tool".
]
"""
A type alias for the role of a message.
"""


class Message(TypedDict):
    """
    A dictionary representing a message used to create a
    chat completion.

    ---

    ### System Messages:

    System messages set the behavior of the assistant.

    Example :
    ```python
    {
        "role": "system",
        "content": "You are a helpful assistant."
    }
    ```

    ### User Messages:

    User messages provide requests or comments for the assistant to respond to.
    They can contain text and image URLs.

    Example :
    ```python
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "What’s in this image?"},
            {"type": "image_url", "image_url": {"url": "[https://example.com/image.png](https://example.com/image.png)"}}
        ]
    }
    # Or simpler text-only:
    # {
    #     "role": "user",
    #     "content": "Hello, world!"
    # }
    ```

    ### Assistant Messages:

    Assistant messages store previous assistant responses. They can also include
    tool calls.

    Example :
    ```python
    {
        "role": "assistant",
        "content": "I can help with that.", # Content can be null if tool_calls are present
        "tool_calls": [
            {
                "id": "call_abc123",
                "type": "function",
                "function": {
                    "name": "get_current_weather",
                    "arguments": "{\\"location\\": \\"Boston, MA\\"}"
                }
            }
        ]
    }
    ```

    ### Tool Messages:

    Tool messages submit the results of tool calls made by the assistant.

    Example :
    ```python
    {
        "role": "tool",
        "content": "{\\"temperature\\": \\"22\\", \\"unit\\": \\"celsius\\"}", # Result of the function call
        "tool_call_id": "call_abc123"
    }
    ```
    """

    role: Required[MessageRole]
    """
    The role of the message author.
    (One of "system", "user", "assistant", or "tool")
    """
    content: Optional[MessageContent] # Content can be null for assistant messages with tool_calls
    """
    The content of the message. Can be a string or an array of content parts.
    Can be null for assistant messages if tool_calls are present.
    """
    name: NotRequired[str]
    """
    An optional name for the participant. Provides the model information to
    differentiate between participants of the same role.
    This is primarily for the `tool` role to identify the function called,
    or for `user` and `assistant` in specific multi-speaker scenarios.
    Deprecated for `function` role messages, use `tool_call_id` on `tool` messages instead.
    """
    tool_calls: NotRequired[List[ToolCall]] # Changed from Iterable[Dict[str, Any]] to List[ToolCall]
    """
    The tool calls generated by the model, such as function calls.
    Only for `assistant` role messages.
    """
    tool_call_id: NotRequired[str]
    """
    Required if role is `tool`. The ID of the tool call resolved by the tool message.
    """
    # function_call is deprecated.
    # function_call: NotRequired[Dict[str, Any]]


# ----------------------------------------------------------------------------
# Primary Response Types
# ----------------------------------------------------------------------------


# NOTE:
# all response types are in pydantic to follow standard python api
# schema
# all models within 'prompted' however, are subscriptable as they
# inherit from this class.
class Subscriptable(BaseModel):
    """
    A light wrapper over a Pydantic BaseModel, that allows for subscriptable
    access to fields in a model.
    """

    def __getitem__(self, key: str) -> Any:
        """
        Get an item from the model.
        """
        return getattr(self, key)

    def __setitem__(self, key: str, value: Any) -> None:
        """
        Set an item in the model.
        """
        setattr(self, key, value)

    def __contains__(self, key: str) -> bool:
        """
        Check if a key exists in the model.
        """
        # Check if the field is set or has a default value
        if key in self.model_fields_set:
            return True
        if key in self.model_fields and self.model_fields[key].default is not None:
            return True
        # For aliased fields, check if the alias is present
        for field_name, field_info in self.model_fields.items():
            if field_info.alias == key:
                if field_name in self.model_fields_set:
                    return True
                if field_info.default is not None:
                    return True
        return False


    def get(self, key: str, default: Any = None) -> Any:
        """
        Get an item from the model, with a default value if the key does not exist.
        """
        if key in self:
            return self[key]
        # Attempt to get by alias if direct key fails
        for field_name, field_info in self.model_fields.items():
            if field_info.alias == key and field_name in self:
                 return getattr(self, field_name)
        return default

    model_config = { # Pydantic v2 config
        "populate_by_name": True, # Allow using field names or aliases
        "extra": "allow" # Allow extra fields not defined in the model
    }


class TopLogprob(Subscriptable):
    """
    Represents the log probabilities of a token chosen as the top choice.
    """

    token: str
    """
    The token.
    """
    bytes: Optional[List[int]] = None
    """
    A list of integers representing the UTF-8 bytes representation of the token.
    Useful in instances where characters are represented by multiple tokens.
    """
    logprob: float
    """
    The log probability of this token.
    """


class TokenLogprob(Subscriptable):
    """
    Represents the logprobs of a specific token, including top alternatives.
    """

    token: str
    """
    The token.
    """
    bytes: Optional[List[int]] = None
    """
    A list of integers representing the UTF-8 bytes representation of the token.
    Useful in instances where characters are represented by multiple tokens.
    """
    logprob: float
    """
    The log probability of this token.
    """
    top_logprobs: List[TopLogprob]
    """
    List of the most likely tokens and their logprobs, at this token position.
    In rare cases, there may be fewer than `top_logprobs` returned.
    """


class ChoiceLogprobs(Subscriptable):
    """
    Log probability information for the choice.
    """

    content: Optional[List[TokenLogprob]] = None
    """
    A list of message content tokens with log probability information.
    """
    # 'refusal' field is not explicitly listed in the latest common API responses for ChoiceLogprobs.
    # It might be specific to certain models or older API versions.
    # If it's still relevant for specific use cases, it can be added back.
    # refusal: Optional[List[TokenLogprob]] = None


class CompletionFunction(Subscriptable): # Corresponds to FunctionCall in the response
    """
    The function that the model wants to call.
    (Deprecated in favor of CompletionToolCall)
    """

    name: str
    """
    The name of the function to be called.
    """
    arguments: str
    """
    The arguments to call the function with, as generated by the model in JSON
    format.
    """

class CompletionToolCallFunction(Subscriptable): # More specific than CompletionFunction
    """
    The function definition within a tool call in a completion response.
    """
    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.
    """

class CompletionToolCall(Subscriptable):
    """
    A tool call in a completion response or chunk.
    """

    id: str
    """
    The ID of the tool call.
    """
    type: Literal["function"]
    """
    The type of the tool. Currently, only `function` is supported.
    """
    function: CompletionToolCallFunction
    """
    The function that the tool calls.
    """
    index: Optional[int] = None # Present in streaming delta
    """
    The index of the tool call in the list of tool calls. (Streamed delta only)
    """


class CompletionMessage(Subscriptable):
    """
    A message object that forms part of the completion response.
    This usually represents the assistant's reply.
    """
    role: Literal["assistant"]
    """
    The role of the message author. Always `assistant` for completion messages.
    """
    content: Optional[MessageContent] # Content can be null if tool_calls is present
    """
    The content of the message. Can be a string or an array of content parts.
    Can be null if tool_calls are present.
    """
    # name: Optional[str] = None # 'name' is not typically part of the assistant's response message.
    function_call: Optional[CompletionFunction] = None # Deprecated
    """
    (Deprecated) The name and arguments of a function that should be called, as generated by the model.
    Replaced by `tool_calls`.
    """
    tool_calls: Optional[List[CompletionToolCall]] = None
    """
    The tool calls generated by the model, such as function calls.
    """
    # tool_call_id is not part of the assistant's response message itself, but rather for tool messages.

class CompletionUsage(Subscriptable):
    """
    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 Completion(Subscriptable):
    """
    A pydantic model representing a chat completion
    response.
    """

    class Choice(Subscriptable):
        """
        A completion choice.
        """

        message: CompletionMessage
        """
        A chat completion message generated by the model.
        """
        finish_reason: Literal["stop", "length", "tool_calls", "content_filter", "function_call"] # "function_call" is deprecated
        """
        The reason the model stopped generating tokens.
        - `stop`: API returned complete message, or a message terminated by a stop sequence.
        - `length`: Incomplete model output due to `max_tokens` parameter or token limit.
        - `tool_calls`: Model called a tool.
        - `content_filter`: Omitted content due to a flag from our content filters.
        - `function_call`: (Deprecated) Model called a function.
        """
        index: int
        """
        The index of this choice in the list of choices.
        """
        logprobs: Optional[ChoiceLogprobs] = None
        """
        Log probability information for the choice.
        """

    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"] # OpenAI SDK uses 'chat.completion'
    """
    The object type, which is always `chat.completion`.
    """
    system_fingerprint: Optional[str] = None
    """
    This fingerprint represents the backend configuration that the model runs with.
    You can use this value to track changes in the backend before comparing output
    from different API calls.
    """
    usage: Optional[CompletionUsage] = None
    """
    Usage statistics for the completion request.
    """
    # service_tier: Optional[Literal["scale", "default"]] = None # Not consistently in latest public API docs for this object

# ----------------------------------------------------------------------------
# Streaming
# ----------------------------------------------------------------------------


class CompletionChunk(Subscriptable):
    """
    A pydantic model representing a chat completion stream chunk.
    """

    class Choice(Subscriptable):
        """
        A choice in a completion chunk.
        """
        class Delta(Subscriptable): # Renamed from CompletionMessage to Delta for clarity in chunks
            """
            The delta message content for a streaming choice.
            """
            role: Optional[Literal["system", "user", "assistant", "tool"]] = None # Role appears in first chunk
            """
            The role of the author of this message.
            """
            content: Optional[str] = None # Content is streamed token by token
            """
            The contents of the chunk message.
            """
            tool_calls: Optional[List[CompletionToolCall]] = None # Tool calls can also be streamed
            """
            Tool calls generated by the model. Can appear incrementally.
            Each tool call in the list can have an `index` field.
            """
            function_call: Optional[CompletionFunction] = None # Deprecated
            """
            (Deprecated) The name and arguments of a function that should be called.
            """

        delta: Delta
        """
        A chat completion delta generated by streamed model responses.
        """
        finish_reason: Optional[Literal["stop", "length", "tool_calls", "content_filter", "function_call"]] = None # "function_call" is deprecated
        """
        The reason the model stopped generating tokens.
        Present in the final chunk of a choice.
        """
        index: int
        """
        The index of this choice in the stream.
        """
        logprobs: Optional[ChoiceLogprobs] = None # Logprobs can also be part of chunks
        """
        Log probability information for the choice.
        """

    id: str
    """
    A unique identifier for the chat completion chunk.
    """
    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 chunk was created.
    """
    model: str
    """
    The model to generate the completion.
    """
    object: Literal["chat.completion.chunk"] # OpenAI SDK uses 'chat.completion.chunk'
    """
    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.
    You can use this value to track changes in the backend before comparing output
    from different API calls.
    """
    usage: Optional[CompletionUsage] = None # Usage is often null in chunks until the very end, or not present at all.
    """
    An optional field that appears when the model is `gpt-3.5-turbo-0125` or `gpt-4-turbo-preview`
    and the `stream_options` parameter is set.
    It contains token usage statistics for the entire request, summed across all chunks.
    """
    # service_tier: Optional[Literal["scale", "default"]] = None # Not consistently in latest public API docs for this object

# ----------------------------------------------------------------------------
# Embeddings
# ----------------------------------------------------------------------------

class EmbeddingsUsage(Subscriptable):
    """
    Usage statistics for the embeddings request.
    """
    prompt_tokens: int
    """
    The number of tokens in the input.
    """
    total_tokens: int
    """
    The total number of tokens used by the request.
    """

class Embedding(Subscriptable): # This 'Embedding' refers to a single embedding object in the response list
    """
    Represents an embedding vector returned by the embedding endpoint.
    """

    embedding: List[float]
    """
    The embedding vector, which is a list of floats. The length of vector depends
    on the model.
    """
    index: int
    """
    The index of the embedding in the list of embeddings.
    """
    object: Literal["embedding"]
    """
    The object type, which is always `embedding`.
    """

# The top-level response for an embedding request
class EmbeddingResponse(Subscriptable): # Added a top-level response type for embeddings
    """
    The response from an OpenAI API embedding endpoint.
    """
    data: List[Embedding]
    """
    A list of embedding objects.
    """
    model: str
    """
    The model ID used to generate the embeddings.
    """
    object: Literal["list"] # The top-level object is often "list" containing embeddings
    """
    The object type, typically "list".
    """
    usage: EmbeddingsUsage
    """
    Usage statistics for the request.
    """