import gradio as gr
import requests
from PIL import Image
from io import BytesIO
import os

# Get token from environment variable (Hugging Face Spaces will provide this)
HF_API_TOKEN = os.getenv("HF_TOKEN")

# Best working model
MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
API_URL = f"https://api-inference.huggingface.co/models/{MODEL}"

headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}

def generate_image(prompt):
    """Generate image using Hugging Face API"""
    if not prompt:
        return None, "Please enter a prompt!"
    
    if not HF_API_TOKEN:
        return None, "❌ Error: HF_TOKEN not found. Please set it in Space settings."
    
    payload = {"inputs": prompt}
    
    try:
        response = requests.post(API_URL, headers=headers, json=payload)
        
        if response.status_code == 200:
            image = Image.open(BytesIO(response.content))
            return image, "✅ Image generated successfully!"
        else:
            return None, f"❌ Error {response.status_code}: {response.text}"
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

# Create Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎨 AI Image Generator
        ### Powered by Stable Diffusion XL
        Enter a text prompt to generate an image!
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Enter your prompt",
                placeholder="e.g., A futuristic city at sunset, Robot holding a red skateboard...",
                lines=3
            )
            generate_btn = gr.Button("🎨 Generate Image", variant="primary", size="lg")
            status_text = gr.Textbox(label="Status", interactive=False)
            
            gr.Markdown("### 💡 Example Prompts:")
            gr.Examples(
                examples=[
                    ["A cute cat wearing sunglasses on a beach"],
                    ["A futuristic city at sunset with flying cars"],
                    ["Robot holding a red skateboard"],
                    ["A magical forest with glowing mushrooms"],
                    ["An astronaut riding a horse on Mars"],
                ],
                inputs=prompt_input
            )
        
        with gr.Column(scale=1):
            image_output = gr.Image(label="Generated Image", type="pil")
    
    # Connect the button
    generate_btn.click(
        fn=generate_image,
        inputs=prompt_input,
        outputs=[image_output, status_text]
    )
    
    # Also allow Enter key to generate
    prompt_input.submit(
        fn=generate_image,
        inputs=prompt_input,
        outputs=[image_output, status_text]
    )

# Launch the app
if __name__ == "__main__":
    print("🚀 Starting AI Image Generator...")
    print("📱 Opening in your browser...")
    demo.launch(share=False)