import streamlit as st
import os
import pathlib
import textwrap
from PIL import Image


import google.generativeai as genai


os.getenv("GOOGLE_API_KEY")
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# function to load gemini Pro model


def get_gemini_response(input, image, prompt):
    model = genai.GenerativeModel('gemini-pro-vision')
    response = model.generate_content([input, image[0], prompt])
    return response.text


def input_image_details(uploaded_file):
    if uploaded_file is not None:
        # Read the file into bytes
        bytes_data = uploaded_file.getvalue()

        image_parts = [
            {
                # Get the mime type of the up]
                "mime_type": uploaded_file.type,
                "data": bytes_data
            }
        ]

        return image_parts
    else:
        raise FileNotFoundError("No file uploaded")


# Initialize streamlit app
st.set_page_config(page_title="Data Extractor")

st.header("Data Extractor")
input = st.text_input("Input Prompt: ", key="input")

uploaded_file = st.file_uploader(
    "Choose an image of ... ", type=["jpg", "jpeg", "png"])
image = ""

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded image.", use_column_width=True)


submit = st.button("Submit")

input_prompt = """
You are an expert in understanding financial reports. We will upload a a image as financial statement
and you will have to answer any questions based on the uploaded invoice image
"""

# If submit button clicked
if submit:
    image_data = input_image_details(uploaded_file)
    reponse = get_gemini_response(
        input_prompt,  image_data, input)
    st.subheader("The response is")
    st.write(reponse)