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

# Use environment variable for API key
API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"

def get_api_key():
    """Retrieve API key from environment variable or Streamlit secrets"""
    # First, try to get from environment variable
    api_key = os.getenv('HUGGINGFACE_API_KEY')
    
    # If not found, try Streamlit secrets
    if not api_key:
        try:
            api_key = st.secrets["HUGGINGFACE_API_KEY"]
        except Exception:
            st.error("API Key not found. Please configure your API key.")
            return None
    
    return api_key

# Lego Minifigurine Style
style_list = [
    {
        "name": "lego minifigurine",
        "prompt": "legoman, lego minifigurine, lego universe, lego, lego figurine, lego style, lego-man",
        "negative_prompt": "photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white",
    }
]

def query_image_generation(image_bytes):
    """Generate Lego minifigurine style image using Hugging Face API"""
    # Get API key
    api_key = get_api_key()
    if not api_key:
        st.error("Cannot proceed without API key")
        return None

    # Prepare headers with API key
    headers = {"Authorization": f"Bearer {api_key}"}
    
    payload = {
        "inputs": style_list[0]["prompt"],
        "negative_prompt": style_list[0]["negative_prompt"],
        "image": base64.b64encode(image_bytes).decode('utf-8')
    }
    
    try:
        response = requests.post(API_URL, headers=headers, json=payload)
        response.raise_for_status()
        return response.content
    except requests.exceptions.RequestException as e:
        st.error(f"API request failed: {e}")
        return None

def main():
    st.title("🧱 Lego Minifigurine Image Generator")
    st.write("Upload an image to transform it into a Lego minifigurine style!")

    # Image upload
    uploaded_file = st.file_uploader(
        "Choose an image...", 
        type=["jpg", "jpeg", "png"], 
        help="Upload an image to convert to Lego style"
    )

    if uploaded_file is not None:
        # Display original image
        original_image = Image.open(uploaded_file)
        st.subheader("Original Image")
        st.image(original_image, use_column_width=True)

        # Convert image to bytes
        img_byte_arr = io.BytesIO()
        original_image.save(img_byte_arr, format='PNG')
        img_byte_arr = img_byte_arr.getvalue()

        # Generate button
        if st.button("Generate Lego Style"):
            with st.spinner('Generating Lego Minifigurine...'):
                generated_image_bytes = query_image_generation(img_byte_arr)
                
                if generated_image_bytes:
                    # Display generated image
                    st.subheader("Lego Minifigurine Style")
                    generated_image = Image.open(io.BytesIO(generated_image_bytes))
                    st.image(generated_image, use_column_width=True)

                    # Download button
                    buffered = io.BytesIO()
                    generated_image.save(buffered, format="PNG")
                    img_str = base64.b64encode(buffered.getvalue()).decode()
                    href = f'<a href="data:image/png;base64,{img_str}" download="lego_minifigurine.png">Download Lego Image</a>'
                    st.markdown(href, unsafe_allow_html=True)
                else:
                    st.error("Failed to generate image. Please try again.")

if __name__ == "__main__":
    main()