import gradio as gr
import os
import requests
import io
import time
import random
from PIL import Image

API_KEY = os.environ.get('IMAGE_API_KEY', '')

def generate_image(prompt, size="1024x1024"):
    if not prompt:
        return None, "Please enter a prompt"
    
    if not API_KEY:
        return None, "API key not configured"
    
    size_map = {'1024x1024': (1024, 1024), '1792x1024': (1792, 1024), '1024x1792': (1024, 1792)}
    width, height = size_map.get(size, (1024, 1024))
    
    try:
        # Submit job
        resp = requests.post(
            'https://api.deapi.ai/api/v2/images/generations',
            headers={
                'Authorization': f'Bearer {API_KEY}',
                'Content-Type': 'application/json'
            },
            json={
                'model': 'Flux1schnell',
                'prompt': prompt,
                'width': width,
                'height': height,
                'steps': 4,
                'seed': random.randint(0, 999999)
            }
        )
        
        if resp.status_code not in [200, 201]:
            return None, f"Error {resp.status_code}: {resp.text[:150]}"
        
        # Get request_id from nested structure
        data = resp.json()
        request_id = data.get('data', {}).get('request_id')
        if not request_id:
            return None, f"No request_id: {resp.text[:200]}"
        
        # Poll for result
        for _ in range(30):
            time.sleep(2)
            result = requests.get(
                f'https://api.deapi.ai/api/v2/images/generations/{request_id}',
                headers={'Authorization': f'Bearer {API_KEY}'}
            )
            if result.status_code == 200:
                rdata = result.json()
                status = rdata.get('data', {}).get('status')
                if status == 'succeeded':
                    output = rdata.get('data', {}).get('output', [])
                    if output and len(output) > 0:
                        img_url = output[0]
                        img_resp = requests.get(img_url)
                        if img_resp.status_code == 200:
                            img = Image.open(io.BytesIO(img_resp.content))
                            return img, "Success!"
                elif status == 'failed':
                    return None, f"Failed: {rdata.get('data', {}).get('error')}"
        
        return None, "Timeout"
        
    except Exception as e:
        return None, str(e)[:200]

interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="A sunset over ocean..."),
        gr.Dropdown(choices=["1024x1024", "1792x1024", "1024x1792"], label="Size", value="1024x1024")
    ],
    outputs=[
        gr.Image(label="Generated Image"),
        gr.Textbox(label="Status")
    ],
    title="AI Image Generator",
    description="Generate images using FLUX model"
)

interface.launch(server_port=7860, server_name="0.0.0.0")