import streamlit as st
import requests
import os

# Load API key from Hugging Face secrets
api_key = os.getenv("KEY")

# Streamlit UI
st.title("AI Image Generator")
st.write("Generate AI-powered images using text prompts.")

prompt = st.text_area("Enter your image description:")
output_type = st.selectbox("Select output format:", ["png", "jpg"])

def generate_image():
    url = "https://ai-image-generator14.p.rapidapi.com/"
    headers = {
        "x-rapidapi-key": api_key,
        "x-rapidapi-host": "ai-image-generator14.p.rapidapi.com",
        "Content-Type": "application/json"
    }
    payload = {
        "jsonBody": {
            "function_name": "image_generator",
            "type": "image_generation",
            "query": prompt,
            "output_type": output_type
        }
    }
    response = requests.post(url, json=payload, headers=headers)
    return response.json()

if st.button("Generate Image"):
    if prompt:
        with st.spinner("Generating image..."):
            result = generate_image()
            
            image_url = result.get("message", {}).get("output_png")
            if image_url:
                st.image(image_url, caption="Generated Image", use_container_width=True)
                
                # Provide download button
                st.download_button(
                    label="Download Image",
                    data=requests.get(image_url).content,
                    file_name=f"generated_image.{output_type}",
                    mime=f"image/{output_type}"
                )
            else:
                st.error("Failed to generate image. No image URL found in response.")
                st.write("Response Data:", result)  # Debugging: Show response data
    else:
        st.warning("Please enter an image description.")