import streamlit as st
import requests
from requests.structures import CaseInsensitiveDict

import json

query_url = "https://api.openai.com/v1/images/generations"

st.title("DALL-E 2 API Image Generation Demo")

st.write("Enter a prompt to generate an image")

prompt = st.text_area("Prompt", "A cat riding a bike")

model = st.selectbox(
    "Select a DALL-E 2 model",
    ["image-alpha-001", ""]
)

num_images = st.slider("Number of images to generate", min_value=1, max_value=10, value=1)

headers = CaseInsensitiveDict()
headers["Content-Type"] = "application/json"
api_key = "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"
headers["Authorization"] = f"Bearer {api_key}"

data = """
{
    """
data += f'"model": "{model}",'
data += f'"prompt": "{prompt}",'
data += f'"num_images": {num_images}'
data += """
}
"""

def generate_images():
    resp = requests.post(query_url, headers=headers, data=data)

    if resp.status_code != 200:
        raise ValueError("Failed to generate image "+resp.text)

    response_text = json.loads(resp.text)
    return response_text['data']

if st.button("Generate Images"):
    with st.spinner("Generating images..."):
        image_data = generate_images()
        for idx, image in enumerate(image_data):
            st.image(image['url'], caption=f"Image {idx+1}", width=400)