import streamlit as st
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
from langchain_google_genai import ChatGoogleGenerativeAI

# ---------------- Streamlit Page Config ----------------
st.set_page_config(page_title="🤖 AI Conversational Data Science Tutor", layout="wide")
st.title("🤖 AI Conversational Data Science Tutor")

# ---------------- Sidebar for Settings ----------------
st.sidebar.header("Settings")
google_api_key = st.sidebar.text_input("Enter your Google API Key", type="password")

# ---------------- Initialize Model ----------------
if google_api_key:
    llm = ChatGoogleGenerativeAI(
        model="gemini-1.5-flash",
        google_api_key=google_api_key,
        temperature=0.3,
    )

    # Conversation memory
    if "memory" not in st.session_state:
        st.session_state.memory = ConversationBufferMemory(return_messages=True)

    conversation = ConversationChain(
        llm=llm,
        memory=st.session_state.memory,
        verbose=True,
    )

    # ---------------- Chat Interface ----------------
    if "messages" not in st.session_state:
        st.session_state.messages = []

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

    user_input = st.chat_input("Ask your Data Science question...")
    if user_input:
        # Store user message
        st.session_state.messages.append({"role": "user", "content": user_input})
        with st.chat_message("user"):
            st.markdown(user_input)

        # Get AI response with memory
        response = conversation.predict(input=user_input)

        st.session_state.messages.append({"role": "assistant", "content": response})
        with st.chat_message("assistant"):
            st.markdown(response)

else:
    st.warning("Please enter your Google API Key in the sidebar to continue.")