import streamlit as st
import tempfile
from PIL import Image
import google.generativeai as genai
import os

MODEL_ID = "gemini-2.0-flash-exp" 
api_key = os.getenv("GEMINI_API_KEY")
model_id = MODEL_ID
genai.configure(api_key=api_key)

if "model" not in st.session_state:
    st.session_state.model = genai.GenerativeModel(MODEL_ID)

model = st.session_state.model
chat = model.start_chat()
mime_type = "image/jpeg"  # Use a consistent MIME type for images

def main():
    st.title("Image Processing with AI")
    img_file_buffer = None
    image_path = ""

    # Use the camera to capture an image
    img_file_buffer = st.camera_input("Take a picture")

    if img_file_buffer is not None:
        # Save the image to a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
            temp_file.write(img_file_buffer.read())
            image_path = temp_file.name
            st.write("Image saved to:", image_path) 
            #st.image(Image.open(image_path), caption="Captured Image", use_container_width=True)

    user_input = st.text_area("Enter your prompt:", height=200)
    prompt =f"Refer to the image. Perform the task or answer the question {user_input}"

    if st.button("Process the image"):
        if prompt:
            with st.spinner("Processing..."):
                # Upload the file with the correct MIME type
                file_data = genai.upload_file(image_path, mime_type=mime_type)
                
                # Send file and prompt to Gemini API
                response = chat.send_message(
                    [
                        prompt, 
                        file_data
                    ]
                )

                # Display Gemini response
                st.markdown(response.text)
        else:
            st.warning("Please provide a valid prompt.")


if __name__ == "__main__":
    main()