import streamlit as st
import requests
import os

# Load API key from environment variables (or manually set it)
api_key = os.getenv("KEY")

# Streamlit UI
st.title("AI Logo Generator 🎨")
st.write("Generate a unique logo for your business using AI!")

# User Inputs
prompt = st.text_input("Enter a description for the logo:", "Make a logo for my gaming company BTZ")
style = st.selectbox("Select a Logo Style:", [28, 29, 30], index=0)  # Add more styles if available
size = st.radio("Select Image Size:", ["1-1", "2-3", "3-4"], index=0)

def generate_logo(prompt, style, size):
    api_url = "https://ai-logo-generator.p.rapidapi.com/aaaaaaaaaaaaaaaaaiimagegenerator/quick.php"
    headers = {
        "x-rapidapi-key": api_key,
        "x-rapidapi-host": "ai-logo-generator.p.rapidapi.com",
        "Content-Type": "application/json"
    }
    payload = {
        "prompt": prompt,
        "style_id": style,
        "size": size
    }
    
    response = requests.post(api_url, json=payload, headers=headers)
    return response.json()

# Generate button
if st.button("Generate Logo"):
    with st.spinner("Generating logo... Please wait!"):
        result = generate_logo(prompt, style, size)
        
        # Extracting image URLs from the JSON response
        image_data = result.get("final_result", [])
        
        if image_data:
            st.subheader("Generated Logo Designs")
            
            for index, image in enumerate(image_data):
                image_url = image.get("origin")  # Extracting the logo link
                
                if image_url:
                    st.image(image_url, caption=f"Logo Design {index+1}", use_container_width=True)
                    
                    # Download button with a unique key
                    st.download_button(
                        label="Download Logo",
                        data=requests.get(image_url).content,
                        file_name=f"logo_design_{index+1}.webp",
                        mime="image/webp",
                        key=f"download_{index}"  # Unique key to prevent duplicate errors
                    )
        else:
            st.error("Failed to generate logo. No image URL found.")
            st.write("Response Data:", result)  # Debugging: Show full response data