import os
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import gradio as gr

# Set up your API key
API_KEY = os.getenv("GOOGLE_API_KEY")
client = genai.Client(api_key=API_KEY)

def edit_image_with_gemini(image, text_input):
    """
    Edits an image using Gemini 2.0 Flash Experimental API based on a given text prompt.
    """
    if image is None or not text_input:
        return "Please upload an image and provide an edit prompt.", None

    # Wrap the text prompt as a tuple (as shown in the docs)
    text_prompt_tuple = (text_input,)
    
    try:
        response = client.models.generate_content(
            model="gemini-2.0-flash-exp-image-generation",
            contents=[text_prompt_tuple, image],
            config=types.GenerateContentConfig(
                response_modalities=['Text', 'Image']
            )
        )
    except Exception as e:
        return f"API call error: {e}", None

    # Process the response: check for both image and text output
    for part in response.candidates[0].content.parts:
        if part.inline_data is not None:
            edited_image = Image.open(BytesIO(part.inline_data.data))
            return "Here is your edited image:", edited_image
        elif part.text is not None:
            # Optionally, print or log the text output
            print("Text output:", part.text)

    return "No image was generated. Try modifying your prompt.", None

def generate_thumbnail_prompt():
    """
    Returns a predefined optimized prompt for creating a YouTube thumbnail.
    """
    return ("Generate a bold, eye-catching YouTube thumbnail with vibrant colors, "
            "large text, and a strong contrast. Make sure it stands out and is attention-grabbing.")

# Build the Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# ✨ AI-Powered Image Editor with Gemini 2.0 Flash Experimental")
    gr.Markdown("Upload an image and describe the edits you want!")
    
    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Image")
        text_input = gr.Textbox(placeholder="Describe your edit...", label="Edit Prompt")
    
    thumbnail_master_btn = gr.Button("🎨 Thumbnail Master")
    output_text = gr.Textbox(label="Status", interactive=False)
    output_image = gr.Image(label="Edited Image")

    with gr.Row():
        submit_btn = gr.Button("Generate Edit")
        clear_btn = gr.Button("Clear")

    text_input.submit(edit_image_with_gemini, [image_input, text_input], [output_text, output_image])
    submit_btn.click(edit_image_with_gemini, [image_input, text_input], [output_text, output_image])
    clear_btn.click(lambda: (None, None), None, [output_text, output_image])
    # Use inputs=None since generate_thumbnail_prompt requires no input
    thumbnail_master_btn.click(generate_thumbnail_prompt, inputs=None, outputs=text_input)

if __name__ == "__main__":
    demo.launch(debug=True)