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 [Gemini Image Generation](https://ai.google.dev/gemini-api/docs/image-generation) API! 
            ⚠️ No images or prompt descriptions are stored.""")
    reference_prompt = """
        **Sample Prompt Description:**
        
        Edit the uploaded image to create a professional profile photo of the person in the picture. Focus on the following:
        - Outfit: Transform the current attire into a fitted, black professional suit. Ensure the fit is flattering and appropriate for a software engineer headshot. 
        - Facial Features & Hair: Retain the original facial features (eyes, nose, mouth, etc.) and individual hair strands as much as possible. The goal is to achieve a realistic, edited look, not an AI-generated or overly smoothed/perfected appearance. Avoid any artificial alterations to these features.
        - Pose Adjustment: Adjust the pose so that the face is facing forward. Make subtle corrections to the body posture to appear slightly straightened and more upright.
        - Framing & Composition: Crop the image as a half-body portrait. Ensure the head and shoulders are well-framed.
        - Background & Lighting: Use a neutral, non-distracting background (e.g., soft gray). Apply soft, natural lighting to create a professional and inviting look. Avoid harsh shadows or overly bright highlights.
        - Overall Style: The final image should be photorealistic and suitable for a professional profile. It should project professionalism and competence.
        """

    # user upload an image and provide a description of the edit they want
    uploaded_image = st.file_uploader("Upload an image and write prompt description to edit", type=["png", "jpg", "jpeg"], help="Choose an image file (PNG, JPG, JPEG) to upload.")
    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!")
            
        with st.expander("See Example Prompt for Professional Profile Photo"):
                    st.info("Note: The example prompt is for reference only. You can provide your own description of the edits you want to make to the image. The quality of the results may vary based on the specificity of your prompt. For more information, see the [Gemini Official Prompting Guide and Strategies](https://ai.google.dev/gemini-api/docs/image-generation#prompt-guide)")
                    st.markdown(reference_prompt)
        edit_description = st.text_area("Describe the edits you want to make to the image", height="content",)
        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)
                        # Convert PIL image to bytes for download
                        img_byte_arr = BytesIO()
                        image_response.save(img_byte_arr, format="JPEG")
                        st.download_button(
                            label="Download Edited Image",
                            data=img_byte_arr.getvalue(),
                            file_name="edited_image.jpeg",
                            mime="image/jpeg"
                        )

                   
