import streamlit as st
import requests
from openai import OpenAI
from PIL import Image
import io
import os
from datetime import datetime

def save_image_from_url(image_url, prompt):
    """Save image from URL to local file with prompt as part of filename"""
    # Create directory if it doesn't exist
    if not os.path.exists("generated_images"):
        os.makedirs("generated_images")
    
    # Download image
    response = requests.get(image_url)
    
    # Create filename from timestamp and shortened prompt
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    # Clean prompt for filename (remove special characters and limit length)
    clean_prompt = "".join(c for c in prompt if c.isalnum() or c.isspace())[:30]
    clean_prompt = clean_prompt.replace(" ", "_")
    
    output_path = f"generated_images/{timestamp}_{clean_prompt}.png"
    
    with open(output_path, "wb") as f:
        f.write(response.content)
    return output_path

def main():
    st.title("DALL-E 2 Image Generator")
    
    # Sidebar for API key
    st.sidebar.header("Settings")
    api_key = st.sidebar.text_input("Enter OpenAI API Key", type="password")
    
    if not api_key:
        st.warning("Please enter your OpenAI API key in the sidebar to continue.")
        return
    
    # Main content
    st.write("Generate images using DALL-E 2")
    
    # Image generation form
    with st.form("image_generation_form"):
        # Prompt input
        prompt = st.text_area("Enter your prompt", 
                            help="Describe the image you want to generate",
                            max_chars=1000)
        
        # Generation options
        col1, col2, col3 = st.columns(3)
        
        with col1:
            size_options = ["1024x1024", "512x512", "256x256"]
            size = st.selectbox("Image size", size_options)
        
        with col2:
            quality = st.selectbox("Image quality", ["standard", "hd"])
        
        with col3:
            n_images = st.slider("Number of images", min_value=1, max_value=4, value=1)
        
        # Submit button
        submit_button = st.form_submit_button("Generate Images")
    
    # Handle form submission
    if submit_button:
        if not prompt:
            st.error("Please enter a prompt.")
            return
        
        try:
            # Initialize OpenAI client
            client = OpenAI(api_key=api_key)
            
            with st.spinner("Generating images..."):
                # Generate images
                response = client.images.generate(
                    model="dall-e-2",
                    prompt=prompt,
                    size=size,
                    quality=quality,
                    n=n_images,
                )
                
                # Display generated images
                st.subheader("Generated Images")
                
                # Create columns based on number of images
                cols = st.columns(n_images)
                
                # Display each image in its own column
                for idx, image_data in enumerate(response.data):
                    # Save image locally
                    saved_path = save_image_from_url(image_data.url, prompt)
                    
                    with cols[idx]:
                        # Display image
                        st.image(saved_path, caption=f"Generated Image {idx+1}", use_container_width=True)
                        
                        # Add download button
                        with open(saved_path, "rb") as file:
                            st.download_button(
                                label=f"Download Image {idx+1}",
                                data=file,
                                file_name=f"dalle2_generation_{idx+1}.png",
                                mime="image/png"
                            )
                
                # Display prompt used
                st.info(f"Prompt used: {prompt}")
                
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
            if "Rate limit" in str(e):
                st.info("You've hit the rate limit. Please wait a moment before trying again.")
            elif "billing" in str(e).lower():
                st.info("Please check your OpenAI account billing status.")

if __name__ == "__main__":
    main()