import streamlit as st
import os
import tempfile

from pdf_processor import process_uploaded_pdf
from embeddings import create_embeddings, setup_vector_store, search_vector_store
from llm_handler import generate_answer, summarize_conversation

st.set_page_config(
    page_title="Arabic Document Chatbot",
    page_icon="🇸🇦",
    layout="centered"
)

# --- SESSION STATE SETUP ---
if "display_history" not in st.session_state:
    st.session_state.display_history = []  # now stores (question, answer, language) tuples
if "recent_history" not in st.session_state:
    st.session_state.recent_history = []
if "conversation_summary" not in st.session_state:
    st.session_state.conversation_summary = ""
if "document_processed" not in st.session_state:
    st.session_state.document_processed = False
if "current_file_name" not in st.session_state:
    st.session_state.current_file_name = None


def render_message(text, is_arabic):
    """
    Renders a chat message's text content, applying right-to-left
    alignment only if THIS specific message was in Arabic.

    We can't rely on a single page-wide CSS rule for this, because
    that would retroactively flip the alignment of OLDER messages
    that were in a different language when they were sent. Instead,
    each message renders independently based on the language it
    was actually written in.
    """
    if is_arabic:
        st.markdown(
            f'<div style="direction: rtl; text-align: right;">{text}</div>',
            unsafe_allow_html=True
        )
    else:
        st.write(text)


# --- SIDEBAR: FILE UPLOAD + SETTINGS ---
with st.sidebar:
    st.header("📄 Document")

    uploaded_file = st.file_uploader("Upload an Arabic PDF", type=["pdf"])

    if uploaded_file is not None:

        if st.session_state.current_file_name != uploaded_file.name:

            with st.spinner("Processing PDF with OCR... this may take a few minutes"):
                try:
                    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
                        tmp_file.write(uploaded_file.read())
                        tmp_path = tmp_file.name

                    chunks = process_uploaded_pdf(tmp_path)

                    if not chunks:
                        raise ValueError("No text could be extracted from this PDF.")

                    embeddings = create_embeddings(chunks)
                    setup_vector_store(chunks, embeddings, collection_name="uploaded_doc")

                    os.unlink(tmp_path)

                    st.session_state.document_processed = True
                    st.session_state.current_file_name = uploaded_file.name
                    st.session_state.display_history = []
                    st.session_state.recent_history = []
                    st.session_state.conversation_summary = ""

                    st.success(f"'{uploaded_file.name}' processed successfully!")

                except Exception as e:
                    st.error(f"Failed to process the PDF: {str(e)}")
                    st.session_state.document_processed = False

    if st.session_state.document_processed:
        st.info(f"📄 Active document: {st.session_state.current_file_name}")

    st.divider()

    st.header("🌐 Language")
    language_choice = st.radio(
        "Choose your language / اختر لغتك",
        options=["English", "العربية"],
        horizontal=False,
        label_visibility="collapsed"
    )
    selected_language = "english" if language_choice == "English" else "arabic"

    # Note: RTL styling is now applied PER MESSAGE via render_message()
    # below, not as a single global rule here - this avoids flipping
    # the alignment of older messages sent in a different language.

    st.divider()

    if st.session_state.document_processed:
        if st.button("🗑️ Clear conversation"):
            st.session_state.display_history = []
            st.session_state.recent_history = []
            st.session_state.conversation_summary = ""
            st.rerun()


# --- MAIN AREA: CHAT INTERFACE ---
st.title("🇸🇦 Arabic Document Chatbot")

if not st.session_state.document_processed:
    st.info("👈 Please upload a PDF in the sidebar to begin chatting.")

else:
    st.write("Ask questions about the uploaded document in Arabic or English")

    # --- DISPLAY FULL CHAT HISTORY ---
    # Each message renders with the alignment matching the language
    # it was actually sent in, not the current sidebar selection
    for past_question, past_answer, past_language in st.session_state.display_history:
        is_arabic = (past_language == "arabic")

        with st.chat_message("user", avatar="🧑"):
            render_message(past_question, is_arabic)
        with st.chat_message("assistant", avatar="🇸🇦"):
            render_message(past_answer, is_arabic)

    # --- CHAT INPUT ---
    user_question = st.chat_input(
        "Type your question..." if selected_language == "english" else "اكتب سؤالك..."
    )

    if user_question and user_question.strip():

        is_arabic_now = (selected_language == "arabic")

        with st.chat_message("user", avatar="🧑"):
            render_message(user_question, is_arabic_now)

        with st.chat_message("assistant", avatar="🇸🇦"):
            with st.spinner("Thinking..." if selected_language == "english" else "جاري التفكير..."):

                answer = None

                try:
                    retrieved_chunks = search_vector_store(
                        user_question,
                        collection_name="uploaded_doc",
                        language=selected_language
                    )

                    if not retrieved_chunks:
                        raise ValueError("No relevant content found for this question.")

                    answer = generate_answer(
                        question=user_question,
                        retrieved_chunks=retrieved_chunks,
                        language=selected_language,
                        recent_history=st.session_state.recent_history,
                        conversation_summary=st.session_state.conversation_summary
                    )

                    render_message(answer, is_arabic_now)

                    with st.expander("View source context"):
                        for i, chunk in enumerate(retrieved_chunks):
                            st.write(f"**Source {i+1}:**")
                            st.write(chunk)
                            st.write("---")

                except ValueError:
                    error_msg = (
                        "I couldn't find relevant information for that question. Try rephrasing it."
                        if selected_language == "english"
                        else "لم أتمكن من العثور على معلومات ذات صلة بهذا السؤال. حاول إعادة صياغته."
                    )
                    st.error(error_msg)

                except Exception as e:
                    error_msg = (
                        "Something went wrong while generating the answer. Please try again."
                        if selected_language == "english"
                        else "حدث خطأ أثناء إنشاء الإجابة. يرجى المحاولة مرة أخرى."
                    )
                    st.error(error_msg)
                    with st.expander("Technical details (for debugging)"):
                        st.code(str(e))

        if answer:
            st.session_state.display_history.append((user_question, answer, selected_language))
            st.session_state.recent_history.append((user_question, answer))

            if len(st.session_state.recent_history) > 2:
                oldest_question, oldest_answer = st.session_state.recent_history.pop(0)
                st.session_state.conversation_summary = summarize_conversation(
                    st.session_state.conversation_summary,
                    oldest_question,
                    oldest_answer
                )