import streamlit as st
import time

# FIRST: Set page config before ANY other Streamlit command
st.set_page_config(page_title="Indian Spiritual RAG")

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

# Define state for input box style
if 'input_focused' not in st.session_state:
    st.session_state.input_focused = False

# Function to update input focus state
def set_input_focus(focused):
    st.session_state.input_focused = focused

# Display title with custom styling
st.markdown("""
<style>
.main-title {
    font-size: 2.5rem;
    color: #FF5722;
    text-align: center;
    margin-bottom: 1rem;
}
.button-style {
    border: 2px solid #FF5722;
    border-radius: 8px;
}
.input-style {
    border: 2px solid #4CAF50;
    border-radius: 8px;
}
.stButton>button {
    border: 2px solid #FF5722 !important;
    border-radius: 8px !important;
}
/* Default style for input (green) */
.stTextInput>div>div>input {
    border: 2px solid #4CAF50 !important;
    border-radius: 8px !important;
}
/* Style for active input (red) - will be applied with JavaScript */
</style>

<script>
document.addEventListener('DOMContentLoaded', function() {
    // Get the input element
    const inputElement = document.querySelector('.stTextInput input');
    
    if (inputElement) {
        // Add event listeners
        inputElement.addEventListener('focus', function() {
            this.style.border = '2px solid #FF5722 !important';
        });
        
        inputElement.addEventListener('blur', function() {
            this.style.border = '2px solid #4CAF50 !important';
        });
        
        // Set initial state
        if (document.activeElement === inputElement) {
            inputElement.style.border = '2px solid #FF5722 !important';
        }
    }
});
</script>

<div class="main-title">Indian Spiritual Texts Q&A</div>
""", unsafe_allow_html=True)

# Initialize session state
if 'initialized' not in st.session_state:
    st.session_state.initialized = False
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

# Setup all authentication
if not st.session_state.initialized:
    try:
        setup_all_auth()
    except Exception as e:
        st.error(f"Authentication error: {str(e)}")

# Preload the model to avoid session state issues
if not st.session_state.initialized:
    try:
        init_message = st.empty()
        init_message.info("Hang in there! We are setting the system up for you. 😊")
        
        # Force model loading at startup to avoid session state issues
        load_model()
        
        # Keep the message for 2 seconds before replacing it
        time.sleep(2)
        init_message.success("System initialized successfully!")
        st.session_state.initialized = True
    except Exception as e:
        st.error(f"Error initializing: {str(e)}")

# Function to process when enter is pressed or button is clicked
def process_input():
    st.session_state.last_query = st.session_state.query_input
    st.session_state.query_input = ""
    st.session_state.submit_clicked = True

# Query input with callback for Enter key
query = st.text_input(
    "Ask your question:",
    key="query_input",
    on_change=process_input
)

# Add JavaScript to change input border color
st.markdown("""
<script>
// Get the input element
const inputElement = document.querySelector('.stTextInput input');

// Add event listeners for focus and blur
if (inputElement) {
    inputElement.addEventListener('focus', function() {
        this.style.borderColor = '#FF5722';
    });
    
    inputElement.addEventListener('blur', function() {
        this.style.borderColor = '#4CAF50';
    });
}
</script>
""", unsafe_allow_html=True)

# Display the current question if there is one
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 button
if st.button("Get Answer") or st.session_state.submit_clicked:
    if st.session_state.last_query:
        # Reset the submit flag
        st.session_state.submit_clicked = False
        
        # Single processing spinner for the entire operation
        with st.spinner("Processing your question..."):
            try:
                result = process_query(st.session_state.last_query, top_k=top_k, word_limit=word_limit)
                st.subheader("Answer:")
                st.write(result["answer_with_rag"])
                st.subheader("Sources:")
                for citation in result["citations"].split("\n"):
                    st.write(citation)
            except Exception as e:
                st.error(f"Error processing query: {str(e)}")
    else:
        st.warning("Please enter a question first.")

# 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.
""")