import gradio as gr
from gradio_client import Client, handle_file
import re

# Instantiate a Client object from gradio_client pointing to the 'selfit-camera/Omni-Image-Editor' Space.
client = Client("selfit-camera/Omni-Image-Editor")

# Define a Python function for text-to-image generation
def generate_image(prompt):
    """
    Generate an image from a text prompt using the Omni Image Editor API.
    
    Args:
        prompt (str): Text description of the image to generate
        
    Returns:
        str: URL of the generated image or error message
    """
    try:
        # Call the client.predict() method with the user's prompt, aspect_ratio='16:9', and api_name='/text_to_image_interface'.
        result = client.predict(
            prompt=prompt,
            aspect_ratio="16:9",
            api_name="/text_to_image_interface"
        )

        # The predict method returns a tuple. The first element of this tuple is an HTML string containing the image.
        # Extract the image URL from this HTML string.
        html_string = result[0]
        match = re.search(r"src='([^']+)'", html_string)
        if match:
            image_url = match.group(1)
            return image_url
        else:
            # Handle cases where the URL might not be found
            return "https://via.placeholder.com/400x200?text=Error:Image+Not+Found"
    except Exception as e:
        return f"Error generating image: {str(e)}"

# Define a Python function for image editing
def edit_image(input_image, edit_prompt):
    """
    Edit an image based on a text prompt using the Omni Image Editor API.
    
    Args:
        input_image (str): Path to the image file or image object
        edit_prompt (str): Text description of the edits to apply
        
    Returns:
        str: URL of the edited image or error message
    """
    try:
        if input_image is None:
            return "Please upload an image first"
        
        # Use handle_file to properly handle the image upload
        result = client.predict(
            input_image=handle_file(input_image),
            prompt=edit_prompt,
            api_name="/edit_image_interface"
        )
        
        # Extract the image URL from the HTML response
        if isinstance(result, tuple) and len(result) > 0:
            html_string = result[0]
            match = re.search(r"src='([^']+)'", html_string)
            if match:
                image_url = match.group(1)
                return image_url
            else:
                return "https://via.placeholder.com/400x200?text=Error:Image+Not+Found"
        else:
            return str(result)
            
    except Exception as e:
        return f"Error editing image: {str(e)}"

# Create a Gradio application using gr.Blocks for more granular control.
with gr.Blocks(
    title='Omni Image Editor with Gradio',
    theme=gr.themes.Soft()
) as demo:
    gr.Markdown("# Omni Image Editor Studio")
    gr.Markdown("Generate images from text descriptions or edit existing images with AI-powered tools.")
    
    with gr.Tabs():
        # Text-to-Image Tab
        with gr.TabItem("Text to Image Generator"):
            gr.Markdown("### Generate Images from Text")
            gr.Markdown("Describe the image you want to generate in detail for best results.")
            
            with gr.Row():
                with gr.Column():
                    prompt_input = gr.Textbox(
                        label='Image Description',
                        placeholder='e.g., A futuristic city at sunset with flying cars, neon lights, cyberpunk style, high quality',
                        lines=3
                    )
                    generate_btn = gr.Button("🎨 Generate Image", variant="primary")
            
            generated_image = gr.Image(label='Generated Image', type='filepath')
            
            # Bind the generate_image function to the button click event
            generate_btn.click(
                fn=generate_image,
                inputs=[prompt_input],
                outputs=[generated_image]
            )
        
        # Image Editing Tab
        with gr.TabItem("Image Editor"):
            gr.Markdown("### Edit Images with AI")
            gr.Markdown("Upload an image and describe the changes you want to make.")
            
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(
                        label='Upload Image',
                        type='filepath'
                    )
                    
            with gr.Row():
                with gr.Column():
                    edit_prompt = gr.Textbox(
                        label='Edit Instructions',
                        placeholder='e.g., Change the sky to sunset colors, add stars, increase contrast',
                        lines=3
                    )
                    edit_btn = gr.Button("✨ Edit Image", variant="primary")
            
            edited_image = gr.Image(label='Edited Image', type='filepath')
            
            # Bind the edit_image function to the button click event
            edit_btn.click(
                fn=edit_image,
                inputs=[input_image, edit_prompt],
                outputs=[edited_image]
            )

# Launch the Gradio application.
if __name__ == "__main__":
    demo.launch(share=True)