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

# Load environment variables from .env.
load_dotenv()

# Configure Google Generative AI
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# Function to get response from Gemini text model
def get_gemini_text_response(input_text):
    model = genai.GenerativeModel('gemini-pro')
    response = model.generate_content(input_text)
    return response.text if response else "Please provide a text prompt."

# Function to get response from Gemini vision model
def get_gemini_image_response(image, input_text=None):
    model = genai.GenerativeModel('gemini-pro-vision')
    if input_text:
        response = model.generate_content([input_text, image])
    else:
        response = model.generate_content(image)
    return response.text if response else "Please provide a image."

# Initialize Streamlit app
st.set_page_config(page_title="Gemini Q&A and Image Analysis", page_icon="🤖")

# Sidebar for additional options
st.sidebar.title("About")
st.sidebar.info("This is a Q&A chatbot powered by Google's Gemini AI model.")
st.sidebar.title("Instructions")
st.sidebar.info("1. Choose either 'Q&A Chatbot' or 'Image Analysis' from the options.\n"
                "2. Enter your question/prompt or upload an image.\n"
                "3. Click the 'Ask' button to get the response.\n"
                "4. The response will be displayed below.")

# Main interface
st.title("Gemini Q&A and Image Analysis Chatbot")
st.write("Select an option to either ask a question or analyze an image.")

# User selection
option = st.selectbox("Choose an option:", ("Q&A Chatbot", "Image Analysis"))

# Input fields based on user selection
if option == "Q&A Chatbot":
    input_text = st.text_input("Enter your question or prompt:", key="input_text")
    if st.button("Ask"):
        with st.spinner("Generating response..."):
            response = get_gemini_text_response(input_text)
        st.subheader("Response:")
        st.write(response)

elif option == "Image Analysis":
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    image_input_text = st.text_input("Enter your question or prompt related to the image (optional):", key="image_input_text")
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image.", use_column_width=True)
        if st.button("Analyze"):
            with st.spinner("Generating response..."):
                response = get_gemini_image_response(image, input_text=image_input_text)
            st.subheader("Response:")
            st.write(response)