import streamlit as st
from app import rag_query, memory, process_feedback


st.title("🛡️ RAG Chatbot MB Ageas Life 🛡️")

# Initialize session history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Hiển thị lại tin nhắn cũ
for i, message in enumerate(st.session_state.messages):
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        if message["role"] == "assistant":
            col1, col2 = st.columns([1, 15])
            with col1:
                if st.button("👍", key=f"thumbs_up_{i}"):
                    process_feedback(st.session_state.messages[i-1]["content"], message["content"], True)
            with col2:
                if st.button("👎", key=f"thumbs_down_{i}"):
                    process_feedback(st.session_state.messages[i-1]["content"], message["content"], False)

# Nhận input người dùng
if prompt := st.chat_input("Ask me any question related to MBAL"):
    # Hiển thị tin nhắn người dùng
    st.chat_message("user").markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})
    memory.chat_memory.add_user_message(prompt)

    # Gọi hàm RAG để trả lời
    response = rag_query(prompt)

    # Hiển thị tin nhắn chatbot
    with st.chat_message("assistant"):
        st.markdown(response)
    st.session_state.messages.append({"role": "assistant", "content": response})
    memory.chat_memory.add_ai_message(response)

    st.rerun()

# Sidebar for additional controls
with st.sidebar:
    st.header("Options")
    if st.button("Clear Chat History"):
        st.session_state.messages = []
        st.rerun()