import streamlit as st
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_community.chat_message_histories import SQLChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
import os

API_KEY = os.getenv("Google_api_key")

template = ChatPromptTemplate(
    messages=[
        ("system", "You're a helpful data science AI chatbot. Answer only questions related to data science and what he told you within a 300-word limit."),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}") 
    ]
)

model = ChatGoogleGenerativeAI(api_key=API_KEY, model="gemini-1.5-pro")
output = StrOutputParser()
chain = template | model | output

def messages_history(session_id):
    return SQLChatMessageHistory(session_id=session_id, connection="sqlite:///sqlite.db")

conversation_chain = RunnableWithMessageHistory(
    chain, messages_history, input_message_key="input", history_messages_key="chat_history"
)

with st.sidebar:
    st.title("🤖 AI Data Science Chatbot")
    st.header("User Login")
    user_id = st.text_input("Enter your User ID:", key="user_id_input")
    # st.header(API_KEY)
if not user_id:
    st.warning("Please enter a User ID to start chatting.")
    st.stop()

if "last_user_id" not in st.session_state or st.session_state.last_user_id != user_id:
    st.session_state.chat_history = []  
    st.session_state.last_user_id = user_id  

chat_history = messages_history(user_id).messages
st.session_state.chat_history = [(msg.type, msg.content) for msg in chat_history]

st.write("Welcome! Start chatting below:")
# st.write(API_KEY)
for role, message in st.session_state.chat_history:
    st.chat_message(role).write(message)

user_input = st.chat_input("Type your message...")

if user_input:
    st.session_state.chat_history.append(("user", user_input))
    st.chat_message("user").write(user_input)

    config = {"configurable": {"session_id": user_id}}
    input_prompt = {"input": user_input} 
    response = conversation_chain.invoke(input_prompt, config=config)

    st.session_state.chat_history.append(("assistant", response))
    st.chat_message("assistant").write(response)