pages/01_Admin.py

"""
Admin interface for document management and reindexing.
"""
import streamlit as st
import os
import shutil
import sys
import backend
from utils.logging_config import setup_logging
import config

# The 01_Admin.py filename pattern automatically makes this page show as "Admin" in the sidebar
# The main file "Public_Chat.py" will similarly display as "Public Chat" without needing special config

# Setup logging
logger = setup_logging()

# Set page config
st.set_page_config(
    page_title="Admin",
    page_icon="🔧",
    layout="wide"
)

# Hide just the deploy button
hide_deploy_button = """
<style>
button[kind="deploy"], 
.stDeployButton {
    display: none !important;
}
</style>
"""
st.markdown(hide_deploy_button, unsafe_allow_html=True)

# Simple authentication
def check_password():
    """Returns `True` if the user had the correct password."""
    def password_entered():
        """Checks whether a password entered by the user is correct."""
        try:
            correct_password = st.secrets.get("admin_password", "admin")
        except Exception:
            # Fallback to default password if secrets are not available
            correct_password = "admin"
            
        if st.session_state["password"] == correct_password:
            st.session_state["password_correct"] = True
            del st.session_state["password"]  # Don't store the password
        else:
            st.session_state["password_correct"] = False

    if "password_correct" not in st.session_state:
        # First run, show input for password
        st.text_input(
            "Password", type="password", on_change=password_entered, key="password"
        )
        return False
    elif not st.session_state["password_correct"]:
        # Password not correct, show input + error
        st.text_input(
            "Password", type="password", on_change=password_entered, key="password"
        )
        st.error("😕 Password incorrect")
        return False
    else:
        # Password correct
        return True

# Initialize chatbot
def initialize_chatbot():
    if "chatbot" not in st.session_state:
        with st.spinner("Initializing chatbot..."):
            # Get configuration from config module
            chatbot_config = config.get_chatbot_config()
            
            # Initialize chatbot
            logger.info("Initializing chatbot...")
            st.session_state.chatbot = backend.Chatbot(chatbot_config)
            
            # Load documents and create index
            documents = st.session_state.chatbot.load_documents()
            st.session_state.chatbot.create_index(documents)
            st.session_state.chatbot.initialize_query_engine()
            logger.info("Chatbot initialized successfully")

# Main admin interface
if check_password():
    # Initialize the chatbot
    initialize_chatbot()
    
    # Admin header
    st.title("🔧 Document Chatbot Admin")
    st.markdown("""
    This is the admin interface for managing documents and index settings.
    """)
    
    # Tabs for different admin functions
    tab1, tab2, tab3 = st.tabs(["Document Management", "Index Management", "Configuration"])
    
    # Document Management Tab
    with tab1:
        st.header("Document Management")
        
        # Show current documents
        data_dir = "data"
        col1, col2 = st.columns([3, 1])
        
        with col1:
            st.subheader("Current Documents")
            if os.path.exists(data_dir):
                documents = [f for f in os.listdir(data_dir) if not f.startswith('.')]
                if documents:
                    for doc in documents:
                        doc_col1, doc_col2 = st.columns([5, 1])
                        doc_col1.write(f"📄 {doc}")
                        if doc_col2.button("Delete", key=f"delete_{doc}"):
                            try:
                                os.remove(os.path.join(data_dir, doc))
                                st.success(f"Deleted {doc}")
                                st.rerun()
                            except Exception as e:
                                st.error(f"Error deleting {doc}: {str(e)}")
                else:
                    st.info("No documents found in data directory.")
            else:
                st.error("Data directory not found.")
                if st.button("Create Data Directory"):
                    os.makedirs(data_dir)
                    st.success("Data directory created.")
                    st.rerun()
        
        with col2:
            st.subheader("Upload Document")
            uploaded_file = st.file_uploader("Choose a file", type=["txt", "pdf", "docx", "md"])
            
            if uploaded_file is not None:
                # Save the uploaded file
                if not os.path.exists(data_dir):
                    os.makedirs(data_dir)
                
                file_path = os.path.join(data_dir, uploaded_file.name)
                with open(file_path, "wb") as f:
                    f.write(uploaded_file.getbuffer())
                
                st.success(f"File {uploaded_file.name} saved successfully!")
    
    # Index Management Tab
    with tab2:
        st.header("Index Management")
        
        # Index Status
        st.subheader("Index Status")
        index_dir = "index"
        if os.path.exists(index_dir):
            index_files = os.listdir(index_dir)
            if index_files:
                st.write(f"Index exists with {len(index_files)} files")
                total_size = sum(os.path.getsize(os.path.join(index_dir, f)) for f in index_files)
                st.write(f"Total index size: {total_size/1024:.2f} KB")
                
                # Display index files
                with st.expander("Index Files"):
                    for file in index_files:
                        st.write(f"- {file}: {os.path.getsize(os.path.join(index_dir, file))/1024:.2f} KB")
            else:
                st.warning("Index directory exists but is empty.")
        else:
            st.error("Index directory not found.")
        
        # Reindex Options
        st.subheader("Reindex Options")
        
        col1, col2 = st.columns(2)
        
        with col1:
            if st.button("Update Index", help="Add new documents to the existing index"):
                with st.spinner("Updating index..."):
                    try:
                        # Load fresh documents
                        documents = st.session_state.chatbot.load_documents()
                        # Update the index with new documents
                        st.session_state.chatbot.update_index(documents)
                        st.success("Index updated successfully!")
                        logger.info("Index updated successfully")
                    except Exception as e:
                        st.error(f"Error updating index: {str(e)}")
                        logger.error(f"Error updating index: {e}")
        
        with col2:
            if st.button("Rebuild Index", help="Delete and recreate the entire index"):
                with st.spinner("Rebuilding index..."):
                    try:
                        # Delete index directory
                        if os.path.exists(index_dir):
                            shutil.rmtree(index_dir)
                            os.makedirs(index_dir)
                            logger.info("Index directory cleared")
                        
                        # Load documents and recreate index
                        documents = st.session_state.chatbot.load_documents()
                        st.session_state.chatbot.create_index(documents)
                        st.session_state.chatbot.initialize_query_engine()
                        
                        st.success("Index rebuilt successfully!")
                        logger.info("Index rebuilt successfully")
                    except Exception as e:
                        st.error(f"Error rebuilding index: {str(e)}")
                        logger.error(f"Error rebuilding index: {e}")
    
    # Configuration Tab
    with tab3:
        st.header("Configuration")
        
        # Load configuration from config module
        chatbot_config = config.get_chatbot_config()
        
        # Configuration form
        with st.form(key="config_form"):
            st.subheader("Model Settings")
            
            col1, col2 = st.columns(2)
            
            # Model settings
            model_col1, model_col2 = st.columns(2)
            model = model_col1.text_input("Model Name", value=chatbot_config.get("model", "claude-3-7-sonnet-20250219"))
            temperature = model_col2.slider("Temperature", min_value=0.0, max_value=1.0, value=chatbot_config.get("temperature", 0.1), step=0.1)
            
            max_tokens = col1.number_input("Max Tokens", value=chatbot_config.get("max_tokens", 2048), min_value=1, max_value=4096)
            
            # Embedding Settings
            st.write("#### Embedding Settings")
            emb_col1, emb_col2 = st.columns(2)
            
            embedding_model = emb_col1.text_input("Embedding Model", value=chatbot_config.get("embedding_model", "sentence-transformers/all-MiniLM-L6-v2"))
            device = emb_col2.selectbox("Device", options=["cpu", "cuda"], index=0 if chatbot_config.get("device", "cpu") == "cpu" else 1)
            
            embed_batch_size = emb_col1.number_input("Embedding Batch Size", value=chatbot_config.get("embed_batch_size", 8), min_value=1, max_value=64)
            
            # Submit button
            submitted = st.form_submit_button("Save Configuration")
            
            if submitted:
                # Update config
                try:
                    # Create new config dictionary
                    new_config = {
                        "model": model,
                        "temperature": temperature,
                        "max_tokens": max_tokens,
                        "embedding_model": embedding_model,
                        "device": device,
                        "embed_batch_size": embed_batch_size
                    }
                    
                    # Update config in session state
                    st.session_state.chatbot.update_config(new_config)
                    
                    st.success("Configuration updated successfully!")
                    logger.info("Configuration updated successfully")
                except Exception as e:
                    st.error(f"Error updating configuration: {str(e)}")
                    logger.error(f"Error updating configuration: {e}")
        
        # System Information
        st.subheader("System Information")
        system_info = {
            "Python Version": sys.version,
            "Operating System": os.name,
            "Index Directory": os.path.abspath(index_dir) if os.path.exists(index_dir) else "Not found",
            "Data Directory": os.path.abspath(data_dir) if os.path.exists(data_dir) else "Not found"
        }
        
        for key, value in system_info.items():
            st.write(f"**{key}:** {value}")
else:
    st.title("🔧 Document Chatbot Admin")
    st.write("Please enter the password to access the admin interface.")