import streamlit as st
#import openai
#import os
import tempfile
#from dotenv import load_dotenv
import requests  # For downloading images

# Load environment variables
#load_dotenv()
#openai_api_key = os.getenv("OPENAI_API_KEY")

# Initialize OpenAI API key
#openai.api_key = openai_api_key

# Set page config
st.set_page_config(page_title='Image Generation', page_icon=':money_with_wings:', layout='wide')

st.title('Image Generation')



# Initialize session state
if 'dynamic_prompt' not in st.session_state:
    st.session_state['dynamic_prompt'] = ""
if 'generate_image' not in st.session_state:
    st.session_state['generate_image'] = False
if 'image_generated' not in st.session_state:
    st.session_state['image_generated'] = False
if 'image_url' not in st.session_state:
    st.session_state['image_url'] = ""

# Main content
with st.form(key='my_form'):
    # Image generation input
    size = st.selectbox('Select size of the images', ('256x256', '512x512', '1024x1024'))
    num_images = st.selectbox('select Number of Image to be generatd ', ('1', '2', "3", '4', '5'))  # Number of images

    # Pizza order input
    pizza_names =["Crispy Chicken", "Chicken Pesto", "Deluxe Pepperoni", "Truffle Mushroom",
                   "Ultimate Cheese", "Spicy Veggie Ranch", "Classic Chicken Ranch",
                   "BBQ Chicken Ranch", "Spicy Chicken Ranch", "Very Veggie",
                   "Super Supreme", "Classic Pepperoni", "Margherita", "Cheeky Chicken",
                   "Chicken Super Supreme", "Chicken BBQ Supreme", "Chicken Fajita",
                   "Chicken Shawerma", "Other"]
    name = st.selectbox("Name of Pizza", pizza_names)
    size_pizza = st.selectbox("Size of Pizza", ["Small", "Medium", "Large", "Extra Large"])
    toppings_options = ["Arugula", "Bacon", "Basil", "Broccoli", "Cheese", "Chicken", "Corn", "Ham", "Mushroom", "Olives", "Onion", "Pepperoni", "Peppers", "Pineapple", "Tomatoes" ]  # rest of the list
    toppings = st.multiselect("Toppings", toppings_options, default=["Mushroom", "Tomatoes", "Onion"])
    prompt = st.text_input(label='Enter additional Descriptions ')
    bg_color_options = ['White', 'Black', 'Light Blue', 'Light Green', 'Light Yellow', 'Light Pink']
    bg_color = st.selectbox('Select Background Color', bg_color_options)
    orient =["Top View", "Side View", "lateral View"]
    orientation_checkbox = st.multiselect("Orientation", orient, default=["Top View"])
    
    generate_prompt_button = st.form_submit_button(label='Generate Prompt')

# Action when the generate prompt button is clicked
if generate_prompt_button:
    st.session_state['dynamic_prompt'] = f"Imagine you are a marketer who wants to post an image on social media for a {size_pizza} {name} Pizza with following toppings on it  {', '.join(toppings)}. I want to generate image for the given descrpition in {orientation_checkbox} with back ground color as {bg_color} "
    st.write(st.session_state['dynamic_prompt'])
    st.session_state['generate_image'] = True

# Display the button to generate the image only if the prompt has been generated
if st.session_state['generate_image']:
    with st.form(key='image_form'):
        generate_image_button = st.form_submit_button(label='Generate Image')

    if generate_image_button:

        # Image generation code using OpenAI's API
        if st.session_state['dynamic_prompt']:
            response = openai.Image.create(
                prompt=st.session_state['dynamic_prompt'],
                n=num_images,
                size=size,
            )
            st.session_state['image_url'] = response["data"][0]["url"]
            st.image(st.session_state['image_url'], caption="Generated image", use_column_width=True)
        
        # Update the session state to reflect that the image has been generated
        st.session_state['image_generated'] = True

        submitted = st.form_submit_button("Submit")
        #st.write("Outside the form")


# Download button for the generated image
if st.session_state['image_generated']:
    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
        response = requests.get(st.session_state['image_url'])
        with open(tmp_file.name, "wb") as f:
            f.write(response.content)  # Saving the image data
        st.download_button(
            label="Download image",
            data=tmp_file,
            file_name="generated_image.png",
            mime="image/png"
        )

    combined_info = f"""
    Image Generation Details:
    - Text Prompt: {st.session_state['dynamic_prompt']}
    - Image Size: {size}
    - Number of Images: {num_images}
    """
    st.write(combined_info)