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

# # Load API key from Hugging Face secrets
# api_key = os.getenv("Gemini_api_key")

# # Initialize Gemini client
# client = genai.Client(api_key=api_key)

# # Function to generate image from prompt
# def generate_image(prompt):
#     response = client.models.generate_content(
#         model="gemini-2.5-flash-image-preview",
#         contents=[prompt],
#     )
#     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 None

# # Gradio interface
# demo = gr.Interface(
#     fn=generate_image,
#     inputs=gr.Textbox(label="Enter your prompt"),
#     outputs=gr.Image(label="Generated Image"),
#     title="Image Generator",
#     description="Enter a text prompt and generate images"
# )

# if __name__ == "__main__":
#     demo.launch()




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

# Load API key
api_key = os.getenv("Gemini_api_key")

# Initialize Gemini client
client = genai.Client(api_key=api_key)

def generate_or_edit_image(prompt, uploaded_img):
    contents = []

    # Add prompt always
    contents.append(
        types.Content(
            role="user",
            parts=[types.Part(text=prompt)]
        )
    )

    # If an image is uploaded, include it as reference
    if uploaded_img is not None:
        buffered = BytesIO()
        uploaded_img.save(buffered, format="PNG")
        image_bytes = buffered.getvalue()

        contents.append(
            types.Content(
                role="user",
                parts=[
                    types.Part(
                        text="Use this image as reference while applying my changes."
                    ),
                    types.Part(
                        inline_data=types.Blob(
                            mime_type="image/png",
                            data=image_bytes
                        )
                    )
                ]
            )
        )

    try:
        response = client.models.generate_content(
            model="gemini-2.5-flash-image-preview",
            contents=contents,
        )

        # Debug log
        print("Gemini Response received.")

        for part in response.candidates[0].content.parts:
            if part.inline_data is not None:
                data = part.inline_data.data
                if isinstance(data, str):
                    image_bytes = base64.b64decode(data)
                else:
                    image_bytes = data
                image = Image.open(BytesIO(image_bytes))
                return image

        return "⚠️ No image generated. Try refining your prompt."

    except Exception as e:
        return f"❌ Error: {e}"

# Gradio app
demo = gr.Interface(
    fn=generate_or_edit_image,
    inputs=[
        gr.Textbox(label="Enter your prompt"),
        gr.Image(label="Upload an image (optional)", type="pil")
    ],
    outputs=gr.Image(label="Generated/Edited Image"),
    title="Image Generator & Editor",
    description=(
        "Enter a text prompt to generate a new image.\n"
        "If you upload an image, Gemini will use it as a reference and apply your requested changes."
    )
)

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