import streamlit as st
import os
from dotenv import load_dotenv

# Load environment variables FIRST
load_dotenv()

from rag import get_answer

# --- PAGE CONFIGURATION ---
st.set_page_config(page_title="Sagar's Personal Assistant", page_icon="🌟")

# Title
st.title("🌟 Sagar's Personal Assistant")
st.write("Hey there! 👋 I'm Sagar's personal assistant! Feel free to ask me anything about Sagar. I'm here to help! 😊")

# Check if API key is configured
if not os.getenv("GOOGLE_API_KEY"):
    st.error("⚠️ GOOGLE_API_KEY is missing! Please add it to your environment variables or .env file.")
    st.stop()

# Initialize Chat History
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display Chat History
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# User Input
if prompt := st.chat_input("What would you like to know?"):
    # Add user message to history
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Generate Response
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            try:
                response = get_answer(prompt)
                st.markdown(response)
                # Add assistant message to history
                st.session_state.messages.append({"role": "assistant", "content": response})
            except Exception as e:
                error_msg = f"❌ Error: {str(e)}"
                st.error(error_msg)
                st.session_state.messages.append({"role": "assistant", "content": error_msg})