import streamlit as st
from google import genai
import os

# Access the API key from environment variables
api_key = os.getenv("GOOGLE_API_KEY")
# Check if the API key is present
if not api_key:
    raise EnvironmentError("GOOGLE_API_KEY not set in environment variables.")

# Set the environment variable for the Google API client
os.environ["GOOGLE_API_KEY"] = api_key
# Initialize the client
client = genai.Client()
MODEL_ID = "gemini-2.0-flash"

# Function to chat with Gemini
def chat_with_gemini(message, history):
    try:
        response = client.models.generate_content(
            model=MODEL_ID,
            contents=message,
            config={"tools": [{"google_search": {}}]},
        )
        return response.text
    except Exception as e:
        return f"Error: {str(e)}"

# Streamlit app
st.title("Gemini Chatbot (with Google Search)")
st.write("Ask anything and get responses powered by Gemini 2.0 Flash.")

# Initialize chat history
if 'history' not in st.session_state:
    st.session_state['history'] = []

# Input field for the user to type their message
user_input = st.text_input("You:", "")

# Button to send the message
if st.button("Send"):
    if user_input:
        # Add user message to history
        st.session_state.history.append({"role": "user", "message": user_input})
        
        # Get response from Gemini
        gemini_response = chat_with_gemini(user_input, st.session_state.history)
        
        # Add Gemini response to history
        st.session_state.history.append({"role": "assistant", "message": gemini_response})
        
        # Display the chat history
        for chat in st.session_state.history:
            if chat["role"] == "user":
                st.chat_message("user").markdown(chat["message"])
            else:
                st.chat_message("assistant").markdown(chat["message"])