import streamlit as st
import requests
import tempfile

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

st.title('Image Generation')

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

# Define your toppings options
toppings_options = ["Mushroom", "Tomatoes", "Onion", "Arugula", "Bacon", "Basil", "Broccoli", "Cheese", "Chicken", "Corn", "Ham", "Olives", "Pepperoni", "Peppers", "Pineapple"]

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

    # Pizza order input
    pizza_names = ["Crispy Chicken", "Chicken Pesto", "..."]  # Add the rest of the pizza names
    name = st.selectbox("Name of Pizza", pizza_names)
    size_pizza = st.selectbox("Size of Pizza", ["Small", "Medium", "Large", "Extra Large"])
    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)
    orientation = ["Top View", "Side View", "Lateral View"]
    orientation_checkbox = st.multiselect("Orientation", orientation, default=["Top View"])
    
    generate_prompt_button = st.form_submit_button(label='Generate Prompt')

# Generate the dynamic prompt
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 the following toppings: {', '.join(toppings)}. I want to generate an image for the given description in {', '.join(orientation_checkbox)} with a background color of {bg_color}."
    st.write("Prompt:", st.session_state['dynamic_prompt'])

# Display the button to generate the image
with st.form(key='image_form'):
    generate_image_button = st.form_submit_button(label='Generate Image')

if generate_image_button and st.session_state['dynamic_prompt']:
    # Here, integrate with your custom model to generate the image
    # For now, using a placeholder URL
    st.session_state['image_url'] = "url_of_your_generated_image"
    st.session_state['image_generated'] = True

    # Display the generated image
    if st.session_state['image_url']:
        st.image(st.session_state['image_url'], caption="Generated Image", use_column_width=True)

# 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)  # Save the image data
        st.download_button(
            label="Download Image",
            data=tmp_file,
            file_name="generated_image.png",
            mime="image/png"
        )