from dotenv import load_dotenv
load_dotenv()
import streamlit as st
import os
import google.generativeai as genai
from PIL import Image

genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
# load gemini model
model=genai.GenerativeModel("gemini-1.5-flash")
def get_gemini_response(input,image,prompt):
    response=model.generate_content([input,image[0],prompt])
    return response.text
def input_image_setup(uploaded_img):
    if uploaded_img is not None:
        bytes_data = uploaded_img.getvalue()
        image_parts=[
            {
                "mime_type": uploaded_img.type,
                "data": bytes_data
            }
        ]
        return image_parts
    else:
        raise FileNotFoundError("Image not found")


st.set_page_config(page_title="Invoice extractor", page_icon="🔮")
st.title("Invoice Extractor using LLM")
st.write("Upload your invoice and we will give you all the information we can based on your query")
input = st.text_input("Ask a question", key="input")
uploaded_img = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
image=""
if uploaded_img is not None:
    image = Image.open(uploaded_img)
    st.image(image, caption="Uploaded Image.", use_column_width=True)
submit=st.button("Submit")

input_prompt="""
You are expert in understanding invoices. We will show you an invoice and you have to answer the following questions based on the invoice:

"""
if submit:
    image_data=input_image_setup(uploaded_img)
    response=get_gemini_response(input_prompt,image_data,input)
    st.subheader("Response:")
    st.write(response)