import gradio as gr
import base64
import io
from PIL import Image
import json
import os
import asyncio
from google import genai
from google.genai import types

# Function to convert PIL Image to bytes
def pil_to_bytes(img, format="PNG"):
    img_byte_arr = io.BytesIO()
    img.save(img_byte_arr, format=format)
    return img_byte_arr.getvalue()

# Function to load image as base64
async def load_image_base64(img):
    if isinstance(img, str):
        # If image is a URL or file path, load it
        raise ValueError("URL loading not implemented in this version")
    else:
        # If image is already a PIL Image
        return pil_to_bytes(img)

# Main function to generate edited image using Gemini
async def generate_image_gemini(prompt, image, api_key, temperature=0.4):
    SAFETY_SETTINGS = {
        types.HarmCategory.HARM_CATEGORY_HARASSMENT: types.HarmBlockThreshold.BLOCK_NONE,
        types.HarmCategory.HARM_CATEGORY_HATE_SPEECH: types.HarmBlockThreshold.BLOCK_NONE,
        types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: types.HarmBlockThreshold.BLOCK_NONE,
        types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: types.HarmBlockThreshold.BLOCK_NONE,
    }
    
    try:
        # Initialize Gemini client with API key
        client = genai.Client(api_key=api_key)
        
        # Convert PIL image to bytes
        image_bytes = await load_image_base64(image)
        
        contents = []

        # Add the image to the contents
        contents.append(
            types.Content(
                role="user",
                parts=[
                    types.Part.from_bytes(
                        data=image_bytes,
                        mime_type="image/png",
                    )
                ],
            )
        )
        
        # Add the prompt to the contents
        edit_prompt = f"Edit this image: {prompt}"
        contents.append(
            types.Content(
                role="user",
                parts=[
                    types.Part.from_text(text=edit_prompt),
                ],
            )
        )

        response = await client.aio.models.generate_content(
            model="gemini-2.0-flash-exp",
            contents=contents,
            config=types.GenerateContentConfig(
                temperature=temperature,
                safety_settings=[
                    types.SafetySetting(
                        category=category, threshold=threshold
                    ) for category, threshold in SAFETY_SETTINGS.items()
                ],
                response_modalities=['Text', 'Image']
            )
        )
        
        edited_images = []
        for part in response.candidates[0].content.parts:
            if part.inline_data is not None:
                image_bytes = part.inline_data.data
                edited_images.append(image_bytes)
                
        # Convert the first returned image bytes to PIL image
        if edited_images:
            result_image = Image.open(io.BytesIO(edited_images[0]))
            return result_image
        else:
            return None
            
    except Exception as e:
        print(f"Google GenAI client failed with error: {e}")
        return None

# Function to process the image edit
def process_image_edit(image, prompt, api_key, image_history, temperature):
    if not image or not prompt or not api_key:
        return None, image_history, "Please provide an image, prompt, and API key"
    
    # Store current image in history if not empty
    if image is not None and image_history is None:
        image_history = []
    if image is not None:
        image_history.append(image)
    
    # Run the async function to edit the image
    try:
        edited_image = asyncio.run(generate_image_gemini(prompt, image, api_key, temperature))
        if edited_image:
            return edited_image, image_history, "Image edited successfully"
        else:
            return image, image_history, "Failed to edit image. Please try again."
    except Exception as e:
        return image, image_history, f"Error: {str(e)}"

# Function to undo the last edit
def undo_edit(image_history):
    if image_history and len(image_history) > 1:
        # Remove current image
        image_history.pop()
        # Return the previous image
        return image_history[-1], image_history, "Reverted to previous image"
    else:
        return None, [], "No previous version available"

# Function to set output image as input for continuous editing
def continue_editing(output_image):
    if output_image is not None:
        return output_image, "Ready to continue editing the current image"
    else:
        return None, "No edited image available to continue editing"

# Create Gradio UI
def create_ui():
    with gr.Blocks(title="Gemini Image Editor") as app:
        gr.Markdown("# Gemini Image Editor")
        gr.Markdown("Upload an image, enter a description of the edit you want, and let Gemini do the rest!")
        
        # Store image history in state
        image_history = gr.State([])
        
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="pil", label="Upload Image")
                prompt = gr.Textbox(label="Edit Description", placeholder="Describe the edit you want...")
                api_key = gr.Textbox(label="Gemini API Key", placeholder="Enter your Gemini API key", type="password")
                
                # Replace hidden settings with accordion
                with gr.Accordion("Advanced Settings", open=False):
                    temperature = gr.Slider(
                        minimum=0.0, 
                        maximum=2.0, 
                        value=1, 
                        step=0.05, 
                        label="Temperature", 
                        info="Controls randomness in generation (0 = deterministic, 1 = creative, 2 = extreme)"
                    )
                
                with gr.Row():
                    edit_btn = gr.Button("Edit Image")
                    undo_btn = gr.Button("Undo Last Edit")
                    continue_btn = gr.Button("Continue Editing")
                
            with gr.Column():
                output_image = gr.Image(type="pil", label="Edited Image")
                status = gr.Textbox(label="Status", interactive=False)
        
        # Set up event handlers
        edit_btn.click(
            fn=process_image_edit,
            inputs=[input_image, prompt, api_key, image_history, temperature],
            outputs=[output_image, image_history, status]
        )
        
        undo_btn.click(
            fn=undo_edit,
            inputs=[image_history],
            outputs=[output_image, image_history, status]
        )
        
        # Add handler for continue editing button
        continue_btn.click(
            fn=continue_editing,
            inputs=[output_image],
            outputs=[input_image, status]
        )
    
    return app

# Launch the app
if __name__ == "__main__":
    app = create_ui()
    app.launch()