from pydantic import BaseModel, Field
from typing import Optional, List, Literal, Union
from enum import Enum
from pydantic import field_validator

"""
All these models are made based on the OpenAI openapi file to make the image generation server compatible with the OpenAI client

OpenAPI Spec: https://app.stainless.com/api/spec/documented/openai/openapi.documented.yml
"""

class ImageModel(str, Enum):
    """I plan to have all of these models optimized for the Aquiles-Image runtime"""
    # Stable Diffusion models
    SD3_MEDIUM = 'stabilityai/stable-diffusion-3-medium'
    SD3_5_LARGE = 'stabilityai/stable-diffusion-3.5-large'
    SD3_5_LARGE_TURBO = 'stabilityai/stable-diffusion-3.5-large-turbo'
    SD3_5_MEDIUM = 'stabilityai/stable-diffusion-3.5-medium'
    
    # FLUX models
    FLUX_1_DEV = 'black-forest-labs/FLUX.1-dev'
    FLUX_1_SCHNELL = 'black-forest-labs/FLUX.1-schnell'
    FLUX_1_KONTEXT_DEV = 'black-forest-labs/FLUX.1-Kontext-dev'
    FLUX_1_KREA_DEV = 'black-forest-labs/FLUX.1-Krea-dev'

    #FLUX2
    FLUX_2_4BNB = 'diffusers/FLUX.2-dev-bnb-4bit'
    FLUX_2 = 'black-forest-labs/FLUX.2-dev'
    
    #Z-Image-Turbo
    Z_IMAGE_TURBO = 'Tongyi-MAI/Z-Image-Turbo'

    QWEN_IMAGE = 'Qwen/Qwen-Image'

    QWEN_IMAGE_2512 = 'Qwen/Qwen-Image-2512'

    QWEN_IMAGE_EDIT_BASE = 'Qwen/Qwen-Image-Edit'

    QWEN_IMAGE_EDIT_2511 = 'Qwen/Qwen-Image-Edit-2511'

    QWEN_IMAGE_EDIT_2509 = 'Qwen/Qwen-Image-Edit-2509'

class VideoModels(str, Enum):
    """I plan to have all of these models optimized for the Aquiles-Image runtime"""
    WAN2_2_API = 'wan2.2'

    WAN2_2_TURBO = 'wan2.2-turbo'
    # HunyuanVideo-1.5
    HY1_5_480_API = 'hunyuanVideo-1.5-480p'

    HY1_5_720_API = 'hunyuanVideo-1.5-720p'

    HY1_5_480_API_FP8 = 'hunyuanVideo-1.5-480p-fp8'

    HY1_5_720_API_FP8 = 'hunyuanVideo-1.5-720p-fp8'

    HY1_5_480_API_TURBO = 'hunyuanVideo-1.5-480p-turbo'

    HY1_5_480_API_TURBO_FP8 = 'hunyuanVideo-1.5-480p-turbo-fp8'

    WAN_2_1 = 'wan2.1'

    WAN_2_1_TURBO = 'wan2.1-turbo'

    WAN_2_1_3B = 'wan2.1-3B'

    WAN_2_1_TURBO_FP8 = 'wan2.1-turbo-fp8'

class ResponseFormat(str, Enum):
    URL = "url"
    B64_JSON = "b64_json"


class OutputFormat(str, Enum):
    PNG = "png"
    JPEG = "jpeg"


class BackgroundType(str, Enum):
    TRANSPARENT = "transparent"
    OPAQUE = "opaque"
    AUTO = "auto"


class QualityType(str, Enum):
    STANDARD = "standard"
    HD = "hd"
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    AUTO = "auto"


class StyleType(str, Enum):
    VIVID = "vivid"
    NATURAL = "natural"


class ImageSize(str, Enum):
    SIZE_256x256 = "256x256"
    SIZE_512x512 = "512x512"
    SIZE_1024x1024 = "1024x1024"
    SIZE_1536x1024 = "1536x1024"
    SIZE_1024x1536 = "1024x1536"
    SIZE_1792x1024 = "1792x1024"
    SIZE_1024x1792 = "1024x1792"
    AUTO = "auto"


class InputFidelity(str, Enum):
    HIGH = "high"
    LOW = "low"


class ModerationLevel(str, Enum):
    LOW = "low"
    AUTO = "auto"


class ImageGenInputUsageDetails(BaseModel):
    text_tokens: int = Field(..., description="The number of text tokens in the input prompt.")
    image_tokens: int = Field(..., description="The number of image tokens in the input prompt.")


class ImageGenUsage(BaseModel):
    input_tokens: int = Field(..., description="The number of tokens (images and text) in the input prompt.")
    total_tokens: int = Field(..., description="The total number of tokens (images and text) used for the image generation.")
    output_tokens: int = Field(..., description="The number of output tokens generated by the model.")
    input_tokens_details: ImageGenInputUsageDetails = Field(..., description="The input tokens detailed information for the image generation.")


class Image(BaseModel):
    b64_json: Optional[str] = Field(None, description="The base64-encoded JSON of the generated image. Default value for `gpt-image-1`, and only present if `response_format` is set to `b64_json` for `dall-e-2` and `dall-e-3`.")
    url: Optional[str] = Field(None, description="When using `dall-e-2` or `dall-e-3`, the URL of the generated image if `response_format` is set to `url` (default value). Unsupported for `gpt-image-1`.")
    revised_prompt: Optional[str] = Field(None, description="For `dall-e-3` only, the revised prompt that was used to generate the image.")


class ImagesResponse(BaseModel):
    created: int = Field(..., description="The Unix timestamp (in seconds) of when the image was created.")
    data: List[Image] = Field(..., description="The list of generated images.")
    background: Optional[Literal["transparent", "opaque"]] = Field(None, description="The background parameter used for the image generation. Either `transparent` or `opaque`.")
    output_format: Optional[Literal["png", "webp", "jpeg"]] = Field(None, description="The output format of the image generation. Either `png`, `webp`, or `jpeg`.")
    size: Optional[Literal["1024x1024", "1536x1024", "1024x1536", "256x256", "512x512", "1792x1024", "1024x1792"]] = Field(None, description="The size of the image generated. Either `1024x1024`, `1024x1536`, or `1536x1024`.")
    quality: Optional[Literal["low", "medium", "high"]] = Field(None, description="The quality of the image generated. Either `low`, `medium`, or `high`.")
    usage: Optional[ImageGenUsage] = Field(None, description="For `gpt-image-1` only, the token usage information for the image generation.")

    @field_validator('background', mode='before')
    @classmethod
    def validate_background(cls, v):
        return 'transparent' if v == 'auto' else v

    @field_validator('quality', mode='before')
    @classmethod
    def validate_quality(cls, v):
        return 'medium' if v == 'auto' else v


class ImageGenPartialImageEvent(BaseModel):
    type: Literal["image_generation.partial_image"] = Field(..., description="The type of the event. Always `image_generation.partial_image`.")
    b64_json: str = Field(..., description="Base64-encoded partial image data, suitable for rendering as an image.")
    created_at: int = Field(..., description="The Unix timestamp when the event was created.")
    size: Literal["1024x1024", "1024x1536", "1536x1024", "auto"] = Field(..., description="The size of the requested image.")
    quality: Literal["low", "medium", "high", "auto"] = Field(..., description="The quality setting for the requested image.")
    background: Literal["transparent", "opaque", "auto"] = Field(..., description="The background setting for the requested image.")
    output_format: Literal["png", "webp", "jpeg"] = Field(..., description="The output format for the requested image.")
    partial_image_index: int = Field(..., description="0-based index for the partial image (streaming).")


class ImagesUsage(BaseModel):
    total_tokens: int = Field(..., description="The total number of tokens (images and text) used for the image generation.")
    input_tokens: int = Field(..., description="The number of tokens (images and text) in the input prompt.")
    output_tokens: int = Field(..., description="The number of image tokens in the output image.")
    input_tokens_details: ImageGenInputUsageDetails = Field(..., description="The input tokens detailed information for the image generation.")


class ImageGenCompletedEvent(BaseModel):
    type: Literal["image_generation.completed"] = Field(..., description="The type of the event. Always `image_generation.completed`.")
    b64_json: str = Field(..., description="Base64-encoded image data, suitable for rendering as an image.")
    created_at: int = Field(..., description="The Unix timestamp when the event was created.")
    size: Literal["1024x1024", "1024x1536", "1536x1024", "auto"] = Field(..., description="The size of the generated image.")
    quality: Literal["low", "medium", "high", "auto"] = Field(..., description="The quality setting for the generated image.")
    background: Literal["transparent", "opaque", "auto"] = Field(..., description="The background setting for the generated image.")
    output_format: Literal["png", "webp", "jpeg"] = Field(..., description="The output format for the generated image.")
    usage: ImagesUsage = Field(..., description="For `gpt-image-1` only, the token usage information for the image generation.")


class ImageEditPartialImageEvent(BaseModel):
    type: Literal["image_edit.partial_image"] = Field(..., description="The type of the event. Always `image_edit.partial_image`.")
    b64_json: str = Field(..., description="Base64-encoded partial image data, suitable for rendering as an image.")
    created_at: int = Field(..., description="The Unix timestamp when the event was created.")
    size: Literal["1024x1024", "1024x1536", "1536x1024", "auto"] = Field(..., description="The size of the requested edited image.")
    quality: Literal["low", "medium", "high", "auto"] = Field(..., description="The quality setting for the requested edited image.")
    background: Literal["transparent", "opaque", "auto"] = Field(..., description="The background setting for the requested edited image.")
    output_format: Literal["png", "webp", "jpeg"] = Field(..., description="The output format for the requested edited image.")
    partial_image_index: int = Field(..., description="0-based index for the partial image (streaming).")


class ImageEditCompletedEvent(BaseModel):
    type: Literal["image_edit.completed"] = Field(..., description="The type of the event. Always `image_edit.completed`.")
    b64_json: str = Field(..., description="Base64-encoded final edited image data, suitable for rendering as an image.")
    created_at: int = Field(..., description="The Unix timestamp when the event was created.")
    size: Literal["1024x1024", "1024x1536", "1536x1024", "auto"] = Field(..., description="The size of the edited image.")
    quality: Literal["low", "medium", "high", "auto"] = Field(..., description="The quality setting for the edited image.")
    background: Literal["transparent", "opaque", "auto"] = Field(..., description="The background setting for the edited image.")
    output_format: Literal["png", "webp", "jpeg"] = Field(..., description="The output format for the edited image.")
    usage: ImagesUsage = Field(..., description="For `gpt-image-1` only, the token usage information for the image generation.")


ImageGenStreamEvent = Union[ImageGenPartialImageEvent, ImageGenCompletedEvent]
ImageEditStreamEvent = Union[ImageEditPartialImageEvent, ImageEditCompletedEvent]



class CreateImageRequest(BaseModel):
    prompt: str = Field(..., max_length=1000, description="A text description of the desired image(s).")
    model: Optional[Union[str, ImageModel]] = Field(None, description="The model to use for image generation.")
    n: Optional[int] = Field(1, ge=1, le=10, description="The number of images to generate. Must be between 1 and 10. For `dall-e-3`, only `n=1` is supported.")
    quality: Optional[Literal["standard", "hd", "low", "medium", "high", "auto"]] = Field("auto", description="The quality of the image that will be generated.")
    response_format: Optional[Literal["url", "b64_json"]] = Field("url", description="The format in which generated images with `dall-e-2` and `dall-e-3` are returned. Must be one of `url` or `b64_json`. URLs are only valid for 60 minutes after the image has been generated.")
    output_format: Optional[Literal["png", "jpeg", "webp"]] = Field("png", description="The format in which the generated images are returned. Must be one of `png` or `jpeg`.")
    output_compression: Optional[int] = Field(100, description="The compression level (0-100%) for the generated images.")
    stream: Optional[bool] = Field(False, description="Generate the image in streaming mode. Defaults to `false`. See the Image generation guide for more information. ")
    partial_images: Optional[int] = Field(0, ge=0, le=3, description="The number of partial images to generate. This parameter is used for streaming responses that return partial images. Value must be between 0 and 3. When set to 0, the response will be a single image sent in one streaming event.")
    size: Optional[Literal["auto", "1024x1024", "1536x1024", "1024x1536", "256x256", "512x512", "1792x1024", "1024x1792"]] = Field("auto", description="The size of the generated images.")
    moderation: Optional[Literal["low", "auto"]] = Field("auto", description="Control the content-moderation level for images generated by `gpt-image-1`. Must be either `low` for less restrictive filtering or `auto` (default value).")
    background: Optional[Literal["transparent", "opaque", "auto"]] = Field("auto", description="Allows to set transparency for the background of the generated image(s). Must be one of `transparent`, `opaque` or `auto` (default value). When `auto` is used, the model will automatically determine the best background for the image.")
    style: Optional[Literal["vivid", "natural"]] = Field("vivid", description="The style of the generated images. ")
    user: Optional[str] = Field(None, description="A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.")


class CreateImageEditRequest(BaseModel):
    # Note: In FastAPI, these fields would be handled with UploadFile and Form()
    # image: Union[bytes, List[bytes]] - Would be handled with UploadFile on the endpoint
    prompt: str = Field(..., max_length=1000, description="A text description of the desired image(s).")
    # mask: Optional[bytes] - It would be handled with UploadFile on the endpoint
    background: Optional[Literal["transparent", "opaque", "auto"]] = Field("auto", description="Allows to set transparency for the background of the generated image(s).")
    model: Optional[Union[str, ImageModel]] = Field(None, description="The model to use for image generation.")
    n: Optional[int] = Field(1, ge=1, le=10, description="The number of images to generate. Must be between 1 and 10.")
    size: Optional[Literal["256x256", "512x512", "1024x1024", "1536x1024", "1024x1536", "1792x1024", "1024x1792", "auto"]] = Field("1024x1024", description="The size of the generated images. Must be one of `1024x1024`, `1536x1024` (landscape), `1024x1536` (portrait), or `auto` (default value)")
    response_format: Optional[Literal["url", "b64_json"]] = Field("url", description="The format in which the generated images are returned. Must be one of `url` or `b64_json`. URLs are only valid for 60 minutes after the image has been generated. ")
    output_format: Optional[Literal["png", "jpeg", "webp"]] = Field("png", description="The format in which the generated images are returned.")
    output_compression: Optional[int] = Field(100, description="The compression level (0-100%) for the generated images. ")
    user: Optional[str] = Field(None, description="A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.")
    input_fidelity: Optional[Literal["high", "low"]] = Field("low", description="Control how much effort the model will exert to match the style and features, especially facial features, of input images.")
    stream: Optional[bool] = Field(False, description="Edit the image in streaming mode. Defaults to `false`. See the Image generation guide for more information.")
    partial_images: Optional[int] = Field(0, ge=0, le=3, description="The number of partial images to generate. This parameter is used for streaming responses that return partial images. Value must be between 0 and 3.")
    quality: Optional[Literal["standard", "low", "medium", "high", "auto"]] = Field("auto", description="The quality of the image that will be generated.")

class ConfigsServe(BaseModel):
    model: Union[ImageModel, str] = Field(default="stabilityai/stable-diffusion-3.5-medium", description="The model to use for image generation.")
    allows_api_keys: List[str] | None = Field( default_factory=lambda: [""], description="API KEYS allowed to make requests")
    max_concurrent_infer: int | None = Field(default=10, description="Maximum number of inferences running at the same time")
    block_request: bool | None = Field(default=False, description="Block requests during the maximum concurrent inferences")
    load_model: bool | None = Field(default=True, description="Loading the model or not simply allows for faster development without having to load the model constantly.")
    steps_n: int | None = Field(default=None, description="Loading the model or not simply allows for faster development without having to load the model constantly.")
    auto_pipeline: bool | None = Field(default=None, description="Load a model that is compatible with diffusers (Only the Text2Image models) but is not mentioned in the Aquiles-Image documentation.")
    device_map: str | None = Field(default=None, description="device_map option in which to load the model (Only compatible with diffusers/FLUX.2-dev-bnb-4bit)")
    type_model: str | None = Field(default=None, description="This is for video models. There are only 2 options: 'Images' and 'Videos'.")
    dist_inference: bool | None = Field(default=False, description="Use distributed inference (Not yet implemented)")
    max_batch_size: int | None = Field(default=4, description="Maximum number of requests to group in a single batch for inference")
    batch_timeout: float | None = Field(default=0.5, description="Maximum time (in seconds) to wait before processing a batch even if not full")
    worker_sleep: float | None = Field(default=0.05, description="Time (in seconds) the worker sleeps between checking for new batch requests")

class Model(BaseModel):
    id: str = Field(..., description="Unique identifier of the model")
    object: Literal["model"] = Field(default="model", description="Type of object")
    created: int = Field(..., description="Unix creation timestamp")
    owned_by: str = Field(..., description="Owner organization")

class ListModelsResponse(BaseModel):
    object: Literal["list"] = Field(default="list", description="Type of object")
    data: list[Model] = Field(..., description="List of available models")

class OrderEnum(str, Enum):
    asc = "asc"
    desc = "desc"


class VideoStatus(str, Enum):
    queued = "queued"
    processing = "processing"
    completed = "completed"
    failed = "failed"


class VideoContentVariant(str, Enum):
    video = "video"
    preview = "preview"


class VideoQuality(str, Enum):
    standard = "standard"
    hd = "hd"

class ListVideosParams(BaseModel):
    limit: int | None = Field(None, ge=0, le=100, description="Number of items to retrieve")
    order: OrderEnum | None = Field(None, description="Order of results by timestamp")
    after: str | None = Field(None, description="Identifier of the last item in the previous pagination request")


class RetrieveVideoContentParams(BaseModel):
    variant: Optional[VideoContentVariant] = Field(
        VideoContentVariant.video,
        description="Which downloadable asset to return. By default, the MP4 video."
    )

class CreateVideoBody(BaseModel):
    model: str = Field(default="sora-2", description="Model to use to generate the video")
    prompt: str = Field(..., description="Prompt of the video to be generated")
    size: Literal['1024x1792', '1280x720', '1792x1024', '720x1280'] | None = Field(None, description="Video size, e.g., '1024x1808'")
    seconds: str | None = Field(None, description="Video duration in seconds")
    quality: VideoQuality | None = Field(VideoQuality.standard, description="Video quality")

class VideoResource(BaseModel):
    id: str = Field(..., description="Unique identifier for the video")
    object: Literal["video"] = Field("video", description="Object type")
    model: str = Field(..., description="Model used to generate the video")
    status: VideoStatus = Field(..., description="Current status of the video job")
    progress: int | None = Field(None, ge=0, le=100, description="Video generation progress (0-100)")
    created_at: int = Field(..., description="Unix timestamp of creation")
    size: Literal['1024x1792', '1280x720', '1792x1024', '720x1280'] | None = Field(None, description="Video dimensions")
    seconds: str | None = Field(None, description="Video duration in seconds")
    quality: VideoQuality | None = Field(None, description="Video quality")
    error: dict | None = Field(None, description="Error information if the job failed")
    prompt: str | None = Field(None, description="Original prompt used to generate the video")


class VideoListResource(BaseModel):
    data: List[VideoResource] = Field(..., description="List of videos")
    object: Literal["list"] = Field("list", description="Object type")
    has_more: bool | None = Field(None, description="Indicates if more results are available")
    first_id: str | None = Field(None, description="ID of the first element in the list")
    last_id: str | None = Field(None, description="ID of the last element in the list")


class DeletedVideoResource(BaseModel):
    id: str = Field(..., description="Identifier of the deleted video")
    object: Literal["video.deleted"] = Field("video.deleted", description="Object type")
    deleted: bool = Field(True, description="Indicates if the video was successfully deleted")


class VideoPathParams(BaseModel):
    video_id: str = Field(..., description="Video identifier")