import streamlit as st
import os

from embedding import load_embeddings
from vectorstore import load_or_build_vectorstore
from chain_setup import build_conversational_chain

def main():
    st.title("💬 المحادثة التفاعلية - إدارة البيانات وحماية البيانات الشخصية")

    # Apply RTL custom CSS for right-to-left text alignment
    st.markdown(
        """
        <style>
            .rtl {
                direction: rtl;
                text-align: right;
            }
        </style>
        """,
        unsafe_allow_html=True
    )

    # Paths and constants
    local_file = "Policies001.pdf"
    index_folder = "faiss_index"

    # Step 1: Load Arabic Embeddings
    embeddings = load_embeddings()

    # Step 2: Build or load the VectorStore
    vectorstore = load_or_build_vectorstore(local_file, index_folder, embeddings)

    # Step 3: Build the Conversational Retrieval Chain
    qa_chain = build_conversational_chain(vectorstore)

    # Step 4: Session State for UI Chat
    if "messages" not in st.session_state:
        st.session_state["messages"] = [
            {"role": "assistant", "content": "👋 مرحبًا! اسألني أي شيء عن إدارة البيانات وحماية البيانات الشخصية!"}
        ]

    # Display existing messages with RTL styling
    for msg in st.session_state["messages"]:
        with st.chat_message(msg["role"]):
            st.markdown(f'<div class="rtl">{msg["content"]}</div>', unsafe_allow_html=True)

    # Step 5: Chat Input
    user_input = st.chat_input("اكتب سؤالك هنا")

    if user_input:
        # a) Display user message
        st.session_state["messages"].append({"role": "user", "content": user_input})
        with st.chat_message("user"):
            st.markdown(f'<div class="rtl">{user_input}</div>', unsafe_allow_html=True)

        # b) Run chain
        response = qa_chain({"question": user_input})
        answer = response["answer"]

        # c) Display assistant response
        st.session_state["messages"].append({"role": "assistant", "content": answer})
        with st.chat_message("assistant"):
            st.markdown(f'<div class="rtl">{answer}</div>', unsafe_allow_html=True)

if __name__ == "__main__":
    main()