import gradio as gr
import requests
import json
import time

url = "https://nexra.aryahcr.cc/api/image/complements"
headers = {
    "Content-Type": "application/json"
}

def generate_image(prompt, model):
    data = {
        "prompt": prompt,
        "model": model,
    }
    try:
        response = requests.post(url, headers=headers, data=json.dumps(data))
        
        if response.status_code == 200:
            response_data = response.json()
            image_id = response_data.get("id")
            
            if not image_id:
                return "Error: No image ID returned in the response."

            while True:
                status_response = requests.get(f"{url}/{image_id}")
                
                if status_response.status_code == 200:
                    status_data = status_response.json()
                    status = status_data.get("status")
                    
                    if status == "completed":
                        images = status_data.get("images")
                        if images and isinstance(images, list):
                            image_url = images[0]
                            return image_url
                        else:
                            return "Error: No images found in the response."
                    elif status == "error":
                        return "Error: Image generation failed."
                    elif status == "not_found":
                        return "Error: Image ID not found."
                    elif status == "pending":
                        time.sleep(1)
                    else:
                        return f"Error: Unexpected status '{status}'."
                else:
                    return f"Error: Status check failed with code {status_response.status_code}."
        else:
            return f"Error: Initial request failed with code {response.status_code}."
    except json.JSONDecodeError:
        return "Error: Invalid JSON response."
    except Exception as e:
        return f"Exception occurred: {str(e)}"

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Enter prompt"),
        gr.Radio(["dalle2"], label="Select Model")
    ],
    outputs=gr.Image(type="filepath"),
    title="DALLE2 Generation",
    description="DALLE-2 Generation from Nextra API, NOTE: DALL-E 3 is not supported yet."
)
iface.launch()