"""
FluxTokens API Types
"""

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


# Models
Model = Literal["gpt-4.1-mini", "gpt-4.1-nano", "gemini-2.5-flash"]


class ModelInfo(BaseModel):
    """Information about an available model."""
    
    id: str
    name: str
    provider: Literal["OpenAI", "Google"]
    input_price: float  # USD per 1M tokens
    output_price: float  # USD per 1M tokens
    max_tokens: int
    supports_vision: bool
    supports_audio: bool
    supports_video: bool


MODELS: Dict[str, ModelInfo] = {
    "gpt-4.1-mini": ModelInfo(
        id="gpt-4.1-mini",
        name="GPT-4.1 Mini",
        provider="OpenAI",
        input_price=0.28,
        output_price=1.12,
        max_tokens=16384,
        supports_vision=True,
        supports_audio=False,
        supports_video=False,
    ),
    "gpt-4.1-nano": ModelInfo(
        id="gpt-4.1-nano",
        name="GPT-4.1 Nano",
        provider="OpenAI",
        input_price=0.07,
        output_price=0.28,
        max_tokens=16384,
        supports_vision=True,
        supports_audio=False,
        supports_video=False,
    ),
    "gemini-2.5-flash": ModelInfo(
        id="gemini-2.5-flash",
        name="Gemini 2.5 Flash",
        provider="Google",
        input_price=0.21,
        output_price=1.75,
        max_tokens=65536,
        supports_vision=True,
        supports_audio=True,
        supports_video=True,
    ),
}


# Message Content Types
class TextContent(BaseModel):
    """Text content in a message."""
    
    type: Literal["text"] = "text"
    text: str


class ImageURL(BaseModel):
    """Image URL details."""
    
    url: str
    detail: Optional[Literal["auto", "low", "high"]] = None


class ImageURLContent(BaseModel):
    """Image URL content in a message."""
    
    type: Literal["image_url"] = "image_url"
    image_url: ImageURL


class InputAudio(BaseModel):
    """Audio input details."""
    
    data: str  # Base64 encoded
    format: Literal["wav", "mp3", "ogg", "flac"]


class AudioContent(BaseModel):
    """Audio content in a message."""
    
    type: Literal["input_audio"] = "input_audio"
    input_audio: InputAudio


class VideoURL(BaseModel):
    """Video URL details."""
    
    url: str


class VideoContent(BaseModel):
    """Video content in a message."""
    
    type: Literal["video_url"] = "video_url"
    video_url: VideoURL


MessageContent = Union[str, List[Union[TextContent, ImageURLContent, AudioContent, VideoContent]]]


# Tool Types
class FunctionDefinition(BaseModel):
    """Function definition for tools."""
    
    name: str
    description: Optional[str] = None
    parameters: Optional[Dict[str, Any]] = None


class Tool(BaseModel):
    """Tool definition."""
    
    type: Literal["function"] = "function"
    function: FunctionDefinition


class FunctionCall(BaseModel):
    """Function call in a tool call."""
    
    name: str
    arguments: str


class ToolCall(BaseModel):
    """Tool call from the model."""
    
    id: str
    type: Literal["function"] = "function"
    function: FunctionCall


# Messages
class ChatMessage(BaseModel):
    """A message in the conversation."""
    
    role: Literal["system", "user", "assistant", "tool"]
    content: Optional[MessageContent] = None
    name: Optional[str] = None
    tool_call_id: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None


# Request
class ChatCompletionRequest(BaseModel):
    """Request for chat completion."""
    
    model: str
    messages: List[ChatMessage]
    temperature: Optional[float] = Field(default=None, ge=0, le=2)
    max_tokens: Optional[int] = Field(default=None, gt=0)
    top_p: Optional[float] = Field(default=None, ge=0, le=1)
    top_k: Optional[int] = Field(default=None, gt=0)
    frequency_penalty: Optional[float] = Field(default=None, ge=-2, le=2)
    presence_penalty: Optional[float] = Field(default=None, ge=-2, le=2)
    stop: Optional[Union[str, List[str]]] = None
    stream: Optional[bool] = None
    user: Optional[str] = None
    tools: Optional[List[Tool]] = None
    tool_choice: Optional[Union[str, Dict[str, Any]]] = None


# Response
class AssistantMessage(BaseModel):
    """Assistant message in response."""
    
    role: Literal["assistant"] = "assistant"
    content: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None


class Choice(BaseModel):
    """A completion choice."""
    
    index: int
    message: AssistantMessage
    finish_reason: Optional[Literal["stop", "length", "tool_calls", "content_filter"]] = None


class PromptTokensDetails(BaseModel):
    """Details about prompt tokens."""
    
    cached_tokens: Optional[int] = None
    audio_tokens: Optional[int] = None


class CompletionTokensDetails(BaseModel):
    """Details about completion tokens."""
    
    reasoning_tokens: Optional[int] = None
    audio_tokens: Optional[int] = None
    accepted_prediction_tokens: Optional[int] = None
    rejected_prediction_tokens: Optional[int] = None


class Usage(BaseModel):
    """Token usage information."""
    
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    prompt_tokens_details: Optional[PromptTokensDetails] = None
    completion_tokens_details: Optional[CompletionTokensDetails] = None


class ChatCompletionResponse(BaseModel):
    """Response from chat completion."""
    
    id: str
    object: Literal["chat.completion"] = "chat.completion"
    created: int
    model: str
    choices: List[Choice]
    usage: Usage
    system_fingerprint: Optional[str] = None


# Streaming
class DeltaMessage(BaseModel):
    """Delta message in streaming response."""
    
    role: Optional[Literal["assistant"]] = None
    content: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None


class StreamChoice(BaseModel):
    """A streaming choice."""
    
    index: int
    delta: DeltaMessage
    finish_reason: Optional[Literal["stop", "length", "tool_calls", "content_filter"]] = None


class ChatCompletionChunk(BaseModel):
    """A chunk in streaming response."""
    
    id: str
    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
    created: int
    model: str
    choices: List[StreamChoice]
    usage: Optional[Usage] = None
