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

# Load environment variables
load_dotenv()
API_KEY = os.getenv("GOOGLE_API_KEY")

# Configure Gemini API
genai.configure(api_key=API_KEY)
model = genai.GenerativeModel('gemini-pro')

# Streamlit Page Config
st.set_page_config(page_title="Enhanced Gemini Q&A App", layout="wide")

# Initialize Chat History in Session State with the Correct Format
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []  # Empty list for chat history

# Function to get response from Gemini API
def get_gemini_response(question):
    # Reformat the session chat history to match the API requirements
    formatted_history = []
    for chat in st.session_state.chat_history:
        if "user" in chat:
            formatted_history.append({"parts": [{"text": chat["user"]}], "role": "user"})
        else:
            formatted_history.append({"parts": [{"text": chat["bot"]}], "role": "bot"})

    chat = model.start_chat(history=formatted_history)  # Use the correctly formatted history
    try:
        response = chat.send_message(question, stream=True)
        full_response = ""
        for chunk in response:
            full_response += chunk.text + " "
        st.session_state.chat_history.append({"user": question, "bot": full_response})  # Append to chat history
        return full_response
    except Exception as e:
        return f"Error: {str(e)}"

# UI Layout
st.title("🤖 Gemini AI - Interactive Q&A")
st.write("Ask me anything and I'll try to answer!")

# Sidebar for Settings
with st.sidebar:
    st.header("Settings")
    if st.button("Clear Chat History"):
        st.session_state.chat_history = []
        st.success("Chat history cleared!")

# User Input
user_input = st.text_input("Your Question:", placeholder="Type your question here...")
submit = st.button("Ask Gemini")

# Response Handling
if submit and user_input:
    if not API_KEY:
        st.error("API Key is missing. Please check your .env file.")
    else:
        with st.spinner("Thinking..."):
            response = get_gemini_response(user_input)
            st.subheader("Response:")
            st.markdown(response)  # Render response in Markdown format

# Display Chat History
if st.session_state.chat_history:
    st.subheader("Chat History")
    for chat in st.session_state.chat_history:
        st.markdown(f"**You:** {chat['user']}")
        st.markdown(f"**Gemini:** {chat['bot']}")
        st.write("---")