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

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

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'
    
    # Qwen models
    QWEN_IMAGE_EDIT = 'Qwen/Qwen-Image-Edit'
    QWEN_IMAGE = 'Qwen/Qwen-Image'


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", "1024x1536", "1536x1024"]] = 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.")


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, Literal["dall-e-2", "gpt-image-1"]]] = 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", "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 CreateImageVariationRequest(BaseModel):
    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.")
    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.")
    size: Optional[Literal["256x256", "512x512", "1024x1024"]] = Field("1024x1024", description="The size of the generated images. Must be one of `256x256`, `512x512`, or `1024x1024`.")
    user: Optional[str] = Field(None, description="A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.")


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] = Field( default_factory=lambda: [""], description="API KEYS allowed to make requests")
    max_concurrent_infer: int = Field(default=10, description="Maximum number of inferences running at the same time")
    block_request: bool = Field(default=False, description="Block requests during the maximum concurrent inferences")