import streamlit as st
from PIL import Image
import io
from utils import query_hf_api

# Set page configuration
st.set_page_config(
    page_title="AI Image Generator",
    page_icon="🎨",
    layout="centered",
    initial_sidebar_state="collapsed",
)

# App Header
st.title("🎨 AI Image Generator")
st.markdown(
    """
    <style>
    .main-title {
        font-size: 2.5rem;
        color: #4CAF50;
        font-weight: bold;
        text-align: center;
    }
    .sub-title {
        font-size: 1.2rem;
        color: #555;
        text-align: center;
        margin-bottom: 20px;
    }
    .footer {
        font-size: 0.9rem;
        color: #888;
        text-align: center;
        margin-top: 50px;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

st.markdown('<h3 class="sub-title">Generate images from your imagination!</h3>', unsafe_allow_html=True)

# Input prompt
prompt = st.text_input(
    "Enter your image prompt:",
    placeholder="e.g., 'Astronaut riding a bike on Mars at sunset'",
)

# Generate button
if st.button("✨ Generate Image"):
    if prompt.strip():
        with st.spinner("🖌️ Generating your masterpiece... Please wait!"):
            try:
                # Query the Hugging Face API
                image_bytes = query_hf_api(prompt)
                image = Image.open(io.BytesIO(image_bytes)).convert("RGB")  # Ensure it's RGB for JPEG
                
                # Display the generated image
                st.image(image, caption="Generated Image", use_container_width=True)
                
                # Save the image as JPEG to a BytesIO buffer
                buffer = io.BytesIO()
                image.save(buffer, format="JPEG", quality=95)
                buffer.seek(0)
                
                # Add a download button
                st.download_button(
                    label="📥 Download Image",
                    data=buffer,
                    file_name="generated_image.jpg",
                    mime="image/jpeg",
                )
            except Exception as e:
                st.error(f"🚨 Error: {str(e)}")
    else:
        st.warning("Please enter a prompt to generate an image.")

# Footer
st.markdown(
    """
    <div class="footer">
        Made with ❤️ by <b>Sumit Yadav</b>. Credits to 🤗 Spaces for hosting this app.
    </div>
    """,
    unsafe_allow_html=True,
)