from __future__ import annotations

from typing import List, Optional

from zai.core import BaseModel


class GeneratedImage(BaseModel):
	"""
	Generated image data

	Attributes:
		b64_json (Optional[str]): Base64 encoded image data
		url (Optional[str]): URL of the generated image
		revised_prompt (Optional[str]): Revised prompt used for generation
	"""

	b64_json: Optional[str] = None
	url: Optional[str] = None
	revised_prompt: Optional[str] = None


class ImagesResponded(BaseModel):
	"""
	Image generation response

	Attributes:
		created (int): Creation timestamp
		data (List[GeneratedImage]): List of generated images
	"""

	created: int
	data: List[GeneratedImage]


class AsyncImagesResponded(BaseModel):
	"""
	Async image generation response

	Attributes:
		id (Optional[str]): The task order number generated by the Z.ai open platform.
							Please use this order number when calling the request result interface.
		model (str): Model name
		request_id (str): The task number submitted by the user when requesting on the client
							or the task number generated by the platform.
		task_status (str): Processing status, PROCESSING (processing), SUCCESS (successful), FAIL (failed).
		image_result (Optional[List[GeneratedImage]]): Image generation result, available when task_status is SUCCESS
	"""

	id: Optional[str] = None
	model: str
	request_id: str
	task_status: str
	image_result: Optional[List[GeneratedImage]] = None
