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

def get_api_key():
    return os.environ.get("PIXAZO_API_KEY")

def generate_image(prompt):
    api_key = get_api_key()
    if not api_key:
        return None, "❌ API key not set"
    
    url = "https://gateway.pixazo.ai/flux-1-schnell/v1/getData"
    
    headers = {
        'Content-Type': 'application/json',
        'Cache-Control': 'no-cache',
        'Ocp-Apim-Subscription-Key': api_key,
    }
    
    body = {
        "prompt": prompt,
        "num_steps": 4,
        "height": 512,
        "width": 512,
    }
    
    try:
        response = requests.post(url, headers=headers, json=body, timeout=60)
        
        if response.status_code == 200:
            result = response.json()
            
            if "output" in result:
                image_url = result["output"]
                
                img_response = requests.get(image_url, timeout=30)
                if img_response.status_code == 200:
                    image = Image.open(io.BytesIO(img_response.content))
                    if image.mode in ('RGBA', 'LA', 'P'):
                        image = image.convert('RGB')
                    return image, "✅ Image generated successfully!"
                else:
                    return None, f"❌ Failed to download image: HTTP {img_response.status_code}"
            else:
                return None, "❌ No image URL in response"
        else:
            return None, f"❌ API error: {response.status_code}"
            
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

# Simple interface
with gr.Blocks(title="Pixazo Image Generator") as demo:
    gr.Markdown("# 🎨 Pixazo Image Generator")
    
    api_key = get_api_key()
    if api_key:
        gr.Markdown("✅ API Key: Configured")
    else:
        gr.Markdown("❌ API Key: Not configured. Set PIXAZO_API_KEY in secrets.")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=3)
            btn = gr.Button("Generate", variant="primary")
        
        with gr.Column():
            image = gr.Image(label="Generated Image")
            status = gr.Textbox(label="Status", interactive=False)
    
    btn.click(generate_image, inputs=prompt, outputs=[image, status])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)