import streamlit as st
from rag_query import query_rag, vector_store  # Added vector_store import

# Custom CSS for styling
st.markdown("""
    <style>
     @import url('https://fonts.googleapis.com/css2?family=Roboto:wght@300;400;700&display=swap');
    .main {
        background-color: #f5f5f5;
        font-family: 'Roboto', sans-serif;
    }
    .stTextInput > div > div > input {border: 2px solid #C8102E; 
        border-radius: 5px;
        background-color: #FFFFFF;
        padding: 10px;
        font-size: 16px
    }
    .stSelectbox > div > div > select {
        border: 2px solid #C8102E;
        border-radius: 5px;
        background-color: #FFFFFF;
        padding: 10px;
        font-size: 16px;
    }
    .stButton > button {
        background-color: #333333; /* Dark grey */
        color: white;
        border-radius: 5px;
        border: none;
        padding: 10px 20px;
        font-size: 16px;
        transition: background-color 0.3s;
    }
    .stButton > button:hover {
        background-color: #4CAF50; /* Green on hover */
        color: white;
    }
    .new-chat-button > button {
        background-color: #FF0000; /* Red */
        color: white;
        border-radius: 5px;
        border: none;
        padding: 10px 20px;
        font-size: 16px;
        transition: background-color 0.3s;
    }
    .new-chat-button > button:hover {
        background-color: #CC0000; /* Darker red on hover */
        color: white;
    }
    .umc-header {
        font-size: 24px; /* Adjusted font size for sidebar */
        font-weight: bold;
        margin-top: 10px; /* Space between logo and text */
        margin-bottom: 20px; /* Space before Instructions */
    }
    .spacer {
        margin-top: 30px; /* Space between text and Instructions */
    }
    .chunk-dropdown {
        margin-top: 20px;
        margin-bottom: 20px;
    }
    .query-card {
        background-color: #FFFFFF;
        border-radius: 10px;
        padding: 15px;
        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
        margin-bottom: 20px;
    }
    .query-title {
        font-size: 18px;
        font-weight: bold;
        color: #333333;
        border-bottom: 2px solid #C8102E;
        padding-bottom: 5px;
        margin-bottom: 10px;
    }
    .sidebar {
        background-color: #E8ECEF;
        padding: 20px;
    }
    .divider {
        border-top: 1px solid #C8102E;
        margin: 10px 0;
    }
    .footer {
        background-color: #E8ECEF;
        padding: 10px;
        text-align: center;
        box-shadow: 0 -2px 4px rgba(0, 0, 0, 0.1);
        margin-top: 20px;
        font-size: 14px;
        color: #333333;
    }
    .footer a {
        color: #C8102E;
        text-decoration: none;
    }
    .footer a:hover {
        text-decoration: underline;
    }
    @media (max-width: 768px) {
        .header-banner {
            flex-direction: column;
            align-items: flex-start;
        }
        .header-text {
            font-size: 20px;
            margin-top: 10px;
        }
    }
    </style>
""", unsafe_allow_html=True)

# Initialize session state for password and query history
if "authenticated" not in st.session_state:
    st.session_state.authenticated = False
if "query_history" not in st.session_state:
    st.session_state.query_history = []  # Store queries and responses
if "retrieved_chunks" not in st.session_state:
    st.session_state.retrieved_chunks = []  # Store retrieved chunks
if "new_query_input" not in st.session_state:
    st.session_state.new_query_input = ""  # Store new query input

# Main app title (shown on all screens)
st.title("United Methodist Church Assistant")

# Header with description and New Chat button
col1, col2 = st.columns([3, 1])
with col1:
    st.write("Ask questions about The Book of Discipline 2020-2024!")
with col2:
    if st.button("New Chat", key="new_chat"):
        st.session_state.query_history = []
        st.session_state.retrieved_chunks = []
        st.session_state.new_query_input = ""
        st.rerun()

# Password form (shown only if not authenticated)
if not st.session_state.authenticated:
    with st.form(key="password_form"):
        password = st.text_input("Enter password:", type="password", key="password")
        submit_password = st.form_submit_button("Login")
        if submit_password:
            if password == "umc2024":  # Password as set
                st.session_state.authenticated = True
                st.rerun()  # Refresh to show main app
            else:
                st.error("Incorrect password. Please try again.")

# Main app content (shown only if authenticated)
if st.session_state.authenticated:
    # Interactive input with sample questions
    sample_questions = [
        "Can you tell me the history of United Methodist Church",
        "What does the Constitution say about inclusiveness?",
        "Tell me about John Wesley.",
        "Who were the key figures in the early history of Methodism in America, and what did they do?",
        "What does the United Methodist Church Constitution say about inclusiveness and racial justice?",
        "What is the role of the General Conference in the United Methodist Church?",
        "How did the United Methodist Church form through mergers over time?",
        "What challenges did the UMC face during the Civil War, and how did they impact the church?",
        "What is the UMC’s stance on ecumenical relations according to the Constitution?",
        "How has the UMC addressed issues of racism throughout its history?",
        "What changes has the UMC faced in the 21st century, particularly around gender and sexuality?",
        "How has the UMC expanded globally since its formation?",
        "What are the Special Sundays in the UMC, and how are they celebrated?",
        "How does the UMC observe Heritage Sunday, and what is its significance?",
        "What is the role of a Certified Lay Minister in the UMC, and how do they contribute to the church?",
        "How does the UMC organize its local churches, particularly in terms of financial responsibilities?",
        "How does the UMC support community outreach through its local church organization?",
        "What are the UMC’s guidelines for organizing a new local church or mission congregation?",
        "How has the UMC’s position on social issues like abortion evolved over time?",
        "What does the UMC say about social justice through its Special Sundays?"
    ]
    query = st.selectbox("Pick a question or type your own:", [""] + sample_questions, key="query_select") or st.text_input("Your question:", key="query_input")

    # Process query and display response
    if st.button("Submit"):
        with st.spinner("UMC AI Bot is thinking..."):
            # Retrieve chunks and response
            results = vector_store.similarity_search(query, k=8)  # Assuming vector_store is accessible
            st.session_state.retrieved_chunks = results  # Store retrieved chunks
            response = query_rag(query)
            st.session_state.query_history.append((query, response))  # Store query and response
            st.session_state.new_query_input = ""  # Clear the new query input
            st.rerun()  # Refresh to show response
            
    # Display query history
    for i, (q, r) in enumerate(st.session_state.query_history):
        st.markdown('<div class="query-card">', unsafe_allow_html=True)
        st.markdown(f'<div class="query-title">Query {i+1}: {q}</div>', unsafe_allow_html=True)
        st.markdown(r)
        # Display retrieved chunks in a dropdown
        if st.session_state.retrieved_chunks:
            chunk_options = [f"Chunk {j+1}: {chunk.page_content[:100]}..." for j, chunk in enumerate(st.session_state.retrieved_chunks)]
            selected_chunk = st.selectbox("View Retrieved Chunks:", ["Select a chunk"] + chunk_options, key=f"chunk_select_{i}")
            if selected_chunk != "Select a chunk":
                chunk_index = int(selected_chunk.split(":")[0].split(" ")[1]) - 1
                chunk = st.session_state.retrieved_chunks[chunk_index]
                st.markdown("**Retrieved Chunk Details:**")
                st.markdown(f"**Source:** {chunk.metadata['source']}")
                if chunk.metadata["part"]:
                    st.markdown(f"**Part:** {chunk.metadata['part']}")
                st.markdown(f"**Heading:** {chunk.metadata['heading']}")
                if chunk.metadata["paragraph_number"]:
                    st.markdown(f"**Paragraph {chunk.metadata['paragraph_number']}:** {chunk.metadata['paragraph_title']}")
                st.markdown(f"**Text:** {chunk.page_content}")
        st.markdown('</div>', unsafe_allow_html=True)
        st.markdown("---")

    # New prompt window with sample questions
    if st.session_state.query_history:  # Show only after at least one query
        st.markdown("### Explore More Insights")
        selected_query = st.selectbox("Pick another question or type your own:", [""] + sample_questions, key="new_query_select")
        custom_query = st.text_input("What else would you like to know?", value=st.session_state.new_query_input, key=f"new_query_input_{len(st.session_state.query_history)}")
        # Use custom query if provided, otherwise use selected query
        new_query = custom_query if custom_query else selected_query
        if st.button("Submit New Query"):
            if new_query:  # Ensure a query is provided
                with st.spinner("UMC AI Bot is thinking..."):
                    results = vector_store.similarity_search(new_query, k=8)
                    st.session_state.retrieved_chunks = results
                    response = query_rag(new_query)
                    st.session_state.query_history.append((new_query, response))
                    st.session_state.new_query_input = ""  # Clear the input for the next query
                    st.rerun()  # Refresh to show new response
            else:
                st.warning("Please select a question or enter a new one.")

# Sidebar with logo, text, and instructions (always shown)
st.sidebar.image("UMC_Logo.png", width=200)  # Adjusted logo size for sidebar
st.sidebar.markdown('<p class="umc-header">United Methodist Communications</p>', unsafe_allow_html=True)
st.sidebar.markdown('<div class="spacer"></div>', unsafe_allow_html=True)  # Add space
st.sidebar.header("Instructions")
st.sidebar.write("Enter a question about The United Methodist Church's Book Of Discipline 2020-2024. The assistant will retrieve relevant sections and explain them in simple language, citing parts, paragraphs, and titles as needed.")
st.sidebar.header("About This AI Bot")
st.sidebar.markdown("""
The United Methodist Church Discipline Assistant is an AI-powered tool designed to help you explore *The Book of Discipline of The United Methodist Church 2020-2024*. By leveraging Retrieval-Augmented Generation (RAG), the assistant retrieves relevant sections from the Book of Discipline and explains them in simple language, making it an invaluable resource for church leaders, members, and researchers.

**How It Can Help:**
- **Historical Insights:** Learn about the UMC’s origins, key figures like John Wesley, and its global expansion.
- **Constitutional Guidance:** Understand the UMC’s governance, inclusiveness, and social principles.
- **Practical Applications:** Explore guidelines for local church organization, Special Sundays, and community outreach.
- **Social and Ethical Stances:** Discover the UMC’s positions on issues like marriage, social justice, and more.

**Limitations:**
Please note that the dataset for this assistant covers *The Book of Discipline 2020-2024* up to Paragraph 269 on page 223. Content beyond this point, including later paragraphs or sections, is not included in the current dataset. For the most comprehensive information, refer directly to the full Book of Discipline.
""")
# Footer
st.markdown('<div class="footer">Powered by <a href="https://www.eksource.com/" target="_blank">ekSource Technologies, Inc</a></div>', unsafe_allow_html=True)