import streamlit as st
import time

# FIRST: Set page config before ANY other Streamlit command
st.set_page_config(page_title="Spirituality Q&A", page_icon="🕉️")

# Initialize ALL session state variables right at the beginning
if 'initialized' not in st.session_state:
    st.session_state.initialized = False
if 'model' not in st.session_state:
    st.session_state.model = None
if 'tokenizer' not in st.session_state:
    st.session_state.tokenizer = None
if 'last_query' not in st.session_state:
    st.session_state.last_query = ""
if 'submit_clicked' not in st.session_state:
    st.session_state.submit_clicked = False
if 'init_time' not in st.session_state:
    st.session_state.init_time = None
if 'form_key' not in st.session_state:
    st.session_state.form_key = 0  # This will help us reset the form
# New variable for debouncing: whether processing is in progress
if 'is_processing' not in st.session_state:
    st.session_state.is_processing = False
# Store the answer so that it persists on screen
if 'last_answer' not in st.session_state:
    st.session_state.last_answer = None
# Add new session state for showing/hiding acknowledgment
if 'show_acknowledgment' not in st.session_state:
    st.session_state.show_acknowledgment = False
# Add page change detection
if 'page_loaded' not in st.session_state:
    st.session_state.page_loaded = True
    # Reset query state when returning to home page
    st.session_state.last_query = ""
    st.session_state.last_answer = None

# THEN: Import your modules
# from rag_engine import process_query, load_model, cached_load_data_files
# from utils import setup_all_auth

# Function to toggle acknowledgment visibility
def toggle_acknowledgment():
    st.session_state.show_acknowledgment = not st.session_state.show_acknowledgment

# Custom styling (pure CSS)
st.markdown("""
<style>
.main-title {
    font-size: 2.5rem;
    color: #c0392b;
    text-align: center;
    margin-bottom: 1rem;
}
/* Button styling */
.stButton>button {
    background-color: #fff0f0 !important;
    color: #3f51b5 !important;
    border: 1px solid #e1e4f2 !important;
    border-radius: 20px !important;
    padding: 8px 16px !important;
    box-shadow: 0 1px 2px rgba(0,0,0,0.03) !important;
    white-space: nowrap !important;
    overflow: hidden !important;
    text-overflow: ellipsis !important;
}
/* Form submit button specific styling */
button[type="submit"], 
.stFormSubmit>button, 
[data-testid="stFormSubmitButton"]>button {
    background-color: #fff0f0 !important;
    color: #3f51b5 !important;
    border: 1px solid #e1e4f2 !important;
    border-radius: 8px !important;
}
.stButton>button:hover {
    background-color: #fafbff !important;
    border-color: #c5cae9 !important;
}
/* Special styling for thank you button - Modified to match Get Answer */
.thank-you-button {
    display: flex;
    justify-content: center;
    margin: 10px 0;
}
.thank-you-button > div > button {
    background-color: #fff0f0 !important;
    color: #3f51b5 !important;
    border: 1px solid #e1e4f2 !important;
    border-radius: 8px !important;
    padding: 8px 16px !important;
    box-shadow: 0 1px 2px rgba(0,0,0,0.03) !important;
    font-weight: normal !important;
    min-width: 150px !important;
}
.thank-you-button > div > button:hover {
    background-color: #fafbff !important;
    border-color: #c5cae9 !important;
}
/* Input field styling */
div[data-baseweb="input"] {
    border: 1px solid #fff0f0 !important;
    border-radius: 8px !important;
    background-color: #ffffff !important;
}
div[data-baseweb="input"]:focus-within {
    border: 1px solid #3f51b5 !important;
}
div[data-baseweb="input"]:active {
    border: 1px solid #fff0f0 !important;
}
/* Style the st.info boxes */
div.stInfo {
    background-color: #f8faff !important;
    color: #3f51b5 !important;
    border: 1px solid #e1e4f2 !important;
    border-radius: 8px !important;
}
/* COMBINED SCROLL CONTAINER */
.questions-scroll-container {
    width: 100%;
    overflow-x: auto;
    scrollbar-width: none; /* Firefox */
    -ms-overflow-style: none; /* IE and Edge */
}
/* Hide scrollbar for Chrome, Safari and Opera */
.questions-scroll-container::-webkit-scrollbar {
    display: none;
}
/* Inner content that holds both rows */
.questions-content {
    display: inline-flex;
    flex-direction: column;
    min-width: max-content;
    gap: 10px;
    padding: 5px 0;
}
/* Individual rows */
.questions-row {
    display: flex;
    flex-direction: row;
    gap: 10px;
}
/* Placeholder for buttons */
.button-placeholder {
    min-height: 38px;
    min-width: 120px;
    margin: 0 5px;
}
/* Acknowledgment section styling - fully fixed */
.acknowledgment-container {
    background-color: #f8f5ff;
    border: 1px solid #e0d6fe;
    border-radius: 8px;
    padding: 15px;
    margin: 20px 0;
    box-shadow: 0 2px 5px rgba(0,0,0,0.05);
}
.acknowledgment-header {
    color: #6a1b9a;
    font-size: 1.3rem;
    margin-bottom: 10px;
    text-align: center;
}
.more-info-link {
    text-align: center;
    margin-top: 10px;
    font-style: italic;
}
.citation-note {
    font-size: 0.8rem;
    font-style: italic;
    color: #666;
    padding: 10px;
    border-top: 1px solid #eee;
    margin-top: 30px;
}
</style>
<div class="main-title">Spirituality Q&A</div>
""", unsafe_allow_html=True)

# Preload resources during initialization
init_message = st.empty()
if not st.session_state.initialized:
    init_message.info("Hang in there! We are setting the system up for you. 😊")
    try:
        # Setup authentication and preload heavy resources
        # setup_all_auth()
        # load_model()  # This uses cached_load_model via alias
        # cached_load_data_files()  # Preload FAISS index, text chunks, and metadata
        st.session_state.initialized = True
        st.session_state.init_time = time.time()
        init_message.success("System initialized successfully!")
        time.sleep(2)
        init_message.empty()
    except Exception as e:
        init_message.error(f"Error initializing: {str(e)}")
elif st.session_state.init_time is not None:
    elapsed_time = time.time() - st.session_state.init_time
    if elapsed_time >= 2.0:
        init_message.empty()
        st.session_state.init_time = None

# Thank You button that matches Get Answer button styling
st.markdown('<div class="thank-you-button">', unsafe_allow_html=True)
col1, col2, col3 = st.columns([1, 2, 1])
with col2:
    if st.session_state.show_acknowledgment:
        button_text = "Hide Thank You Note"
    else:
        button_text = "Show Thank You Note"
    if st.button(button_text, key="thank_you_button", use_container_width=True, disabled=st.session_state.is_processing):
        toggle_acknowledgment()
st.markdown('</div>', unsafe_allow_html=True)

if st.session_state.show_acknowledgment:
    st.markdown('<div class="acknowledgment-container">', unsafe_allow_html=True)
    st.markdown('<div class="acknowledgment-header">A Heartfelt Thank You</div>', unsafe_allow_html=True)
    st.markdown("""
    With deepest reverence, we express our gratitude to:
    
    **The Saints and Spiritual Masters** whose timeless wisdom illuminates this application. From ancient sages to modern masters, 
    their selfless dedication to uplift humanity through spiritual knowledge continues to guide seekers on the path.
    
    **The Sacred Texts** that have preserved the eternal truths across millennia, offering light in times of darkness 
    and clarity in times of confusion.
    
    **The Publishers** who have diligently preserved and disseminated these precious teachings, making them accessible 
    to spiritual aspirants worldwide. Their work ensures these wisdom traditions endure for future generations.
    
    **The Authors** who have dedicated their lives to interpreting and explaining complex spiritual concepts, 
    making them accessible to modern readers.
    
    This application is merely a humble vessel for the ocean of wisdom they have shared with the world. We claim no 
    ownership of these teachings—only profound gratitude for the opportunity to help make them more accessible.
    """)
    
    st.markdown('<div class="more-info-link">For detailed information about our sources, please visit the "Sources" page in the navigation menu.</div>', unsafe_allow_html=True)
    st.markdown('</div>', unsafe_allow_html=True)

# Function to handle query selection from the common questions buttons
def set_query(query):
    # If already processing, ignore further input
    if st.session_state.is_processing:
        return
    st.session_state.last_query = query
    st.session_state.submit_clicked = True
    st.session_state.is_processing = True
    st.experimental_rerun()

# Function to group questions into rows based on length
def group_buttons(questions, max_chars_per_row=100):
    rows = []
    current_row = []
    current_length = 0
    for q in questions:
        # Add some buffer for button padding/margins
        q_length = len(q) + 5
        if current_length + q_length > max_chars_per_row and current_row:
            rows.append(current_row)
            current_row = [q]
            current_length = q_length
        else:
            current_row.append(q)
            current_length += q_length
    if current_row:
        rows.append(current_row)
    return rows

# All common questions in a single list
common_questions = [
    "What is the Atman or the soul?",
    "Are there rebirths?",
    "What is Karma?",
    "What is the ultimate truth?",
    "What is Swami Vivekananda's opinion about the SELF?",
    "Explain moksha or salvation. Is that for real?",
    "Destiny or free will?",
    "What is the ultimate goal of human life?",
    "Do we really die?",
    "How can you attain self-realization?"
]

# Display heading for common questions
st.markdown("### Common questions to try:")

# Group questions into rows and create buttons (disabled if processing)
question_rows = group_buttons(common_questions, max_chars_per_row=80)
for row_idx, row in enumerate(question_rows):
    cols = st.columns(len(row))
    for i, (col, q) in enumerate(zip(cols, row)):
        with col:
            if st.button(q, key=f"r{row_idx}_q{i}", use_container_width=True, disabled=st.session_state.is_processing):
                set_query(q)

# Function to handle form submission
def handle_form_submit():
    # If already processing, ignore further input
    if st.session_state.is_processing:
        return
    if st.session_state.query_input and st.session_state.query_input.strip():
        st.session_state.last_query = st.session_state.query_input.strip()
        st.session_state.submit_clicked = True
        st.session_state.is_processing = True
        # Increment the form key to force a reset
        st.session_state.form_key += 1

# Create a form with a dynamic key (to allow resetting)
with st.form(key=f"query_form_{st.session_state.form_key}"):
    query = st.text_input("Ask your question:", key="query_input", 
                          placeholder="Press enter to submit your question", disabled=st.session_state.is_processing)
    submit_button = st.form_submit_button("Get Answer", on_click=handle_form_submit, disabled=st.session_state.is_processing)

# Display the current question if available
if st.session_state.last_query:
    st.markdown("### Current Question:")
    st.info(st.session_state.last_query)

# Sliders for customization
col1, col2 = st.columns(2)
with col1:
    top_k = st.slider("Number of sources:", 3, 10, 5)
with col2:
    word_limit = st.slider("Word limit:", 50, 500, 200)

# Process the query only if it has been explicitly submitted
if st.session_state.submit_clicked and st.session_state.last_query:
    st.session_state.submit_clicked = False
    with st.spinner("Processing your question..."):
        try:
            # result = process_query(st.session_state.last_query, top_k=top_k, word_limit=word_limit)
            # Simulated result for testing
            result = {"answer_with_rag": "This is a sample answer.", "citations": "Sample citation 1\nSample citation 2"}
            st.session_state.last_answer = result  # Store result in session state
        except Exception as e:
            st.session_state.last_answer = {"answer_with_rag": f"Error processing query: {str(e)}", "citations": ""}
    # Reset debouncing after processing and force a rerun to re-enable buttons
    st.session_state.is_processing = False
    st.experimental_rerun()

# Display the answer if available
if st.session_state.last_answer is not None:
    st.subheader("Answer:")
    st.write(st.session_state.last_answer["answer_with_rag"])
    st.subheader("Sources:")
    for citation in st.session_state.last_answer["citations"].split("\n"):
        st.write(citation)

# Add helpful information
st.markdown("---")
st.markdown("""
### About this app
This application uses a Retrieval-Augmented Generation (RAG) system to answer questions about Indian spiritual texts.
It searches through a database of texts to find relevant passages and generates answers based on those passages.

For detailed information about the sources used in this application, including saints, sacred texts, and publishers, 
please visit the "Sources" page in the navigation menu.
""")

# Citation note at the bottom - improved with support message
st.markdown('<div class="citation-note">', unsafe_allow_html=True)
st.markdown("""
The answers presented in this application are re-presented summaries of relevant passages from the listed citations. 
For the original works in their complete and authentic form, users are respectfully encouraged to purchase 
the original print or digital works from their respective publishers. Your purchase helps support these publishers 
who have brought such important spiritual works to the world.
""")
st.markdown('</div>', unsafe_allow_html=True)