from typing import Any, Dict, List, Optional

from zai.core import BaseModel


class ChoiceDeltaFunctionCall(BaseModel):
	"""
	Function call delta information in streaming response

	Attributes:
		arguments: Function call arguments
		name: Function name
	"""

	arguments: Optional[str] = None
	name: Optional[str] = None


class ChoiceDeltaToolCallFunction(BaseModel):
	"""
	Tool call function delta information in streaming response

	Attributes:
		arguments: Function call arguments
		name: Function name
	"""

	arguments: Optional[str] = None
	name: Optional[str] = None


class ChoiceDeltaToolCall(BaseModel):
	"""
	Tool call delta information in streaming response

	Attributes:
		index: Index of the tool call
		id: Unique identifier for the tool call
		function: Function call information
		type: Type of the tool call
	"""

	index: int
	id: Optional[str] = None
	function: Optional[ChoiceDeltaToolCallFunction] = None
	type: Optional[str] = None


class AudioCompletionChunk(BaseModel):
	"""
	Audio completion chunk information

	Attributes:
		id: Unique identifier for the audio chunk
		data: Audio data content
		expires_at: Timestamp when the audio expires
	"""

	id: Optional[str] = None
	data: Optional[str] = None
	expires_at: Optional[int] = None


class ChoiceDelta(BaseModel):
	"""
	Delta information for streaming chat completion choice

	Attributes:
		content: Content delta
		role: Role of the message sender
		reasoning_content: Reasoning content delta, it's recommended to return the model's reasoning_content in next request to achieve better results in multi-turn conversations.
		tool_calls: List of tool call deltas
		audio: Audio completion chunk
	"""

	content: Optional[str] = None
	role: Optional[str] = None
	reasoning_content: Optional[str] = None
	tool_calls: Optional[List[ChoiceDeltaToolCall]] = None
	audio: Optional[AudioCompletionChunk] = None


class Choice(BaseModel):
	"""
	Choice information in streaming chat completion

	Attributes:
		delta: Delta information for the choice
		finish_reason: Reason why the completion finished
		index: Index of the choice
	"""

	delta: ChoiceDelta
	finish_reason: Optional[str] = None
	index: int

class PromptTokensDetails(BaseModel):
	"""
	Detailed breakdown of token usage for the input prompt

	Attributes:
		cached_tokens: Number of tokens reused from cache
	"""

	cached_tokens: int


class CompletionTokensDetails(BaseModel):
	"""
	Detailed breakdown of token usage for the model completion

	Attributes:
	    reasoning_tokens: Number of tokens used for reasoning steps
	"""

	reasoning_tokens: int

class CompletionUsage(BaseModel):
	"""
	Token usage information for completion

	Attributes:
		prompt_tokens: Number of tokens in the prompt
		completion_tokens: Number of tokens in the completion
		total_tokens: Total number of tokens used
	"""

	prompt_tokens: int
	prompt_tokens_details: Optional[PromptTokensDetails] = None
	completion_tokens: int
	completion_tokens_details: Optional[CompletionTokensDetails] = None
	total_tokens: int


class ChatCompletionChunk(BaseModel):
	"""
	Streaming chat completion chunk response

	Attributes:
		id: Unique identifier for the completion
		choices: List of completion choices
		created: Timestamp when the completion was created
		model: Model used for the completion
		usage: Token usage information
		extra_json: Additional JSON data
	"""

	id: Optional[str] = None
	choices: List[Choice]
	created: Optional[int] = None
	model: Optional[str] = None
	usage: Optional[CompletionUsage] = None
	extra_json: Dict[str, Any]
