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! ⚠️ No images or descriptions are stored.")
    # 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"], 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!")

        reference_prompt = """
            Please 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.  Change glasses to frameless fashion glasses. Change tie color to deep blue
        - 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.
        """
        with st.expander("See Example Prompt for Professional Profile Photo"):
            st.markdown(reference_prompt)
            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.         Gemini Official Prompting Guide and Strategies: https://ai.google.dev/gemini-api/docs/image-generation#prompt-guide")
        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)