##Importing All the modules
import streamlit as st
import os
from PIL import Image
import google.generativeai as genai
from dotenv import load_dotenv


# Load all environment Variables
load_dotenv()

##Configuring the api key
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

## Function to load Gemini Vison Pro Vision Model and Get response

def get_gemini_response(input,image,prompt):

    ##Loading the desired Model
    model= genai.GenerativeModel("gemini-pro-vision")
    response=model.generate_content([input,image[0],prompt])
    return response.text

## Function to extract data from Image Uploaded
def input_image_setup(uploaded_file):
    # Check if a file has been uploaded
    if uploaded_file is not None:
        # Read the file into bytes
        bytes_data = uploaded_file.getvalue()

        image_parts = [
            {
                "mime_type": uploaded_file.type,  # Get the mime type of the uploaded file
                "data": bytes_data
            }
        ]
        return image_parts
    else:
        raise FileNotFoundError("No file uploaded")
    
#       Initializing our Streamlit Prompt

st.set_page_config(page_title="Medical Prescription Extractor")

st.header("Prescription Analyzer")
input=st.text_input("Optional Hint for the Model (e.g., Extract dosage): ",key="input")
uploaded_file = st.file_uploader("Choose a Prescription Image...", type=["jpg", "jpeg", "png"])
image=""   
if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Prescription.", use_column_width=True)


submit=st.button("Analyze Prescription")

##  DEFINING A SYSTEM PROMPT

input_prompt = "You are a licensed medical professional trained to interpret prescriptions. You will receive prescription images as input and you will have to extract key information like dosage, strength, frequency, and medication name."

if submit:
    image_data = input_image_setup(uploaded_file)
    response = get_gemini_response(input,image_data,input_prompt)

    st.subheader("Analysis by your model Gemini Pro Vision: ")
    st.write(response)

    st.balloons()