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


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


## Function to load OpenAI model and get respones

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

    if response.candidates and response.candidates[0].safety_ratings:
        for rating in response.candidates[0].safety_ratings:
            if rating.category in ("HARM_CATEGORY_DEROGATORY", "HARM_CATEGORY_TOXICITY", "HARM_CATEGORY_VIOLENCE",
                                  "HARM_CATEGORY_SEXUALLY_SUGGESTIVE", "HARM_CATEGORY_MEDICAL", "HARM_CATEGORY_DANGEROUS"):
                return f"Response blocked due to safety concerns: {rating.category}"
    if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
        return response.candidates[0].content.parts[0].text
    else:
        return "No valid response generated."

    

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")


# Streamlit App

# Set Page Configuration
st.set_page_config(
    page_title="Image Reading Gemini",
    page_icon=":mag:",
    layout="wide"
)

# Custom Theme
st.markdown(
    """
    <style>
    .reportview-container {
        background: linear-gradient(to right, #4CAF50, #2196F3); 
    }
    .sidebar .sidebar-content {
        background: #f0f0f5;
    }
    .stButton button {
        background-color: #4CAF50;
        color: white;
    }
    .stTextInput input {
        background-color: #e8f0fe;
    }
    </style>
    """,
    unsafe_allow_html=True
)

st.title("Debyez Image Extractor :eyes:")

# Input Elements with Session State
if "input" not in st.session_state:
    st.session_state.input = ""
if "clear_input" not in st.session_state:  # Add a flag for clearing
    st.session_state.clear_input = False

# Callback function for the "Clear" button
def clear_input_callback():
    st.session_state.input = ""
    st.session_state.clear_input = True  # Set the clear flag
    
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
input = st.text_input("Input Prompt: ", key="input", value=st.session_state.input)


# Submit and Clear Buttons in the Same Row
submit = st.button("Submit")
clear = st.button("Clear", on_click=clear_input_callback)

# Input Prompt for Gemini
input_prompt = """
                You are an expert in understanding writings, colours and designs from an image.
                You will receive input images &
                you will have to answer questions based on the input image only
                """

if submit:
    if uploaded_file is not None:
        image_data = input_image_setup(uploaded_file)
        response = get_gemini_response(input_prompt, image_data, input)
        st.markdown(f"**The Response is:**")
        st.write(response)
        st.markdown("---")
    else:
        st.warning("Please upload an image first.")

if st.session_state.clear_input:
    st.session_state.clear_input = False 
    st.experimental_rerun()  # Refresh the UI