"""Provider-agnostic types for multi-LLM support."""

from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional, Union

# Import OpenAI content part types
from openai.types.chat import ChatCompletionContentPartImageParam as ImageContentPart
from openai.types.chat import (
    ChatCompletionContentPartInputAudioParam as AudioContentPart,
)
from openai.types.chat import ChatCompletionContentPartTextParam as TextContentPart
from openai.types.chat import ChatCompletionMessageParam, ChatCompletionMessageToolCall
from pydantic import BaseModel, Field


# Streaming Tool Call Types (for handling partial data)
# ------------------------------------------------------------
class ToolCallDeltaFunction(BaseModel):
    """Function details in a streaming tool call delta."""

    name: Optional[str] = Field(None, description="The name of the function")
    arguments: Optional[str] = Field(
        None, description="The arguments for the function as JSON string"
    )


class ToolCallDelta(BaseModel):
    """Represents a partial tool call from streaming responses.

    This type handles incremental tool call data during streaming,
    where fields may be incomplete or arrive in chunks.
    """

    id: Optional[str] = Field(None, description="The ID identifying a tool call")
    type: Optional[str] = Field("function", description="The type of tool call")
    function: ToolCallDeltaFunction = Field(default_factory=ToolCallDeltaFunction)
    index: int = Field(0, description="The index of the tool call in the message")

    def to_completion_tool_call(self) -> ChatCompletionMessageToolCall:
        """Convert to a complete tool call for non-streaming contexts.

        At this stage all fields should be populated.
        """
        from openai.types.chat.chat_completion_message_tool_call import Function

        return ChatCompletionMessageToolCall(
            id=self.id or f"call_{id(self)}",
            type="function",
            function=Function(
                name=self.function.name or "unknown",
                arguments=self.function.arguments or "{}",
            ),
        )


class AudioResponse(BaseModel):
    """Audio response from the model."""

    id: str
    expires_at: datetime
    data: str  # base64 encoded audio bytes
    transcript: str


# OpenAIMessageContentPart - union of all message content parts
OpenAIMessageContentPart = Union[
    TextContentPart,
    ImageContentPart,
    AudioContentPart,
]


class ProviderType(str, Enum):
    """Supported provider types."""

    OPENAI = "openai"
    LITELLM = "litellm"
    ANTHROPIC = "anthropic"
    GOOGLE = "google"
    COHERE = "cohere"


class ProviderConfig(BaseModel):
    """Configuration for a provider."""

    provider_type: ProviderType
    api_key: Optional[str] = None
    api_base: Optional[str] = None
    api_version: Optional[str] = None
    organization: Optional[str] = None
    project: Optional[str] = None
    default_model: Optional[str] = None
    timeout: int = 60
    max_retries: int = 3

    class Config:
        extra = "allow"


class CompletionRequest(BaseModel):
    """Provider-agnostic completion request."""

    messages: List[ChatCompletionMessageParam]
    model: List[str]
    tools: Optional[List[Dict[str, Any]]] = None
    tool_choice: Optional[Union[str, Dict[str, Any]]] = None
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    top_p: Optional[float] = None
    frequency_penalty: Optional[float] = None
    presence_penalty: Optional[float] = None
    stop: Optional[List[str]] = None
    stream: bool = False
    response_format: Optional[Dict[str, Any]] = None
    seed: Optional[int] = None
    parallel_tool_calls: Optional[bool] = None
    include_usage: Optional[bool] = None
    timeout: Optional[int] = None
    max_retries: Optional[int] = None

    class Config:
        extra = "allow"


#   Did NOT bring (should use OpenAI SDK instead):
#     - OpenAITextMessagePart - Use openai.types.chat.ChatCompletionContentPartText
#     - OpenAIImageMessagePart - Use openai.types.chat.ChatCompletionContentPartImage
#     - OpenAIUserMessage - Use openai.types.chat.ChatCompletionUserMessageParam
#     - OpenAIAssistantMessage - Use openai.types.chat.ChatCompletionAssistantMessageParam
#     - OpenAIToolMessage - Use openai.types.chat.ChatCompletionToolMessageParam
#     - OpenAIMessage - Union of the above
#     - OpenAIAssistantMessageResponse - Use openai.types.chat.ChatCompletionMessage
