import gradio as gr
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import base64
from dotenv import load_dotenv
import os
import numpy as np

# Load environment variables
load_dotenv()

# Initialize the client
client = genai.Client(api_key=os.environ.get('GEMINI_API_KEY'))
print('done')

def generate_image(prompt):
    """Generate an image based on text prompt"""
    try:
        response = client.models.generate_content(
            model="gemini-2.0-flash-exp-image-generation",
            contents=prompt,
            config=types.GenerateContentConfig(
                response_modalities=['Text', 'Image']
            )
        )
        
        for part in response.candidates[0].content.parts:
            if part.inline_data is not None:
                image = Image.open(BytesIO(part.inline_data.data))
                return image
        return "No image generated"
    
    except Exception as e:
        return f"Error generating image: {str(e)}"

def edit_uploaded_image(images, edit_prompt):
    """Edit the uploaded images based on the prompt"""
    try:
        if not images:
            return "No images uploaded"
        
        edited_images = []
        for image in images:
            # Convert Gradio image input (numpy array) to PIL Image
            pil_image = Image.fromarray(image) if isinstance(image, np.ndarray) else image # Handle both PIL Image and numpy array inputs

            # Prepare content with text prompt and image
            text_input = f"Hi, this is an image. {edit_prompt}"
            contents = [text_input, pil_image]
        
            # Generate edited image using Gemini API
            response = client.models.generate_content(
                model="gemini-2.0-flash-exp-image-generation",
                contents=contents,
                config=types.GenerateContentConfig(
                    response_modalities=['Text', 'Image']
                )
            )
        
            for part in response.candidates[0].content.parts:
                if part.inline_data is not None:
                    edited_image = Image.open(BytesIO(part.inline_data.data))
                    edited_images.append(edited_image)
                    break  # Exit the inner loop once an image is found
            else:
                 edited_images.append("No edited image generated")  # Append a message if no image part found in the response
        return edited_images
    
    except Exception as e:
        return f"Error editing image: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Image Generation and Editing App") as app:
    gr.Markdown("# Image Generation and Editing App")
    gr.Markdown("Generate images from text prompts or edit uploaded images")
    
    with gr.Tab("Generate Image"):
        prompt_input = gr.Textbox(
            label="Enter your prompt",
            placeholder="e.g., 'A pig with wings and a top hat flying over a futuristic city'"
        )
        generate_button = gr.Button("Generate")
        generated_output = gr.Image(label="Generated Image")
        
    with gr.Tab("Edit Uploaded Image"):
        image_upload = gr.Image(label="Upload your image(s)", image_mode="RGBA", type="numpy", height=300, width=300,  multiple=True)  # Allow multiple images
        edit_prompt = gr.Textbox(
            label="Edit prompt",
            placeholder="e.g., 'Add a llama next to the subject'"
        )
        edit_button = gr.Button("Edit Image")
        edited_output = gr.Gallery(label="Edited Images")  # Use Gallery for multiple images
    
    # Connect the components with functions
    generate_button.click(
        fn=generate_image,
        inputs=prompt_input,
        outputs=generated_output
    )
    
    edit_button.click(
        fn=edit_uploaded_image,
        inputs=[image_upload, edit_prompt],
        outputs=edited_output
    )

print('done','******************************')
# Launch the app
app.launch()