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

def generate_image(prompt, width, height, api_key):
    headers = {
        'accept': 'application/json',
        'x-key': api_key,
        'Content-Type': 'application/json',
    }
    json_data = {
        'prompt': prompt,
        'width': int(width),
        'height': int(height),
    }
    # Create image generation request
    response = requests.post(
        'https://api.bfl.ml/v1/flux-pro-1.1',
        headers=headers,
        json=json_data,
    )
    if response.status_code != 200:
        return f"Error: {response.status_code} {response.text}"
    request = response.json()
    request_id = request.get("id")
    if not request_id:
        return "Failed to get request ID."

    # Poll for result
    while True:
        time.sleep(0.5)
        result_response = requests.get(
            'https://api.bfl.ml/v1/get_result',
            headers={
                'accept': 'application/json',
                'x-key': api_key,
            },
            params={
                'id': request_id,
            },
        )
        if result_response.status_code != 200:
            return f"Error: {result_response.status_code} {result_response.text}"
        result = result_response.json()
        status = result.get("status")
        if status == "Ready":
            image_url = result['result'].get('sample')
            if not image_url:
                return "No image URL found in the result."
            # Fetch the image from the URL
            image_response = requests.get(image_url)
            if image_response.status_code != 200:
                return f"Failed to fetch image: {image_response.status_code}"
            image = Image.open(BytesIO(image_response.content))
            return image
        elif status == "Failed":
            return "Image generation failed."
        elif status == "Pending":
            pass  # Continue polling
        else:
            return f"Unexpected status: {status}"

# Create Gradio interface with sliders for width and height
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(
            label="Prompt",
            placeholder="Describe the image you want to generate",
            lines=3,
        ),
        gr.Slider(
            label="Width",
            minimum=256,
            maximum=1440,
            step=32,
            value=1024,
        ),
        gr.Slider(
            label="Height",
            minimum=256,
            maximum=1440,
            step=32,
            value=768,
        ),
        gr.Textbox(
            label="API Key",
            type="password",
            placeholder="Enter your BFL API key",
        ),
    ],
    outputs="image",
    title="BFL Image Generation Demo",
    description="""
    Generate images using the BFL FLUX 1.1 [pro] model.

    **Instructions:**
    1. Enter a descriptive prompt for the image you want to generate.
    2. Use the sliders to specify the width and height (in pixels) for the image. Values must be multiples of 32.
    3. Enter your BFL API key. You can obtain one from [BFL API Portal](https://api.bfl.ml).
    4. Click "Submit" to generate the image.

    **Note:** Ensure you have sufficient credits and adhere to the API usage limits.
    """
)

iface.launch()