import os
from typing import Literal

from openai import AsyncOpenAI

from ...errors import APIKeyNotFoundError
from ...types.field import Field, resolve_default
from ...types.file import File


async def gen_images(
    prompt: str = Field(
        description="A text description of the desired image(s). Max length: 32000 characters.",
        max_length=32000,
    ),
    background: Literal["transparent", "opaque", "auto"] = Field(
        default=None,
        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."
            "If transparent, the output format needs to support transparency, so it should be set to either png (default value) or webp.",
    ),
    moderation: Literal["low", "auto"] = Field(
        default=None,
        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).",
    ),
    n: int = Field(
        default=None,
        description="The number of images to generate. Must be between 1 and 10.",
        ge=1,
        le=10,
    ),
    output_compression: int = Field(
        default=None,
        description="The compression level (0-100%) for the generated images. " 
            "This parameter is only supported for gpt-image-1 with the webp or jpeg output formats, and defaults to 100.",
    ),
    output_format: Literal["png", "jpeg", "webp"] = Field(
        default=None,
        description="The format in which the generated images are returned. "
            "This parameter is only supported for gpt-image-1. Must be one of png, jpeg, or webp.",
    ),
    quality: Literal["auto", "high", "medium", "low"] = Field(
        default=None,
        description="The quality of the image that will be generated.",
    ),
    size: Literal["1024x1024", "1536x1024", "1024x1536", "auto"] = Field(
        default=None,
        description="The size of the generated image.",
    ),
    # model -> we only support gpt-image-1 (the other ones are bad)
    # response_format -> we have this fixed, we return File either way
    # style -> not available for gpt-image-1
    # user -> not used
) -> list[File]:

    prompt = resolve_default("prompt", prompt)
    background = resolve_default("background", background)
    moderation = resolve_default("moderation", moderation)
    n = resolve_default("n", n)
    output_compression = resolve_default("output_compression", output_compression)
    output_format = resolve_default("output_format", output_format)
    quality = resolve_default("quality", quality)
    size = resolve_default("size", size)

    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        raise APIKeyNotFoundError("OPENAI_API_KEY not found")
    
    client = AsyncOpenAI(api_key=api_key)

    kwargs = {}

    if prompt is not None:
        kwargs["prompt"] = prompt
    if background is not None:
        kwargs["background"] = background
    if moderation is not None:
        kwargs["moderation"] = moderation
    if n is not None:
        kwargs["n"] = n
    if output_compression is not None:
        kwargs["output_compression"] = output_compression
    if output_format is not None:
        kwargs["output_format"] = output_format
    if quality is not None:
        kwargs["quality"] = quality
    if size is not None:
        kwargs["size"] = size

    res = await client.images.generate(
        model="gpt-image-1",
        **kwargs,
    )

    files = [File.validate(f"data:image/png;base64,{image.b64_json}") for image in res.data]
    return files


async def edit_images(
    images: list[File] = Field(
        description="The image(s) to edit. Must be a supported image file or an array of images. "
            "Each image should be a png, webp, or jpg file less than 50MB. You can provide up to 16 images."
    ),
    prompt: str = Field(
        description="A text description of the desired image(s). Max length: 32000 characters.",
        max_length=32000,
    ),
    background: Literal["transparent", "opaque", "auto"] = Field(
        default=None,
        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."
            "If transparent, the output format needs to support transparency, so it should be set to either png (default value) or webp.",
    ),
    mask: File = Field(
        default=None,
        description="An additional image whose fully transparent areas (e.g. where alpha is zero) indicate where image should be edited. "
            "If there are multiple images provided, the mask will be applied on the first image. "
            "Must be a valid PNG file, less than 4MB, and have the same dimensions as image."
    ),
    n: int = Field(
        default=None,
        description="The number of images to generate. Must be between 1 and 10.",
        ge=1,
        le=10,
    ),
    quality: Literal["auto", "high", "medium", "low"] = Field(
        default=None,
        description="The quality of the image that will be generated.",
    ),
    size: Literal["1024x1024", "1536x1024", "1024x1536", "auto"] = Field(
        default=None,
        description="The size of the generated image.",
    ),
    # model -> we only support gpt-image-1 (the other ones are bad)
    # response_format -> we have this fixed, we return File either way
    # user -> not used
) -> list[File]:

    images = resolve_default("images", images)
    prompt = resolve_default("prompt", prompt)
    background = resolve_default("background", background)
    mask = resolve_default("mask", mask)
    n = resolve_default("n", n)
    quality = resolve_default("quality", quality)
    size = resolve_default("size", size)

    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        raise APIKeyNotFoundError("OPENAI_API_KEY not found")
    
    client = AsyncOpenAI(api_key=api_key)

    kwargs = {}

    if images is not None:
        kwargs["image"] = images
    if prompt is not None:
        kwargs["prompt"] = prompt
    if background is not None:
        kwargs["background"] = background
    if mask is not None:
        kwargs["mask"] = mask
    if n is not None:
        kwargs["n"] = n
    if quality is not None:
        kwargs["quality"] = quality
    if size is not None:
        kwargs["size"] = size

    res = await client.images.edit(
        model="gpt-image-1",
        **kwargs,
    )

    files = [File.validate(f"data:image/png;base64,{image.b64_json}") for image in res.data]
    return files
