import openai
import requests
import urllib.parse
import json
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from io import BytesIO
import os

openai.api_key = os.environ["OPEN_AI_KEY"]



png_bytes_array = []
images = []

system_prompt = """Transcript of a dialog, where the User interacts with an Assistant named Stablediffy. Stablediffy knows much about prompt engineering for stable diffusion (an open-source image generation software). The User asks Stablediffy about prompts for stable diffusion Image Generation.

Possible keywords for stable diffusion: "cinematic, colorful background, concept art, dramatic lighting, high detail, highly detailed, hyper realistic, intricate, intricate sharp details, octane render, smooth, studio lighting, trending on artstation, landscape, scenery, cityscape, underwater, salt flat, tundra, jungle, desert mountain, ocean, beach, lake, waterfall, ripples, swirl, waves, avenue, horizon, pasture, plateau, garden, fields, floating island, forest, cloud forest, grasslands, flower field, flower ocean, volcano, cliff, snowy mountain
city, cityscape, street, downtown"

-- Transcript --

USER: suggest a prompt for a young girl from Swiss sitting by the window with headphones on
ASSISTANT: gorgeous young Swiss girl sitting by window with headphones on, wearing white bra with translucent shirt over, soft lips, beach blonde hair, octane render, unreal engine, photograph, realistic skin texture, photorealistic, hyper realism, highly detailed, 85mm portrait photography, award winning, hard rim lighting photography

USER: suggest a prompt for an mysterious city
ASSISTANT: Mysterious city, cityscape, urban, downtown, street, noir style, cinematic lightning, dramatic lightning, intricate, sharp details, octane render, unreal engine, highly detailed, night scene, dark lighting, gritty atmosphere

USER: suggest a prompt for a high quality render of a car in 1950
ASSISTANT: Car in 1950, highly detailed, classic car, 1950's, highly detailed, dramatic lightning, cinematic lightning, unreal engine

USER:"""

def get_modified_text_response(user_query):
    try:
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "system", "content": system_prompt},
                    {"role": "user", "content": "suggest a prompt for" +  user_query}]
        )
        res = response["choices"][0]["message"]["content"].replace('\'', '')
        print(res)
        return res
    except:
        return "no gpt response"

def text_to_image(text, selected_value):
    global png_bytes_array
    images = []
    if(selected_value == "New"):
        png_bytes_array = []
        image_urls = generate_image(text)
        for img in image_urls:
            response = requests.get(img['url'])
            image = Image.open(BytesIO(response.content))
            images.append(image)
            bytesIO = BytesIO()
            image.save(bytesIO, format="PNG")
            png_bytes = bytesIO.getvalue()
            png_bytes_array.append(png_bytes)
        return images
    
    else:
        index = int(selected_value[-1])
        image_urls = variation_image(png_bytes_array[index])
        for img in image_urls:
            response = requests.get(img['url'])
            image = Image.open(BytesIO(response.content))
            images.append(image)
        return images
    
def variation_image(image):
    response = openai.Image.create_variation(
    image = image,
    n=4,
    size="1024x1024"
        )
    #print(response)
    #image_url = response['data'][0]['url']
    return response['data']
def generate_image(prompt):
    better_prompt = get_modified_text_response(prompt)
    response = openai.Image.create(
    prompt=better_prompt,
    n=4,
    size="1024x1024"
        )

    return response['data']
def main():
    radio_buttons = gr.inputs.Radio(["New", "Var0","Var1", "Var2", "Var3"], label="Select a variation option")
    out = gr.Gallery(
            label="Generated images", show_label=False, elem_id="gallery"
        ).style(columns=[2], rows=[2], object_fit="contain", height="auto")
    iface = gr.Interface(fn=text_to_image, inputs=["text",radio_buttons], 
                         #outputs=gr.Gallery([gr.Image(type="pil").style(height=300,width=300),gr.Image(type="pil").style(height=300,width=300),gr.Image(type="pil").style(height=300,width=300),gr.Image(type="pil").style(height=300,width=300)]))
                         outputs = out)
    iface.launch()  

if __name__ == "__main__":
    main()