import streamlit as st
import time
from app import get_rag_chain

# --- CONFIG ---
st.set_page_config(page_title="The Message AI", page_icon="📖", layout="centered")

# --- CUSTOM CSS ---
st.markdown("""
    <style>
    .stChatMessage { border-radius: 10px; border: 1px solid #E0E0E0; }
    div[data-testid="stChatMessage"]:nth-child(even) { background-color: #FFF; border-left: 4px solid #8B5E3C; }
    div[data-testid="stChatMessage"]:nth-child(odd) { background-color: #F9F9F9; }
    </style>
""", unsafe_allow_html=True)

# --- SIDEBAR ---
with st.sidebar:
    st.title("About")
    st.info("This AI simulates the persona of William Branham using a local database of sermon transcripts.")
    if st.button("Clear Chat"):
        st.session_state.messages = []
        st.rerun()

# --- INIT ---
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "assistant", "content": "God bless you. What is on your heart?"}]

# --- LAUNCHER (With Loading Spinner) ---
@st.cache_resource
def load_chain():
    return get_rag_chain()

# This is where the magic happens. 
# It will show a spinner while unzipping 1.5GB, preventing the "Timeout" error.
with st.spinner("Opening the books (Unzipping database)... This may take 2 minutes..."):
    try:
        chain = load_chain()
    except Exception as e:
        st.error(f"Failed to load database: {e}")
        st.stop()

# --- CHAT UI ---
st.title("The Message AI")
st.caption("Interactive Archive • Powered by Gemini Flash")
st.divider()

for msg in st.session_state.messages:
    with st.chat_message(msg["role"], avatar="📖" if msg["role"] == "assistant" else "👤"):
        st.markdown(msg["content"])
        if "sources" in msg:
            with st.expander("Sermon References"):
                for src in msg["sources"]:
                    st.markdown(f"- *{src}*")

if prompt := st.chat_input("Ask a question..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user", avatar="👤"):
        st.markdown(prompt)

    with st.chat_message("assistant", avatar="📖"):
        with st.spinner("Searching the archives..."):
            try:
                response = chain.invoke({"query": prompt})
                result_text = response['result']
                
                # Extract Sources
                source_docs = response.get('source_documents', [])
                sources = list(set([doc.metadata.get('source', 'Unknown') for doc in source_docs]))

                # Typing Effect
                placeholder = st.empty()
                full_response = ""
                for chunk in result_text.split():
                    full_response += chunk + " "
                    time.sleep(0.04)
                    placeholder.markdown(full_response + "▌")
                placeholder.markdown(full_response)

                # Save History
                st.session_state.messages.append({
                    "role": "assistant", 
                    "content": full_response,
                    "sources": sources
                })
                
                if sources:
                    with st.expander("Sermon References"):
                        for src in sources:
                            st.markdown(f"- *{src}*")

            except Exception as e:
                st.error(f"Error: {e}")