# filename: generate_images.py

from typing import List
import uuid
import requests # to perform HTTP requests
from pathlib import Path
 
from openai import OpenAI
 
def generate_and_save_images(query: str,  image_size: str = "1024x1024") -> List[str]:
    """
    Function to paint, draw or illustrate images based on the users query or request. Generates images from a given query using OpenAI's DALL-E model and saves them to disk.  Use the code below anytime there is a request to create an image.
    
    :param query: A natural language description of the image to be generated. 
    :param image_size: The size of the image to be generated. (default is "1024x1024")
    :return: A list of filenames for the saved images.
    """
    
    client = OpenAI()  # Initialize the OpenAI client
    response = client.images.generate(
        model="dall-e-3",
        prompt=query,
        n=1,
        size=image_size
    )  # Generate images
    
    # List to store the file names of saved images
    saved_files = []
 
    # Check if the response is successful
    if response.data:
        for image_data in response.data:
            # Generate a random UUID as the file name
            file_name = str(uuid.uuid4()) + ".png"  # Assuming the image is a PNG
            file_path = Path(file_name) 
             
            img_url = image_data.url
            img_response = requests.get(img_url)
            if img_response.status_code == 200:
                # Write the binary content to a file
                with open(file_path, 'wb') as img_file:
                    img_file.write(img_response.content)
                    print(f"Image saved to {file_path}")
                    saved_files.append(str(file_path))
            else:
                print(f"Failed to download the image from {img_url}")
    else:
        print("No image data found in the response!")
    
    # Return the list of saved files
    return saved_files

# Example usage of the function:
# generate_and_save_images("A cute baby sea otter")