import gradio as gr
import fal_client
import time
import requests
from io import BytesIO
from PIL import Image
import base64

def generate_image(api_key, prompt):
    try:
        # Set the API key provided by the user
        fal_client.api_key = api_key

        # Default parameters
        image_size = "landscape_4_3"
        num_images = 1
        enable_safety_checker = True
        safety_tolerance = "2"
        sync_mode = True  # Enable sync_mode to get the image directly

        handler = fal_client.submit(
            "fal-ai/flux-pro/v1.1",
            arguments={
                "prompt": prompt,
                "image_size": image_size,
                "num_images": num_images,
                "enable_safety_checker": enable_safety_checker,
                "safety_tolerance": safety_tolerance,
                "sync_mode": sync_mode,
            },
        )

        # Since sync_mode is True, handler.get() will wait for the result
        result = handler.get()

        # For debugging, print the result
        print("API Response Result:", result)

        if not result.get("images"):
            return None, "No images were generated."

        image_info = result["images"][0]

        # Check if image data is in the 'data' field (base64 encoded)
        if 'data' in image_info:
            # Image data is base64 encoded
            image_data = image_info['data']
            image_bytes = base64.b64decode(image_data)
            image = Image.open(BytesIO(image_bytes))
            return image, None
        elif 'url' in image_info:
            image_url = image_info['url']
            print("Image URL:", image_url)
            # Download the image
            response = requests.get(image_url)
            if response.status_code == 200:
                image = Image.open(BytesIO(response.content))
                return image, None
            else:
                return None, f"Failed to download the generated image. Status code: {response.status_code}"
        else:
            return None, "Image data not found in API response."
    except Exception as e:
        # For debugging, print the exception
        print("Exception occurred:", e)
        return None, str(e)

# Set up the Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="API Key", placeholder="Enter your API key here...", type="password"),
        gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
    ],
    outputs=[
        gr.Image(label="Generated Image"),
        gr.Textbox(label="Error Message")
    ],
    title="FLUX1.1 [pro] Image Generation",
    description="Generate an image using the FLUX1.1 [pro] model by providing your API key and a text prompt."
)

iface.launch()