import streamlit as st
from PIL import Image
import requests
import io
import base64

# Set page config
st.set_page_config(page_title="AI Image Generator", page_icon="🎨", layout="wide")

# Custom CSS
st.markdown("""
<style>
.stApp {
    background-color: #f0f0f5;
}
.stButton>button {
    background-color: #4CAF50;
    color: white;
    font-weight: bold;
}
</style>
""", unsafe_allow_html=True)

# Title
st.title("🎨 AI Image Generator (Dalle-mini)")

# Input area
prompt = st.text_area("Enter your image description:", height=100)

# Generate button
if st.button("Generate Image"):
    if prompt:
        try:
            with st.spinner("Generating your masterpiece..."):
                # Make a request to the Dalle-mini API (Craiyon)
                response = requests.post("https://backend.craiyon.com/generate", json={"prompt": prompt})
                
                if response.status_code == 200:
                    image_data = response.json()['images'][0]
                    image = Image.open(io.BytesIO(base64.b64decode(image_data)))
                    
                    # Display the image
                    st.image(image, caption="Generated Image", use_column_width=True)
                    
                    # Add download button
                    buf = io.BytesIO()
                    image.save(buf, format="PNG")
                    st.download_button(
                        label="Download Image",
                        data=buf.getvalue(),
                        file_name="generated_image.png",
                        mime="image/png"
                    )
                else:
                    st.error("Failed to generate image. Please try again.")
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
    else:
        st.warning("Please enter a prompt to generate an image.")

# Footer
st.markdown("---")
st.markdown("Created with ❤️ for 3EyeDimensions AI Engineer Internship")
st.markdown("Note: This app uses the Dalle-mini model for image generation.")