import streamlit as st
from PIL import Image
from google import genai
import json
from io import BytesIO


        
def edit_image_with_ai(image, description):
    with open('app_config.json') as config_file:
        config = json.load(config_file)
        api_key = config["nano-banana"]["api_key"]

    client = genai.Client(api_key=api_key)

    prompt = description

    response = client.models.generate_content(
        model="gemini-2.5-flash-image-preview",
        contents=[prompt, image],
    )

    text_response = None
    image_response = None

    for part in response.candidates[0].content.parts:
        if part.text is not None:
            text_response = part.text
        
        if part.inline_data is not None:
            image_response = Image.open(BytesIO(part.inline_data.data))
    
    return text_response, image_response
        
def ai_image_editor():
    st.markdown("<h1 style='text-align: center; color: #451002;'>AI Image Editor🖼️</h1>", unsafe_allow_html=True)
    st.markdown("<h5 style='text-align: center;'> Edit and enhance your images with the power of AI! 🎨 </h3>", unsafe_allow_html=True)
    st.info("Powered by the online [Gemini Image Generation](https://ai.google.dev/gemini-api/docs/image-generation) API!")
    # user upload an image and provide a description of the edit they want
    uploaded_image = st.file_uploader("Upload an image to edit", type=["png", "jpg", "jpeg"])
    if uploaded_image is not None:
        with st.spinner("Loading image..."):
            # display the uploaded image
            image = Image.open(uploaded_image)
            st.image(image, width= 300)
            st.success("Image uploaded successfully!")

        edit_description = st.text_area("Describe the edits you want to make to the image")
        if st.button("AI Editing"):
            if edit_description.strip() == "":
                st.error("Please provide a description of the edits you want to make.")
            else:
                with st.spinner("Editing image..."):
                    # Call the AI image editing function (to be implemented)
                    text_response, image_response = edit_image_with_ai(image, edit_description)
                    if text_response:
                        st.info(text_response)
                    if image_response:
                        st.image(image_response, width=300)