app_streamlit.py.bak

import os
import streamlit as st
import asyncio
import tempfile
from pathlib import Path
import sys
import requests
import logging
import traceback
import time
from threading import Thread
from queue import Queue
import sqlite3
import hashlib
import json
import shutil
from datetime import datetime
import chromadb

# Configure logging
logging.basicConfig(level=logging.INFO, 
                   format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Add the RAG-Agent source code to the Python path
rag_agent_path = os.path.join(os.path.dirname(__file__), "rag-agent/src")
sys.path.append(rag_agent_path)

# Single import indicator to prevent repeated import logging
if "imports_done" not in st.session_state:
try:
    # Import Ollama with error handling
    import ollama
    logger.info("Successfully imported ollama package")
except ImportError as e:
    logger.error(f"Failed to import ollama package: {str(e)}")
    st.error(f"Failed to import ollama package: {str(e)}")
    sys.exit(1)

try:
    # Import local modules from the rag-agent source
    from ollama_pure_chat import OllamaChat
    from fileloader import Loader_Local
    logger.info("Successfully imported modules from rag-agent")
except ImportError as e:
    logger.error(f"Failed to import modules from rag-agent: {str(e)}")
    st.error(f"Failed to import modules from rag-agent: {str(e)}")
    sys.exit(1)

    # Mark imports as completed
    st.session_state.imports_done = True
else:
    # Silent imports when already done
    import ollama
    from ollama_pure_chat import OllamaChat
    from fileloader import Loader_Local

# Check if Ollama API is available
def is_ollama_available():
    """Check if the Ollama server is running and accessible"""
    # Only check once per session unless explicitly refreshed
    if "ollama_available" not in st.session_state or "ollama_checked_time" not in st.session_state:
    try:
        response = requests.get("http://localhost:11434/api/tags", timeout=2)
            status = response.status_code == 200
            st.session_state.ollama_available = status
            st.session_state.ollama_checked_time = time.time()
            return status
        except (requests.ConnectionError, requests.Timeout):
            logger.warning("Ollama server is not running or not accessible")
            st.session_state.ollama_available = False
            st.session_state.ollama_checked_time = time.time()
        return False
        except Exception as e:
            logger.error(f"Unexpected error checking Ollama availability: {str(e)}")
            st.session_state.ollama_available = False
            st.session_state.ollama_checked_time = time.time()
            return False
    else:
        # Check if we need to refresh (every 60 seconds)
        if time.time() - st.session_state.ollama_checked_time > 60:
            try:
                response = requests.get("http://localhost:11434/api/tags", timeout=2)
                status = response.status_code == 200
                st.session_state.ollama_available = status
                st.session_state.ollama_checked_time = time.time()
                return status
            except Exception:
                # Just keep existing status on error
                return st.session_state.ollama_available
        return st.session_state.ollama_available

# Default model to use if none is specified
default_model = "deepseek-r1:1.5b"

# Persistent storage paths - Use absolute path for consistency
PERSISTENT_STORAGE_BASE = os.getenv("PERSISTENT_STORAGE_BASE", os.path.join(os.path.dirname(os.path.abspath(__file__)), 'localdatabase'))
USERS_DB_PATH = os.path.join(PERSISTENT_STORAGE_BASE, "users.db")

# Ensure persistent storage directories exist
os.makedirs(PERSISTENT_STORAGE_BASE, exist_ok=True)

# Function to get user-specific persistent storage path
def get_user_storage_path(username):
    user_path = os.path.join(PERSISTENT_STORAGE_BASE, username)
    os.makedirs(user_path, exist_ok=True)
    return user_path

# Modify init_user_database to use persistent path
def init_user_database():
    """Initialize SQLite database for user management"""
    conn = sqlite3.connect(USERS_DB_PATH)
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            username TEXT PRIMARY KEY,
            password TEXT NOT NULL,
            databases TEXT DEFAULT '["default"]',
            last_login DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()
    conn.close()

# Modify get_available_databases to use persistent storage
def get_available_databases():
    """Get available databases for the current user from ChromaDB"""
    if not hasattr(st.session_state, 'username'):
        return []
    
    try:
        # Get user-specific storage path
        user_storage_path = get_user_storage_path(st.session_state.username)
        
        # List all directories in the user's storage path
        # Each directory represents a database
        if os.path.exists(user_storage_path):
            databases = [d for d in os.listdir(user_storage_path) 
                        if os.path.isdir(os.path.join(user_storage_path, d))]
            
            # If no databases found, return empty list
            if not databases:
                return []
                
            return databases
        else:
            return []
        except Exception as e:
        logger.error(f"Error retrieving databases from ChromaDB: {str(e)}")
        return []

# Modify get_context to improve retrieval
def get_context(question):
    """Improved context retrieval with better error handling and persistent storage support"""
    if not init_loader():
        st.error("Document loader could not be initialized.")
        return "", []
    
    context = ""
    context_chunks = []
    
    try:
        # Get the current database
        current_db = st.session_state.current_database
        
        # Get user-specific storage path
        user_storage_path = get_user_storage_path(st.session_state.username)
        
        # Set absolute persist directory path for the loader
        absolute_persist_dir = os.path.join(user_storage_path, current_db)
        st.session_state.loader.persist_directory = absolute_persist_dir
        
        # Debug output
        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
            logger.info(f"Getting context for question: {question}")
            logger.info(f"Using database: {current_db}")
            logger.info(f"Absolute persist directory: {absolute_persist_dir}")
            logger.info(f"Loader persist directory set to: {st.session_state.loader.persist_directory}")
        
        # Retrieve collections for the current database
        database_collections = st.session_state.database_collections.get(current_db, {})
        
        # Show debug info if enabled
        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
            logger.info(f"Found {len(database_collections)} collections in current database")
        
        # Determine all collections to search in
        collections_to_search = []
        
        # First, add regular database collections
        for collection_name, collection_info in database_collections.items():
            collections_to_search.append({
                "id": collection_name,  # Use plain collection name without prefix
                "name": collection_name,
                "is_temporary": False
            })
        
        # Then, add temporary collections if any exist
        if hasattr(st.session_state, 'temp_collections') and st.session_state.temp_collections:
            for temp_coll_name, file_info in st.session_state.temp_collections.items():
                collections_to_search.append({
                    "id": temp_coll_name,
                    "name": file_info.get('file_name', temp_coll_name),
                    "is_temporary": True
                })
        
        # If there are no collections, return empty context
        if not collections_to_search:
            st.info("No collections available to search.")
            return "", []
            
        # Show debug info about collections if debug mode is on
        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
            logger.info(f"Searching in {len(collections_to_search)} collections")
            for coll in collections_to_search:
                logger.info(f"Collection: {coll['id']} ({coll['name']}), Temporary: {coll['is_temporary']}")
                
        # Check if the ChromaDB client is properly initialized with the correct directory
        try:
            if hasattr(st.session_state.loader, 'chroma_client'):
                client_path = getattr(st.session_state.loader.chroma_client, '_path', 'unknown')
                if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                    logger.info(f"ChromaDB client path: {client_path}")
                    
                # Reinitialize client if path doesn't match
                if client_path != absolute_persist_dir:
                    if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                        logger.info(f"Reinitializing ChromaDB client with path: {absolute_persist_dir}")
                    import chromadb
                    st.session_state.loader.chroma_client = chromadb.PersistentClient(path=absolute_persist_dir)
        except Exception as client_error:
            logger.error(f"Error checking ChromaDB client: {str(client_error)}")
        
        # Search in all collections
        for collection in collections_to_search:
            try:
                collection_id = collection["id"]
                display_name = collection["name"]
                is_temp = collection["is_temporary"]
                
                # Try to retrieve context
                if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                    logger.info(f"Querying collection: {collection_id}")
                
                # Get collection results, handling errors gracefully
                try:
                    results = st.session_state.loader(
                        question, 
                        collection_id, 
                        top_k=5  # Retrieve up to 5 most relevant chunks
                    )
                except Exception as query_error:
                    if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                        logger.error(f"Error querying collection {collection_id}: {str(query_error)}")
                        logger.error(traceback.format_exc())
                    continue  # Skip this collection and try the next one
                
                # Check if results are valid and contain documents
                if results and 'documents' in results and results['documents'] and results['documents'][0]:
                    # Add documents to the context for the model
                    for i, doc in enumerate(results['documents'][0]):
                        # Truncate very long documents to prevent context overflow
                        truncated_doc = doc[:1000] + '...' if len(doc) > 1000 else doc
                        
                        source_label = "Temporary File" if is_temp else "Database"
                        context += f"doc_{i} from {display_name} ({source_label}): {truncated_doc}\n"
                        context_chunks.append(f"**Document Chunk {i+1} from {display_name}:**\n{truncated_doc}")
                
            except Exception as collection_error:
                logger.error(f"Error retrieving from collection {collection['name']}: {str(collection_error)}")
                if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                    logger.error(f"Traceback: {traceback.format_exc()}")
        
        # Log context retrieval details
        logger.info(f"Retrieved {len(context_chunks)} context chunks for query: {question}")
        
        # If no context found, return empty
        if not context_chunks:
            st.info("No relevant context found for the query.")
    
    except Exception as e:
        logger.error(f"Comprehensive error getting context: {str(e)}")
        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
            logger.error(f"Traceback: {traceback.format_exc()}")
        st.error(f"Error retrieving context: {str(e)}")
    
    return context, context_chunks

# Modify show_create_database_modal to use persistent storage
def show_create_database_modal():
    """Create a new database for the current user"""
    if not hasattr(st.session_state, 'username'):
        st.error("Please log in first")
        return
    
    with st.form(key="create_database_form"):
        st.subheader("Create New Database")
        new_db_name = st.text_input("Database Name")
        submit_button = st.form_submit_button(label="Create")
        
        if submit_button and new_db_name:
            try:
                # Check if database already exists in session state
                current_databases = st.session_state.available_databases if hasattr(st.session_state, 'available_databases') else []
                
                # Check if database already exists
                if new_db_name in current_databases:
                    st.error(f"Database '{new_db_name}' already exists")
            else:
                    # Add new database to available databases list
                    current_databases.append(new_db_name)
                    st.session_state.available_databases = current_databases
                    
                    # Create a directory for the new ChromaDB database in user's persistent storage
                    user_storage_path = get_user_storage_path(st.session_state.username)
                    database_path = os.path.join(user_storage_path, new_db_name)
                    os.makedirs(database_path, exist_ok=True)
                    
                    # Update session state
                    st.session_state.current_database = new_db_name
                    st.success(f"ChromaDB database '{new_db_name}' created for file loading!")
                    
                    # Initialize collections for this ChromaDB database
                    if "database_collections" not in st.session_state:
                        st.session_state.database_collections = {}
                    st.session_state.database_collections[new_db_name] = {}
                    
                    # Refresh the page to update database list
                    st.rerun()
            
    except Exception as e:
                st.error(f"Error creating ChromaDB database: {str(e)}")

# Modify login_user to update last login and load persistent collections
def login_user(username, password):
    """Authenticate user and load persistent collections"""
    if len(username) < 3:
        return False, "Username must be at least 3 characters long"
    
    # Hash the password
    hashed_password = hashlib.sha256(password.encode()).hexdigest()
    
    try:
        conn = sqlite3.connect(USERS_DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute("SELECT * FROM users WHERE username = ? AND password = ?", 
                       (username, hashed_password))
        user = cursor.fetchone()
        
        if user:
            # Update last login timestamp
            cursor.execute(
                "UPDATE users SET last_login = CURRENT_TIMESTAMP WHERE username = ?", 
                (username,)
            )
            conn.commit()
            
            # Load persistent collections for the user
            user_storage_path = get_user_storage_path(username)
            
            # Initialize database collections from persistent storage
            database_collections = {}
            
            # Get user's databases
            cursor.execute("SELECT databases FROM users WHERE username = ?", (username,))
            result = cursor.fetchone()
            
            if result and result[0]:
                databases = json.loads(result[0])
                
                # Load collections for each database
                for db_name in databases:
                    db_path = os.path.join(user_storage_path, db_name)
                    if os.path.exists(db_path):
                        # Here you might want to add logic to load collection metadata
                        database_collections[db_name] = {}
            
            # Load temporary collections
            temp_collections = {}
            file_collections = {}
            
            # Load any persisted temporary collections from ChromaDB
            if init_loader():
                loader = Loader_Local()
                
                # Check each database directory for temporary collections
                for db_name in databases:
                    db_path = os.path.join(user_storage_path, db_name)
                    if os.path.exists(db_path):
                        # Set up the loader to check this database
                        loader.persist_directory = db_path
                        
                        try:
                            # Get all collection names from ChromaDB
                            chroma_client = chromadb.PersistentClient(path=db_path)
                            all_collections = chroma_client.list_collections()
                            
                            # Find temporary collections for this user
                            for collection in all_collections:
                                collection_name = collection.name
                                
                                # Only load temp collections for this user
                                if collection_name.startswith(f"temp_{username}_"):
                                    # Extract file ID from collection name
                                    file_id = collection_name.replace(f"temp_{username}_", "")
                                    
                                    # Try to get collection metadata
                                    try:
                                        collection_obj = chroma_client.get_collection(name=collection_name)
                                        # Get any file in this collection to extract metadata
                                        results = collection_obj.get(limit=1)
                                        
                                        if results and results['metadatas'] and results['metadatas'][0]:
                                            source = results['metadatas'][0].get('source', '')
                                            file_name = os.path.basename(source)
                                            file_extension = Path(file_name).suffix.lstrip(".").lower()
                                            
                                            # Store in temp_collections
                                            temp_collections[collection_name] = {
                                                "file_name": file_name,
                                                "file_type": file_extension,
                                                "added_date": "Previously uploaded",
                                                "is_temporary": True
                                            }
                                            
                                            # Add to file_collections for retrieval
                                            file_collections[file_id] = collection_name
            except Exception as e:
                                        logger.error(f"Error loading temporary collection metadata: {str(e)}")
    except Exception as e:
                            logger.error(f"Error listing collections in database {db_name}: {str(e)}")
            
            conn.close()
            
            return True, "Login successful", database_collections, temp_collections, file_collections
        else:
            conn.close()
            return False, "Invalid username or password", {}, {}, {}
    except Exception as e:
        return False, f"Login error: {str(e)}", {}, {}, {}

# Initialize the loader only once
def init_loader():
    if not hasattr(st.session_state, 'loader') or st.session_state.loader is None:
        try:
            # Check if Ollama API is available
            try:
                models = ollama.list()
                logger.info(f"Ollama API available with {len(models.get('models', []))} models")
            except Exception as e:
                logger.warning(f"Ollama API not available: {str(e)}")
            
            # Initialize loader with a default persist directory (will be updated later per user/database)
            st.session_state.loader = Loader_Local()
            
            # Set user-specific storage path if user is logged in
            if hasattr(st.session_state, 'username') and st.session_state.username:
                user_storage_path = get_user_storage_path(st.session_state.username)
                
                # If current database is set, use its path as the persist directory
                if hasattr(st.session_state, 'current_database') and st.session_state.current_database:
                    absolute_persist_dir = os.path.join(user_storage_path, st.session_state.current_database)
                    # Ensure directory exists
                    os.makedirs(absolute_persist_dir, exist_ok=True)
                    
                    # Update loader with proper persist directory
                    st.session_state.loader.persist_directory = absolute_persist_dir
                    
                    # Re-initialize ChromaDB client with the correct path
                    import chromadb
                    st.session_state.loader.chroma_client = chromadb.PersistentClient(path=absolute_persist_dir)
                    
                    if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                        logger.info(f"Initialized loader with persist directory: {absolute_persist_dir}")
            
            logger.info("Successfully initialized Loader_Local")
            return True
        except Exception as e:
            logger.error(f"Failed to initialize Loader_Local: {str(e)}")
            if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                logger.error(traceback.format_exc())
            return False
    return True

# Modify main to handle persistent collections
def main():
    # Set page config
    st.set_page_config(
        page_title="RAG-Agent Chat",
        page_icon="🤖",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Initialize session state variables
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    if "initialized" not in st.session_state:
        st.session_state.initialized = False
    
    if "selected_model" not in st.session_state:
        st.session_state.selected_model = default_model
    
    if "enable_retrieval" not in st.session_state:
        st.session_state.enable_retrieval = True
        
    if "file_collections" not in st.session_state:
        st.session_state.file_collections = {}
    
    # Check Ollama status only once per session load
    # Use session variable instead of calling API multiple times
    if "ollama_available" not in st.session_state:
        # Only make the API call once
    ollama_status = is_ollama_available()
    else:
        ollama_status = st.session_state.ollama_available
        
    # Initialize the document loader
    loader_success = init_loader()
    if not loader_success:
        logger.warning("Failed to initialize document loader. Some features may not be available.")
    
    # Initialize the database
    init_user_database()
    
    # Check if the user is authenticated
    if "authenticated" not in st.session_state or not st.session_state.authenticated:
        show_login_page()
        return
    
    # Ensure chat_instance exists and is initialized at the very beginning
    if "chat_instance" not in st.session_state or st.session_state.chat_instance is None:
        if ollama_status:
            try:
                # Create a new chat instance
                model_name = st.session_state.selected_model
                logger.info(f"Initializing chat instance with model: {model_name}")
                st.session_state.chat_instance = OllamaChat(
                    model=model_name,
                    system_prompt="You are a helpful chatbot assistant designed to answer questions about the given context. <context>"
                )
                st.session_state.initialized = True
                logger.info(f"Successfully initialized chat instance with model: {model_name}")
            except Exception as e:
                logger.error(f"Error initializing chat instance: {str(e)}\n{traceback.format_exc()}")
                st.session_state.chat_instance = None
                st.session_state.initialized = False
                # We'll show an error to the user in the UI later
        else:
            logger.warning("Ollama server is not available. Chat functionality will be limited.")
            st.session_state.chat_instance = None
            st.session_state.initialized = False
    
    # Set current database to first available database or default
    available_dbs = get_available_databases()

    if available_dbs:
        st.session_state.current_database = available_dbs[0]
    
    # Load persistent collections
    if "database_collections" not in st.session_state:
        user_storage_path = get_user_storage_path(st.session_state.username)
        database_collections = {}
        
        for db_name in available_dbs:
            db_path = os.path.join(user_storage_path, db_name)
            if os.path.exists(db_path):
                # Load collections from files.db
                files_db_path = os.path.join(db_path, "files.db")
                if os.path.exists(files_db_path):
                    try:
                        files_conn = sqlite3.connect(files_db_path)
                        files_cursor = files_conn.cursor()
                        files_cursor.execute("SELECT * FROM files")
                        files = files_cursor.fetchall()
                        
                        database_collections[db_name] = {}
                        for file in files:
                            file_id = file[0]
                            file_name = file[1]
                            file_type = file[2]
                            collection_name = file[3]
                            added_date = file[4]
                            description = file[5]
                            
                            database_collections[db_name][file_id] = {
                                "collection_name": collection_name,
                                "file_name": file_name,
                                "file_type": file_type,
                                "added_date": added_date,
                                "description": description or ""
                            }
                        
                        files_conn.close()
                    except Exception as e:
                        logger.error(f"Error loading files from database {db_name}: {str(e)}")
                        database_collections[db_name] = {}
                else:
                    database_collections[db_name] = {}
        
        st.session_state.database_collections = database_collections
    
    # Sidebar Configuration
    with st.sidebar:
        st.title("RAG-Agent Chat")
        
        # User info
        st.markdown(f"**Logged in as:** {st.session_state.username}")
        if st.button("Logout"):
            st.session_state.authenticated = False
            st.rerun()
            
        st.divider()
        
        # Add debug mode toggle
        if "debug_mode" not in st.session_state:
            st.session_state.debug_mode = False
            
        debug_mode = st.toggle("Debug Mode", value=st.session_state.debug_mode, key="debug_toggle")
        if debug_mode != st.session_state.debug_mode:
            st.session_state.debug_mode = debug_mode
            if debug_mode:
                st.info("Debug mode enabled. Additional logging information will be shown.")
                # Set logging level to DEBUG when debug mode is on
                logger.setLevel(logging.DEBUG)
            else:
                # Reset to INFO level when debug mode is off
                logger.setLevel(logging.INFO)
        
        # Check if Ollama is running - show clear status
        if not ollama_status:
            st.error("⚠️ Ollama server is not running")
            st.info("Please start Ollama by running 'ollama serve' in a terminal")
            
            # Add manual check button
            if st.button("Check Ollama Status"):
                # Force refresh the Ollama status
                if "ollama_available" in st.session_state:
                    del st.session_state.ollama_available
                if "ollama_checked_time" in st.session_state:
                    del st.session_state.ollama_checked_time
                if "cached_models" in st.session_state:
                    del st.session_state.cached_models
                st.rerun()
        else:
            st.success("✅ LLM  is running")
        
        # Database Selection Section
        st.markdown("### Database Selection")
        
        # Get available databases
        available_databases = get_available_databases()
        
        # Create a column layout for the database controls
        db_col1, db_col2 = st.columns([3, 1])
        
        with db_col1:
            # Dropdown to select database - only show if there are databases
            if available_databases:
            selected_db = st.selectbox(
                "Select Database:",
                available_databases,
                index=available_databases.index(st.session_state.current_database) if st.session_state.current_database in available_databases else 0
            )
            
            if selected_db != st.session_state.current_database:
                st.session_state.current_database = selected_db
                # Update collections for the new database
                st.session_state.file_collections = get_collections_for_database(selected_db)
                st.success(f"Switched to database: {selected_db}")
            else:
                st.info("No databases available. Create a new one.")
        
        with db_col2:
            # Button to create new database
            if st.button("New DB"):
                show_create_database_modal()
        
        # Database Editor/Viewer
        with st.expander("Database Details", expanded=False):
            show_database_editor()
            
        st.divider()
        
        # Retrieval Toggle
        st.markdown("### Document Retrieval")
        retrieval_status = "🔍 ON" if st.session_state.enable_retrieval else "🚫 OFF"
        toggle_label = f"Database Retrieval: {retrieval_status}"
        
        if st.toggle(toggle_label, value=st.session_state.enable_retrieval, key="retrieval_toggle"):
            st.session_state.enable_retrieval = True
        else:
            st.session_state.enable_retrieval = False
        
        st.divider()
        
        # Model Selection Dropdown
        st.markdown("### Model Selection")
        
        # Only show model dropdown if Ollama is running
        if ollama_status:
            # Get available models
            models = get_available_models()
            
            # Create dropdown for model selection
            if models:
            selected_model = st.selectbox(
                "Select a model:",
                models,
                index=models.index(st.session_state.selected_model) if st.session_state.selected_model in models else 0
            )
            
            # Apply model change
            if selected_model != st.session_state.selected_model:
                try:
                    # Update the session state
                    st.session_state.selected_model = selected_model
                    
                    # Create new chat instance with selected model
                    st.session_state.chat_instance = OllamaChat(
                        model=selected_model,
                            system_prompt="You are a helpful chatbot assistant designed to answer questions about the given context. <context>"
                    )
                    st.session_state.initialized = True
                    st.success(f"Model changed to: {selected_model}")
                    logger.info(f"Model changed to: {selected_model}")
                except Exception as e:
                    logger.error(f"Error changing model: {str(e)}\n{traceback.format_exc()}")
                    st.error(f"Error changing model: {str(e)}")
            else:
                st.warning("No models available. Please pull models using Ollama CLI.")
        else:
            st.warning("Model selection unavailable - Ollama is not running")
        
        st.divider()
        
        # ChromaDB Collections Section
        with st.expander("ChromaDB Collections", expanded=True):
            st.markdown("### Add Documents to ChromaDB")
            
            # File Upload for ChromaDB
            uploaded_files = st.file_uploader(
                "Upload documents to add to ChromaDB:",
                accept_multiple_files=True,
                type=["pdf", "txt", "csv", "md", "doc", "docx"],
                key="chromadb_uploader"
            )
            
            # Process files for ChromaDB
            if uploaded_files:
                # Get current database and existing collections
                current_db = st.session_state.current_database
                existing_collections = list(st.session_state.database_collections.get(current_db, {}).keys())
                
                st.markdown("### Configure Collection for Uploaded Files")
                
                # Collection selection
                collection_choice = st.radio(
                    "Choose Collection Option", 
                    ["Create New Collection", "Use Existing Collection"],
                    key="upload_collection_choice"
                )
                
                if collection_choice == "Use Existing Collection":
                    # Dropdown for existing collections
                    if existing_collections:
                        selected_collection = st.selectbox(
                            "Select Existing Collection:", 
                            existing_collections,
                            key="upload_existing_collection"
                        )
                        
                        # No description needed for existing collections
                        collection_description = None
                    else:
                        st.warning("No existing collections. Please create a new one.")
                        collection_choice = "Create New Collection"
                
                if collection_choice == "Create New Collection":
                    # Input for custom collection name
                    st.markdown("#### Create New Collection")
                    st.info("Enter a descriptive name for your collection. This name should reflect the content or topic of your documents.")
                    
                    new_collection_name = st.text_input(
                        "Collection Name", 
                        key="upload_new_collection_name",
                        help="Choose a descriptive name related to the document content. For example: 'machine_learning', 'research_papers', 'company_policies', etc."
                    )
                    
                    # Validate collection name
                    if new_collection_name and new_collection_name in existing_collections:
                        st.warning(f"Collection '{new_collection_name}' already exists. Files will be added to this existing collection.")
                    
                    # Optional: Add a description or tags
                    collection_description = st.text_area(
                        "Collection Description (Optional)", 
                        key="upload_collection_description",
                        help="Provide additional context about what this collection contains or represents."
                    )
                    
                    selected_collection = new_collection_name
                
                # Process the files
                if st.button("Process Files", key="upload_process_files_btn", disabled=not (selected_collection and selected_collection.strip())):
                    with st.spinner("Processing files..."):
                        success_count = 0
                        for file in uploaded_files:
                            success = process_file(file, selected_collection, 
                                                  collection_description if collection_choice == "Create New Collection" else None)
                            if success:
                                success_count += 1
                        
                        if success_count == len(uploaded_files):
                            st.success(f"Successfully processed all {success_count} files to collection '{selected_collection}'")
                        elif success_count > 0:
                            st.warning(f"Processed {success_count} out of {len(uploaded_files)} files to collection '{selected_collection}'")
                        else:
                            st.error("Failed to process any files. Please check the logs for details.")
            
            # Show help text 
            st.markdown("#### How to use file upload")
            st.markdown("""
            1. Upload one or more documents using the file uploader
            2. Choose to create a new collection or use an existing one
            3. If creating a new collection, enter a descriptive name for your documents
            4. Click "Process Files" to add the documents to the collection
            5. Use the chat interface to ask questions about your documents
            """)
            
            # Show available collections in current database
            collections = get_collections_for_database(st.session_state.current_database)
            if collections:
                st.markdown("### Available Collections")
                collection_names = list(collections.keys())
                
                # Allow selecting a collection
                selected_collection = st.selectbox(
                    "Select Collection to Use:",
                    collection_names,
                    key="selected_chromadb_collection"
                )
                
                if selected_collection:
                    st.session_state.active_collection = selected_collection
                    st.success(f"Using collection: {selected_collection}")
                    
                    # Show collection details
                    if selected_collection in collections:
                        coll_info = collections[selected_collection]
                        st.markdown(f"**Collection:** {selected_collection}")
                        
                        # Show all files in the collection if available
                        if "files" in coll_info and len(coll_info["files"]) > 0:
                            st.markdown(f"**Files:** {len(coll_info['files'])} document(s)")
                            
                            # Display first 5 files directly
                            files_to_show = min(5, len(coll_info["files"]))
                            for idx, file in enumerate(coll_info["files"][:files_to_show]):
                                st.markdown(f"- {file}")
                                
                            # If there are more files, provide option to view more
                            remaining_files = len(coll_info["files"]) - files_to_show
                            if remaining_files > 0:
                                display_option = st.radio(
                                    f"Show all files ({remaining_files} more)",
                                    ["Hide Additional Files", "Show All Files"],
                                    key=f"show_all_{selected_collection}",
                                    horizontal=True
                                )
                                if display_option == "Show All Files":
                                    for idx, file in enumerate(coll_info["files"][files_to_show:]):
                                        st.markdown(f"- {file}")
                        else:
                            st.markdown(f"**File:** {coll_info.get('file_name', 'Unknown')}")
                            
                        st.markdown(f"**Type:** {coll_info.get('file_type', 'Unknown')}")
                        st.markdown(f"**Added:** {coll_info.get('added_date', 'Unknown')}")
                        
                        # Inspect button
                        if st.button("Inspect Collection Data", key="inspect_from_dropdown"):
                            inspect_collection(selected_collection)
            else:
                st.info("No collections available in the current database. Upload documents to create collections.")
        
        st.divider()
        
        # File Upload Section for temporary QA
        st.markdown("### Upload Documents")
        
        # Track upload state to avoid reprocessing
        if "upload_state" not in st.session_state:
            st.session_state.upload_state = {}
            
        if "temp_collections" not in st.session_state:
            st.session_state.temp_collections = {}
        
        # Create an expander to show currently loaded temporary files
        with st.expander("Temporary Files in Memory", expanded=True):
            if st.session_state.temp_collections:
                for temp_coll_name, file_info in st.session_state.temp_collections.items():
                    st.markdown(f"📄 **{file_info['file_name']}** - *Collection: {temp_coll_name}*")
                
                # Add button to clear all temporary collections
                if st.button("Clear All Temporary Files"):
                    st.session_state.temp_collections = {}
                    st.success("All temporary files cleared from memory")
                    st.rerun()
            else:
                st.info("No temporary files loaded. Upload files below for quick Q&A without saving to your database.")
            
        uploaded_files = st.file_uploader(
            "Upload documents for temporary QA (not saved to database):",
            accept_multiple_files=True,
            type=["pdf", "txt", "csv", "md", "doc", "docx"]
        )
        
        # Process newly uploaded files
        if uploaded_files:
            with st.spinner("Processing files for quick Q&A..."):
                # Add a progress bar
                progress_bar = st.progress(0)
                status_text = st.empty()
                
                # Track processing status
                all_processed = True
                
                total_files = len(uploaded_files)
                for idx, file in enumerate(uploaded_files):
                    # Update progress
                    progress = float(idx) / float(total_files)
                    progress_bar.progress(progress)
                    status_text.text(f"Processing file {idx+1}/{total_files}: {file.name}")
                    
                    # Create a unique identifier for this file
                file_id = file.name.replace(" ", "_").replace(".", "_")
                
                # Only process if it's a new file we haven't seen before
                    if file_id not in st.session_state.upload_state:
                        # Create a temporary collection name specific to this user and file
                        temp_collection_name = f"temp_{st.session_state.username}_{file_id}"
                        
                    # Mark as being processed
                    st.session_state.upload_state[file_id] = "processing"
                        
                        # Create a temporary file
                        temp_dir = tempfile.mkdtemp()
                        temp_file_path = os.path.join(temp_dir, file.name)
                        
                        try:
                            # Write file to temporary location
                            with open(temp_file_path, "wb") as f:
                                f.write(file.getbuffer())
                            
                            # Get file extension
                            file_extension = Path(file.name).suffix.lstrip(".").lower()
                            
                            # Get user-specific storage path
                            user_storage_path = get_user_storage_path(st.session_state.username)
                            
                            # Set persistent directory for the loader
                            current_db = st.session_state.current_database
                            current_db_path = os.path.join(user_storage_path, current_db)
                            st.session_state.loader.persist_directory = current_db_path
                            
                            # Process the file with timeout handling
                            logger.info(f"Processing temporary file: {file.name} into collection: {temp_collection_name}")
                            status_text.text(f"Processing file {idx+1}/{total_files}: {file.name} - Generating chunks and embeddings...")
                            
                            # Import the needed modules for timeout handling
                            import concurrent.futures
                            import time
                            
                            # Define a function to process with timeout
                            def process_with_timeout(timeout=120):  # 2-minute timeout
                                try:
                                    # Store local references to avoid session state access in thread
                                    local_loader = st.session_state.loader
                                    
                                    # Define a completely isolated thread function - no Streamlit dependencies
                                    def process_file_thread():
                                        try:
                                            # No session state access in this function
                                            result = local_loader.create_or_insert_collection(
                                                temp_file_path,
                                                temp_collection_name,
                                                file_extension,
                                                'local_parser'
                                            )
                                            return result
                                        except Exception as e:
                                            # Don't use logger inside thread
                                            print(f"Error in processing thread: {str(e)}")
                                            return False

                                    # Try to create the temporary collection in a separate thread
                                    with concurrent.futures.ThreadPoolExecutor() as executor:
                                        future = executor.submit(process_file_thread)
                                        
                                        # Poll and update status while waiting
                                        start_time = time.time()
                                        processing = True
                                        while processing and time.time() - start_time < timeout:
                                            if future.done():
                                                processing = False
                                                return future.result()
                                            else:
                                                # Only update the status message from the main thread
                                                elapsed = time.time() - start_time
                                                # This will work because status_text is in the main thread context
                                                status_text.text(f"Processing file {idx+1}/{total_files}: {file.name} - {elapsed:.1f}s elapsed...")
                                                # Add a small sleep to prevent UI from freezing
                                                time.sleep(0.5)
                                        
                                        # If we get here, we've timed out
                                        if not future.done():
                                            future.cancel()
                                            raise TimeoutError("Processing timed out")

                                except Exception as e:
                                    logger.error(f"Error in processing with timeout: {str(e)}")
                                    return False
                            
                            # Process with timeout
                            success = process_with_timeout()
                            
                            if success:
                                # Store in file_collections for retrieval
                                if "file_collections" not in st.session_state:
                                    st.session_state.file_collections = {}
                                
                                st.session_state.file_collections[file_id] = temp_collection_name
                                
                                # Store metadata for display
                                st.session_state.temp_collections[temp_collection_name] = {
                                    "file_name": file.name,
                                    "file_type": file_extension,
                                    "added_date": time.strftime("%Y-%m-%d %H:%M:%S"),
                                    "is_temporary": True
                                }
                                
                                # Mark as processed
                                st.session_state.upload_state[file_id] = "processed"
                                status_text.text(f"✅ Processed file {idx+1}/{total_files}: {file.name}")
                            else:
                                all_processed = False
                                status_text.text(f"❌ Failed to process file {idx+1}/{total_files}: {file.name}")
                                st.session_state.upload_state[file_id] = "failed"
                        
                        except TimeoutError:
                            all_processed = False
                            logger.error(f"Processing timed out for file: {file.name}")
                            st.error(f"Processing timed out for file: {file.name} - File may be too large or complex")
                            st.session_state.upload_state[file_id] = "failed"
                        except Exception as e:
                            all_processed = False
                            logger.error(f"Error processing temporary file: {str(e)}")
                            st.error(f"Error processing file: {str(e)}")
                            st.session_state.upload_state[file_id] = "failed"
                        finally:
                            # Clean up temporary files
                            try:
                                if os.path.exists(temp_file_path):
                                    os.remove(temp_file_path)
                                if os.path.exists(temp_dir):
                                    os.rmdir(temp_dir)
                            except Exception as cleanup_error:
                                logger.warning(f"Failed to clean up temporary files: {str(cleanup_error)}")
                        
                        # Complete the progress bar
                        progress_bar.progress(1.0)
                        status_text.text("Processing completed")
                        
                        if all_processed:
                            st.success(f"Successfully processed all {total_files} files")
                        else:
                            st.warning("Some files failed to process. See error messages above.")
                            
                        # Force UI to refresh after a short delay
                        time.sleep(1)  # Give user time to see completion message
                        st.rerun()
        
        # Help text for temporary uploads
        st.markdown("""
        **Temporary uploads:**
        - Files uploaded here are processed immediately for Q&A
        - Context is retrieved automatically when you ask questions
        - These files remain available between sessions
        - They don't appear in your regular database collections
        - Use the database upload section above for permanent storage
        """)
                
        st.divider()
        
        # Display collections dropdown
        show_collections_dropdown()
        
        st.divider()
        
        # Clear Chat Button
        if st.button("🗑️ Clear Chat History"):
            st.session_state.chat_history = []
            st.session_state.messages = []
            st.session_state.file_collections = {}
            st.session_state.upload_state = {}
            
            # Don't reinitialize these - just keep them
            # st.session_state.loader - keep this
            # st.session_state.chat_instance - keep this
            
            st.success("Chat history and documents have been cleared.")
        
        # Show Models Information
        with st.expander("🤖 Available Models", expanded=True):
            if ollama_status:
                models = get_available_models()
                if models:
                    st.markdown("### Ollama Models")
                    for model in models:
                        if model == st.session_state.selected_model:
                            st.markdown(f"✅ **{model}** (selected)")
                        else:
                            st.markdown(f"• {model}")
                    
                    # Add a button to refresh models list
                    if st.button("Refresh Models"):
                        if "cached_models" in st.session_state:
                            del st.session_state.cached_models
                        st.rerun()
                else:
                    st.warning("No models found. Is Ollama running?")
            else:
                st.warning("Cannot retrieve models - Ollama is not running")
        
        # Help Information
        with st.expander("ℹ️ Help & Information"):
            st.markdown("""
            ### RAG-Agent Chat Help
            
            **Features:**
            - Chat with selected Ollama model
            - Upload documents to ask questions about them
            - Toggle database retrieval on/off
            - Choose from different models
            - Select or create databases
            
            **Tips:**
            - Upload multiple files to build a larger knowledge base
            - Be specific with your questions for better answers
            - Toggle database retrieval off if you want pure model responses
            - Clear chat history to start fresh
            - Use different databases for different projects
            """)
            
    # Main Chat Area
    chat_container = st.container()
    
    with chat_container:
        # Display header
    st.title("RAG-Agent Chat")
    
    # Show database info in the header
        if hasattr(st.session_state, 'current_database'):
    st.markdown(f"**Current Database:** {st.session_state.current_database}")
    
        # Display Ollama status
        if not st.session_state.ollama_available:
            st.error("⚠️ Ollama server is not available. Please start it by running 'ollama serve' in a terminal and refresh this page.")
        elif "chat_instance" not in st.session_state or st.session_state.chat_instance is None:
            st.warning("⚠️ Chat instance not initialized. Please check that Ollama is running correctly and the selected model is available.")
            
            # Add a retry button
            if st.button("Retry Connection"):
                try:
                    # Attempt to reconnect to Ollama
                    if is_ollama_available():
            model_name = st.session_state.selected_model
            st.session_state.chat_instance = OllamaChat(
                model=model_name,
                            system_prompt="You are a helpful chatbot assistant designed to answer questions about the given context. <context>"
            )
            st.session_state.initialized = True
                        st.success(f"Successfully connected to Ollama with model: {model_name}")
                        st.rerun()
                    else:
                        st.error("Ollama is still not available. Please start it and try again.")
        except Exception as e:
                    st.error(f"Failed to initialize chat: {str(e)}")
    
    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
                st.write(message["content"])
    
    # Display retrieved context
    context_container = st.container()
    
    # Chat input - enable only if chat is initialized and Ollama is running
        chat_disabled = not st.session_state.ollama_available or not st.session_state.initialized
        
        if chat_disabled:
            st.warning("Chat input disabled. Please ensure Ollama is running and a model is loaded.")
        prompt = None
    else:
        prompt = st.chat_input("Ask a question about your documents...")
    
    if prompt:
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display user message
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Display assistant response
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            context = ""
            
            # Get context if retrieval is enabled
                if hasattr(st.session_state, 'enable_retrieval') and st.session_state.enable_retrieval and "file_collections" in st.session_state and st.session_state.file_collections:
                with context_container:
                    context_status = st.empty()
                    context_status.info("🔍 Retrieving relevant context...")
                    
                    try:
                        context, context_chunks = get_context(prompt)
                        
                        if context_chunks:
                            st.success(f"🔍 Found {len(context_chunks)} relevant document chunks.")
                            with st.expander("Show Retrieved Context"):
                                for i, chunk in enumerate(context_chunks):
                                    st.markdown(chunk)
                        else:
                            context_status.info("🔍 No relevant document chunks found for this query.")
                    except Exception as e:
                        context_status.error(f"Error retrieving context: {str(e)}")
                        logger.error(f"Error retrieving context: {str(e)}\n{traceback.format_exc()}")
            
            # Generate response
            full_response = ""
                
                # Double check initialization before proceeding
                if "chat_instance" not in st.session_state or st.session_state.chat_instance is None:
                    message_placeholder.error("Chat instance is not initialized. Please try refreshing the page.")
                    return
            
            try:
                message_placeholder.markdown("Thinking...")
                
                # Get chat history in the format expected by OllamaChat
                    chat_history = []
                    if hasattr(st.session_state, 'chat_history'):
                chat_history = st.session_state.chat_history
                    else:
                        st.session_state.chat_history = chat_history
                        
                # Generate response in a streaming fashion
                result_queue = generate_llm_response(prompt, chat_history, context)
                
                # Stream the response
                timeout_count = 0
                while True:
                    try:
                        result_type, content = result_queue.get(timeout=0.1)
                        
                        if result_type == "chunk":
                            full_response += content
                            message_placeholder.markdown(full_response + "▌")
                            timeout_count = 0  # Reset timeout counter on successful retrieval
                        elif result_type == "done":
                            message_placeholder.markdown(full_response)
                            
                            # Update chat history
                            chat_history.append(f"User: {prompt}\n{full_response}")
                            
                            # Keep only the last 10 interactions
                            if len(chat_history) > 10:
                                chat_history = chat_history[-10:]
                            
                            st.session_state.chat_history = chat_history
                            break
                        elif result_type == "error":
                            message_placeholder.error(f"Error: {content}")
                            break
                    except Exception as queue_error:
                        # No result yet, continue waiting but with a timeout
                        timeout_count += 1
                        # If we've waited too long (5 seconds) with no response, show status
                        if timeout_count > 50 and timeout_count % 50 == 0:  # Every ~5 seconds
                            message_placeholder.markdown(f"{full_response}... _still thinking_▌")
                        continue
                
            except Exception as e:
                logger.error(f"Error generating response: {str(e)}\n{traceback.format_exc()}")
                message_placeholder.error(f"Error: {str(e)}")
            
            # Add assistant response to chat history
            if full_response:
                st.session_state.messages.append({"role": "assistant", "content": full_response})

# Function to show login page
def show_login_page():
    # Initialize user database
    init_user_database()
    
    st.title("RAG-Agent Chat Login")
    
    # Center the login form
    col1, col2, col3 = st.columns([1, 2, 1])
    
    with col2:
        login_mode = st.radio("Choose Action", ["Login", "Register"])
        
        if login_mode == "Login":
            st.markdown("### Login")
        username = st.text_input("Username")
        password = st.text_input("Password", type="password")
        
        if st.button("Login"):
                success, message, database_collections, temp_collections, file_collections = login_user(username, password)
                if success:
                st.session_state.authenticated = True
                st.session_state.username = username
                    
                    # Set persistent collections
                    st.session_state.database_collections = database_collections
                    st.session_state.temp_collections = temp_collections
                    st.session_state.file_collections = file_collections
                    
                    # Set current database
                    available_dbs = list(database_collections.keys()) if database_collections else []
                    if available_dbs:
                        st.session_state.current_database = available_dbs[0]
                    
                st.rerun()
            else:
                    st.error(message)
        
        else:
            st.markdown("### Register")
            new_username = st.text_input("Choose Username")
            new_password = st.text_input("Choose Password", type="password")
            confirm_password = st.text_input("Confirm Password", type="password")
            
            if st.button("Register"):
                if new_password != confirm_password:
                    st.error("Passwords do not match")
                else:
                    success, message = register_user(new_username, new_password)
                    if success:
                        st.success(message)
                        # Optionally, automatically log in after registration
                        st.session_state.authenticated = True
                        st.session_state.username = new_username
                        
                        # Initialize default database for new user
                        st.session_state.database_collections = {"default": {}}
                        st.session_state.current_database = "default"
                        
                        # Initialize empty temp collections
                        st.session_state.temp_collections = {}
                        st.session_state.file_collections = {}
                        
                        st.rerun()
            else:
                        st.error(message)

# Function to get available Ollama models
def get_available_models():
    """Retrieve list of available Ollama models"""
    # Only fetch models once per session, rather than using a TTL-based cache
    if "cached_models" not in st.session_state:
        try:
            # Check if Ollama is available first
            if not is_ollama_available():
                logger.warning("Ollama server is not running. Cannot retrieve models.")
                st.session_state.cached_models = [default_model]  # Store in session state
                return st.session_state.cached_models
            
            # Use ollama library to list models
            response = ollama.list()
            
            # Extract model names
            if response and 'models' in response:
                models = [model['model'] for model in response['models']]
                
                # Ensure default model is in the list
                if default_model not in models:
                    models.insert(0, default_model)
                
                st.session_state.cached_models = models  # Store in session state
                logger.info(f"Retrieved {len(models)} models from Ollama")
                return models
            else:
                logger.warning("No models found. Using default model.")
                st.session_state.cached_models = [default_model]  # Store in session state
                return st.session_state.cached_models
        
        except Exception as e:
            logger.error(f"Error retrieving Ollama models: {str(e)}")
            st.session_state.cached_models = [default_model]  # Store in session state
            return st.session_state.cached_models
    else:
        # Return cached models
        return st.session_state.cached_models

# Function to show database editor
def show_database_editor():
    st.markdown(f"**Current Database:** {st.session_state.current_database}")
    
    # Get collections for current database
    collections = get_collections_for_database(st.session_state.current_database)
    
    if collections:
        st.markdown("**Collections:**")
        
        # Add search filter for collections
        filter_query = st.text_input("Filter collections", key="collection_filter")
        filtered_collections = {k: v for k, v in collections.items() 
                               if not filter_query or filter_query.lower() in k.lower()}
        
        # Display collections in a table format if there are many
        if len(filtered_collections) > 5:
            collection_data = []
            for coll_name, coll_info in filtered_collections.items():
                file_name = coll_info.get("file_name", "Unknown")
                file_type = coll_info.get("file_type", "Unknown")
                added_date = coll_info.get("added_date", "Unknown")
                collection_data.append({
                    "Collection": coll_name,  # Show collection name first
                    "File": file_name,
                    "Type": file_type,
                    "Added": added_date
                })
            
            # Create a DataFrame for display
            import pandas as pd
            df = pd.DataFrame(collection_data)
            st.dataframe(df)
            
            # Handle collection actions
            st.markdown("**Collection Actions:**")
            col1, col2 = st.columns(2)
            with col1:
                collection_to_edit = st.selectbox("Select collection", list(filtered_collections.keys()))
            with col2:
                action = st.selectbox("Action", ["View Details", "Delete"])
            
            if st.button("Execute Action"):
                if action == "Delete":
                    if delete_collection(st.session_state.current_database, collection_to_edit):
                        st.success(f"Deleted collection: {collection_to_edit}")
                        st.rerun()
                elif action == "View Details":
                    show_collection_details(st.session_state.current_database, collection_to_edit)
        else:
            # For fewer collections, show more detailed individual cards
            for coll_name, coll_info in filtered_collections.items():
                with st.container():
                    # Headline using collection name instead of file name
                    st.markdown(f"📄 **Collection: {coll_name}**")
                    col1, col2 = st.columns([3, 1])
                    
                    with col1:
                        # Show all files in the collection if available
                        if "files" in coll_info and len(coll_info["files"]) > 0:
                            st.markdown(f"**Files:** {len(coll_info['files'])} document(s)")
                            
                            # Use a radio button instead of checkbox to avoid nesting issues
                            display_option = st.radio(
                                f"Files in '{coll_name}'",
                                ["Hide Files", "Show Files"],
                                key=f"display_files_{coll_name}",
                                horizontal=True
                            )
                            
                            if display_option == "Show Files":
                                # Display files directly in the main container
                                max_files_displayed = 10
                                for idx, file in enumerate(coll_info["files"][:max_files_displayed]):
                                    st.markdown(f"{idx+1}. {file}")
                                
                                # If there are more files than can comfortably be shown
                                if len(coll_info["files"]) > max_files_displayed:
                                    st.markdown(f"... and {len(coll_info['files']) - max_files_displayed} more files")
                        else:
                            # Fallback to single file
                        st.markdown(f"**File:** {coll_info.get('file_name', 'Unknown')}")
                            
                        # Show type information
                        if "entries" in coll_info:
                            # This is a text input collection
                            st.markdown(f"**Type:** Text Collection")
                            st.markdown(f"**Entries:** {coll_info.get('entries', 1)}")
                        else:
                            # This is a file-based collection
                        st.markdown(f"**Type:** {coll_info.get('file_type', 'Unknown')}")
                            
                        st.markdown(f"**Added:** {coll_info.get('added_date', 'Unknown')}")
                        if "modified_date" in coll_info:
                            st.markdown(f"**Last Modified:** {coll_info.get('modified_date', 'Unknown')}")
                        
                        # Add notes or description field that can be edited
                        if "description" not in coll_info:
                            coll_info["description"] = ""
                        
                        new_description = st.text_area(f"Notes for '{coll_name}'", 
                                                      value=coll_info.get("description", ""), 
                                                      key=f"desc_{coll_name}")
                        
                        if new_description != coll_info.get("description", ""):
                            st.session_state.database_collections[st.session_state.current_database][coll_name]["description"] = new_description
                            st.success("Notes saved!")
                    
                    with col2:
                        if st.button("Delete", key=f"del_{coll_name}"):
                            if delete_collection(st.session_state.current_database, coll_name):
                                st.success(f"Deleted collection: {coll_name}")
                                st.rerun()
                        
                        # Button to view document chunks for debugging
                        if st.button("Inspect", key=f"inspect_{coll_name}"):
                            inspect_collection(coll_name)
    else:
        st.info("No collections in this database yet. Upload documents to create collections.")
    
    # Rename and Import/Export sections
    st.markdown("### Database Management")
    
    # Rename Database
    rename_container = st.container()
    with rename_container:
        new_name = st.text_input("Rename Current Database", key="rename_db_input")
        if st.button("Rename Database"):
            if new_name and new_name != st.session_state.current_database:
                if new_name not in get_available_databases():
                    rename_database(st.session_state.current_database, new_name)
                    st.success(f"Database renamed to '{new_name}'")
                    st.rerun()
                else:
                    st.error(f"Database '{new_name}' already exists")
    
    # Import/Export
    import_export_container = st.container()
    with import_export_container:
        col1, col2 = st.columns(2)
        with col1:
            st.download_button(
                "Export Database",
                data=get_database_export_data(st.session_state.current_database),
                file_name=f"{st.session_state.current_database}_export.json",
                mime="application/json"
            )
        with col2:
            st.file_uploader("Import Database", type=["json"], key="db_import_file", 
                            on_change=import_database_callback)

# Function to inspect a collection
def inspect_collection(collection_name):
    if not init_loader():
        st.error("Loader is not initialized")
        return
    
    collection_id = f"{collection_name}"
    
    try:
        # Get collection details - handle this more gracefully
        try:
            # Try to get the collection directly from ChromaDB
            user_storage_path = get_user_storage_path(st.session_state.username)
            current_db = st.session_state.current_database
            current_db_path = os.path.join(user_storage_path, current_db)
            
            # Create a local client to avoid thread issues
            chroma_client = chromadb.PersistentClient(path=current_db_path)
            
            # Check if collection exists
            try:
                collection = chroma_client.get_collection(name=collection_id)
                
                # Get collection data safely
                collection_data = {}
                try:
                    # Get count
                    count = collection.count()
                    collection_data["count"] = count
                    
                    # Get a sample of documents if available
                    if count > 0:
                        sample = collection.get(limit=min(5, count))
                        
                        # Only include important fields to avoid overwhelming the display
                        collection_data["sample"] = {
                            "ids": sample.get("ids", [])[:5],
                            "metadatas_count": len(sample.get("metadatas", [])),
                            "documents_preview": [doc[:100] + "..." if len(doc) > 100 else doc 
                                               for doc in sample.get("documents", [])[:3]]
                        }
                except Exception as e:
                    collection_data["error"] = f"Error retrieving data: {str(e)}"
                
                st.json(collection_data)
                
            except Exception as coll_error:
                st.error(f"Collection '{collection_id}' not found: {str(coll_error)}")
                
        except Exception as e:
            st.error(f"Error accessing ChromaDB: {str(e)}")
            
    except Exception as e:
        st.error(f"Error inspecting collection: {str(e)}")
        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
            st.error(traceback.format_exc())

# Function to show collection details
def show_collection_details(database_name, collection_name):
    if database_name in st.session_state.database_collections:
        if collection_name in st.session_state.database_collections[database_name]:
            coll_info = st.session_state.database_collections[database_name][collection_name]
            
            st.markdown(f"### Collection: {collection_name}")
            st.markdown(f"**File:** {coll_info.get('file_name', 'Unknown')}")
            st.markdown(f"**Type:** {coll_info.get('file_type', 'Unknown')}")
            st.markdown(f"**Added:** {coll_info.get('added_date', 'Unknown')}")
            st.markdown(f"**Description:** {coll_info.get('description', '')}")
            
            # Add option to inspect collection chunks
            if st.button("Inspect Collection Data"):
                inspect_collection(collection_name)
        else:
            st.error(f"Collection {collection_name} not found")
    else:
        st.error(f"Database {database_name} not found")

# Helper function to export database data
def get_database_export_data(database_name):
    if database_name in st.session_state.database_collections:
        export_data = {
            "database_name": database_name,
            "collections": st.session_state.database_collections[database_name],
            "export_date": time.strftime("%Y-%m-%d %H:%M:%S")
        }
        return json.dumps(export_data, indent=2)
    else:
        return json.dumps({"error": "Database not found"})

# Callback for database import
def import_database_callback():
    uploaded_file = st.session_state.db_import_file
    if uploaded_file is not None:
        try:
            content = uploaded_file.getvalue().decode('utf-8')
            import_data = json.loads(content)
            
            # Check if it's a valid export
            if "database_name" in import_data and "collections" in import_data:
                db_name = import_data["database_name"]
                
                # Add database if it doesn't exist
                if db_name not in st.session_state.available_databases:
                    st.session_state.available_databases.append(db_name)
                
                # Add collections
                if "database_collections" not in st.session_state:
                    st.session_state.database_collections = {}
                
                st.session_state.database_collections[db_name] = import_data["collections"]
                
                # Switch to the imported database
                st.session_state.current_database = db_name
                st.success(f"Successfully imported database: {db_name}")
                
                # Clear the file uploader
                st.session_state.db_import_file = None
            else:
                st.error("Invalid database export format")
        except Exception as e:
            st.error(f"Error importing database: {str(e)}")
            
# Function to get collections for a database
def get_collections_for_database(database_name):
    # In a real app, this would query actual collections
    if "database_collections" not in st.session_state:
        st.session_state.database_collections = {}
    
    if database_name not in st.session_state.database_collections:
        st.session_state.database_collections[database_name] = {}
    
    return st.session_state.database_collections[database_name]

# Function to delete a collection
def delete_collection(database_name, collection_name):
    # Delete from session state
    if database_name in st.session_state.database_collections:
        if collection_name in st.session_state.database_collections[database_name]:
            del st.session_state.database_collections[database_name][collection_name]
            
    # Also delete from file_collections if present
    collection_value = f"collection_{collection_name}"
    keys_to_delete = []
    for key, value in st.session_state.file_collections.items():
        if value == collection_value:
            keys_to_delete.append(key)
            
    for key in keys_to_delete:
        del st.session_state.file_collections[key]
    
    # In a real app, this would also delete from the actual database
    return True

# Function to rename a database
def rename_database(old_name, new_name):
    """Rename a database for the current user"""
    if not hasattr(st.session_state, 'username'):
        st.error("Please log in first")
        return False
    
    try:
        conn = sqlite3.connect(USERS_DB_PATH)
        cursor = conn.cursor()
        
        # Retrieve current databases
        cursor.execute("SELECT databases FROM users WHERE username = ?", 
                       (st.session_state.username,))
        result = cursor.fetchone()
        
        import json
        current_databases = json.loads(result[0]) if result and result[0] else []
        
        # Check if old database exists and new name is unique
        if old_name not in current_databases:
            st.error(f"Database '{old_name}' does not exist")
            return False
        
        if new_name in current_databases:
            st.error(f"Database '{new_name}' already exists")
            return False
        
        # Rename the database
        current_databases = [new_name if db == old_name else db for db in current_databases]
        
        # Update user's databases in SQLite
        cursor.execute(
            "UPDATE users SET databases = ? WHERE username = ?", 
            (json.dumps(current_databases), st.session_state.username)
        )
        conn.commit()
        
        # Update session state
        st.session_state.current_database = new_name
        
        # Rename collections if they exist in session state
    if "database_collections" in st.session_state and old_name in st.session_state.database_collections:
        st.session_state.database_collections[new_name] = st.session_state.database_collections[old_name]
        del st.session_state.database_collections[old_name]
    
        return True
    
    except Exception as e:
        st.error(f"Error renaming database: {str(e)}")
        return False
    finally:
        conn.close()

# User registration function
def register_user(username, password):
    """Register a new user in the SQLite database"""
    if len(username) < 3:
        return False, "Username must be at least 3 characters long"
    
    # Hash the password
    hashed_password = hashlib.sha256(password.encode()).hexdigest()
    
    try:
        conn = sqlite3.connect(USERS_DB_PATH)
        cursor = conn.cursor()
        
        try:
            cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", 
                           (username, hashed_password))
            conn.commit()
            return True, "Registration successful"
        except sqlite3.IntegrityError:
            return False, "Username already exists"
        finally:
            conn.close()
    except Exception as e:
        return False, f"Registration error: {str(e)}"

# Process file uploads
def process_file(file, collection_name=None, collection_description=None):
    """
    Process an uploaded file and insert it into a ChromaDB collection.
    
    Args:
        file (UploadedFile): The file uploaded by the user
        collection_name (str, optional): Name of the collection to insert the file into
        collection_description (str, optional): Description of the collection
    
    Returns:
        bool: True if successful, False otherwise
    """
    if not init_loader():
        st.error("Document loader could not be initialized. Cannot process files.")
        return False
    
    logger.info(f"Processing uploaded file: {file.name}")
    
    # Create a temporary file
    temp_dir = tempfile.mkdtemp()
    temp_file_path = os.path.join(temp_dir, file.name)
    
    try:
        # Write file to temporary location
        with open(temp_file_path, "wb") as f:
            f.write(file.getbuffer())
        
        # Get file extension
        file_extension = Path(file.name).suffix.lstrip(".").lower()
        
        # Get user-specific storage path
        user_storage_path = get_user_storage_path(st.session_state.username)
        
        # If collection name is not provided, we need to prompt the user
        if collection_name is None:
            # Get existing collections for the current database
            current_db = st.session_state.current_database
            existing_collections = list(st.session_state.database_collections.get(current_db, {}).keys())
            
            # Create a container for collection selection
            collection_container = st.container()
            
            with collection_container:
                # Radio button to choose between existing and new collection
                collection_choice = st.radio(
                    "Choose Collection", 
                    ["Use Existing Collection", "Create New Collection"],
                    key="collection_selection_radio"
                )
                
                if collection_choice == "Use Existing Collection":
                    # Dropdown for existing collections
                    if existing_collections:
                        selected_collection = st.selectbox(
                            "Select Existing Collection:", 
                            existing_collections,
                            key="existing_collection_select"
                        )
                    else:
                        st.warning("No existing collections. Please create a new one.")
                        collection_choice = "Create New Collection"
                
                if collection_choice == "Create New Collection":
                    # Input for new collection name with more guidance
                    st.info(f"Uploading file: {file.name}")
                    st.info(f"File type: {file_extension}")
                    
                    # Emphasize entering a meaningful collection name
                    st.markdown("### Enter a Meaningful Collection Name")
                    st.markdown("This name will be used to identify your collection. Choose a name that describes the content or purpose of the document.")
                    
                    # Provide an example rather than a suggestion based on file type
                    example_collection_name = "my_topic_collection"
                    
                    # Add an input for custom collection name with more context
                    new_collection_name = st.text_input(
                        "Collection Name", 
                        value=example_collection_name,
                        key="new_collection_name_input",
                        help="Choose a descriptive name related to the document content. For example: 'machine_learning', 'research_papers', 'company_policies', etc."
                    )
                    
                    # Optional: Add a description or tags
                    collection_description = st.text_area(
                        "Collection Description (Optional)", 
                        key="collection_description_input",
                        help="Provide additional context about what this collection contains or represents."
                    )
                    
                    # Validate collection name
                    if not new_collection_name:
                        st.error("Collection name cannot be empty")
                        return False
                    
                    # Check if collection name already exists
                    if new_collection_name in existing_collections:
                        st.error(f"Collection '{new_collection_name}' already exists")
                        return False
                    
                    selected_collection = new_collection_name
                
                # Confirm button
                if st.button("Confirm Collection", key="confirm_collection_btn"):
                    return process_file_with_collection(file, selected_collection, collection_description if collection_choice == "Create New Collection" else None)
                
                # If not confirmed, stop execution
                st.stop()
        else:
            # If collection name is provided, process directly
            return process_file_with_collection(file, collection_name, collection_description)
            
    except Exception as e:
        logger.error(f"Unexpected error processing file: {str(e)}\n{traceback.format_exc()}")
        st.error(f"Unexpected error processing file '{file.name}': {str(e)}")
        return False
    finally:
        # Clean up
        try:
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)
            if os.path.exists(temp_dir):
                os.rmdir(temp_dir)
        except Exception as cleanup_error:
            logger.warning(f"Failed to clean up temporary files: {str(cleanup_error)}")
            
def process_file_with_collection(file, collection_name, collection_description=None):
    """
    Internal function to process a file with a specified collection
    """
    if not collection_name:
        st.error("Collection name is required")
        return False
        
    logger.info(f"Processing file {file.name} for collection: {collection_name}")
    
    # Create a temporary file
    temp_dir = tempfile.mkdtemp()
    temp_file_path = os.path.join(temp_dir, file.name)
    
    try:
        # Write file to temporary location
        with open(temp_file_path, "wb") as f:
            f.write(file.getbuffer())
            
        # Get file extension
        file_extension = Path(file.name).suffix.lstrip(".").lower()
        
        # Get user-specific storage path and current database
        user_storage_path = get_user_storage_path(st.session_state.username)
        current_db = st.session_state.current_database
        
        # Ensure the current database directory exists
        current_db_path = os.path.join(user_storage_path, current_db)
        os.makedirs(current_db_path, exist_ok=True)
        
        # Set persistent directory for the loader
        st.session_state.loader.persist_directory = current_db_path
        
        # Debug logging
        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
            logger.info(f"Setting loader persist directory to: {current_db_path}")
            logger.info(f"Current ChromaDB client path: {getattr(st.session_state.loader.chroma_client, '_path', 'unknown')}")
        
        # Re-initialize ChromaDB client with the correct path to ensure it's properly set
        try:
            import chromadb
            st.session_state.loader.chroma_client = chromadb.PersistentClient(path=current_db_path)
            
            if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                logger.info(f"Re-initialized ChromaDB client with path: {current_db_path}")
                logger.info(f"Updated ChromaDB client path: {getattr(st.session_state.loader.chroma_client, '_path', 'unknown')}")
        except Exception as client_error:
            logger.error(f"Error re-initializing ChromaDB client: {str(client_error)}")
            if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                logger.error(traceback.format_exc())
        
        # Set collection ID (should match the exact format expected by ChromaDB)
        collection_id = collection_name
        
        # Try to create or insert the collection
        try:
            result = st.session_state.loader.create_or_insert_collection(
                path_local=temp_file_path,
                collection_name=collection_id,
                extension=file_extension
            )
            
            if not result:
                logger.error(f"Failed to create or insert collection: {collection_id}")
                st.error(f"Failed to process file. Please check the logs for details.")
                return False
                
            # After successful processing, add collection to database_collections
            if "database_collections" not in st.session_state:
                st.session_state.database_collections = {}
            
            if current_db not in st.session_state.database_collections:
                st.session_state.database_collections[current_db] = {}
            
            # Generate a new collection ID for display purposes
            display_collection_id = str(int(time.time())) if not collection_id else collection_id
            
            # Store collection metadata
            st.session_state.database_collections[current_db][display_collection_id] = {
                "collection_name": collection_id,
                "file_name": file.name,
                "file_type": file_extension,
                "added_date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "description": collection_description or ""
            }
            
            # Save to database
            files_db_path = os.path.join(current_db_path, "files.db")
            try:
                # Connect to SQLite database
                create_files_table = not os.path.exists(files_db_path)
                conn = sqlite3.connect(files_db_path)
                cursor = conn.cursor()
                
                # Create table if it doesn't exist
                if create_files_table:
                    cursor.execute('''
                    CREATE TABLE IF NOT EXISTS files (
                        id TEXT PRIMARY KEY,
                        file_name TEXT,
                        file_type TEXT,
                        collection_name TEXT,
                        added_date TEXT,
                        description TEXT
                    )
                    ''')
                
                # Insert into database
                cursor.execute(
                    "INSERT OR REPLACE INTO files VALUES (?, ?, ?, ?, ?, ?)",
                    (
                        display_collection_id,
                        file.name,
                        file_extension,
                        collection_id,
                        datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                        collection_description or ""
                    )
                )
                
                conn.commit()
                conn.close()
                
            except Exception as db_error:
                logger.error(f"Error saving collection metadata to database: {str(db_error)}")
                if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                    logger.error(traceback.format_exc())
            
            st.success(f"File '{file.name}' has been processed and added to collection '{collection_name}'")
            return True
            
        except Exception as process_error:
            logger.error(f"Error processing file: {str(process_error)}")
            if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
                logger.error(traceback.format_exc())
            st.error(f"Error processing file '{file.name}': {str(process_error)}")
            return False
            
    except Exception as e:
        logger.error(f"Unexpected error in process_file_with_collection: {str(e)}")
        if hasattr(st.session_state, "debug_mode") and st.session_state.debug_mode:
            logger.error(traceback.format_exc())
        st.error(f"Unexpected error processing file '{file.name}': {str(e)}")
        return False
    finally:
        # Clean up
        try:
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)
            if os.path.exists(temp_dir):
                os.rmdir(temp_dir)
        except Exception as cleanup_error:
            logger.warning(f"Failed to clean up temporary files: {str(cleanup_error)}")

# Function to generate LLM response
def generate_llm_response(prompt, chat_history, context=""):
    """Generate a response from the LLM with streaming support"""
    result_queue = Queue()
    
    # Make local copies of everything needed before starting the thread
    # This completely isolates the thread from needing st.session_state
    if "chat_instance" not in st.session_state or st.session_state.chat_instance is None:
        return Queue_with_error("Chat instance not initialized. Please refresh the page and try again.")
        
    # Store all needed data in local variables to avoid st.session_state in thread
    local_system_prompt = ""
    if context:
        local_system_prompt = f"You are a helpful chatbot assistant designed to answer questions about the given context. Context: {context}"
    else:
        local_system_prompt = "You are a helpful chatbot assistant designed to answer questions."
    
    # Get a reference to the model outside the thread
    # We must set the system prompt before the thread starts
    chat_instance = st.session_state.chat_instance
    chat_instance.system_prompt = local_system_prompt
    
    # Add additional timeout to make sure the thread terminates
    response_timeout = 120  # seconds
    
    # Model name for logging only
    model_name = getattr(chat_instance, "model", "unknown")
    
    # Define a helper function to create a queue with an error
    def Queue_with_error(error_message):
        error_queue = Queue()
        error_queue.put(("error", error_message))
        return error_queue
    
    # Simple thread function that has zero Streamlit dependencies
    def generate_in_thread():
        try:
            # Define a simple async function that only streams from chat_async
            async def stream_response():
                try:
                    # Add a timeout to the async operation
                    async with asyncio.timeout(response_timeout):
                        # Direct call to chat_async with minimal processing
                        async for chunk in chat_instance.chat_async(
                            user_message=prompt,
                            previous_messages=chat_history,
                            context=context
                        ):
                            if chunk and chunk.strip():
                                result_queue.put(("chunk", chunk))
                        result_queue.put(("done", ""))
                except asyncio.TimeoutError:
                    result_queue.put(("error", f"Response generation timed out after {response_timeout} seconds"))
                except Exception as e:
                    # Don't use logger here to avoid StreamlitAPIException
                    result_queue.put(("error", f"Error in async streaming: {str(e)}"))
                    
            # Get a new event loop - create one if needed
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
            # Run the coroutine and close the loop when done
            try:
                loop.run_until_complete(stream_response())
            except Exception as loop_error:
                result_queue.put(("error", f"Loop error: {str(loop_error)}"))
            finally:
                loop.close()
                
        except Exception as e:
            # Don't use logger here to avoid StreamlitAPIException
            result_queue.put(("error", f"Thread error: {str(e)}"))
    
    # Log before starting thread (in the main thread where it's safe)
    logger.debug(f"Starting chat with model {model_name}")
    
    # Run in a separate thread to not block the UI
    thread = Thread(target=generate_in_thread, daemon=True)
    thread.start()
    
    return result_queue

# Function to show collections dropdown in sidebar
def show_collections_dropdown():
    """Display the current database collections in a dropdown menu"""
    if "database_collections" not in st.session_state:
        return
    
    current_db = st.session_state.current_database
    collections = st.session_state.database_collections.get(current_db, {})
    
    if not collections:
        st.info("No collections available in the current database.")
        return
    
    st.markdown("### Collections")
    
    # Extract collection names and details for display
    collection_options = {}
    for coll_name, coll_info in collections.items():
        # Prioritize collection name (which is the key) instead of file name
        # This shows users the collection names they entered
        display_name = coll_name
        collection_options[display_name] = coll_name
    
    # Create the dropdown
    selected_display_name = st.selectbox(
        "Select Collection:",
        options=list(collection_options.keys()),
        key="collection_dropdown"
    )
    
    # Get the actual collection name from the selected display name
    if selected_display_name:
        selected_collection = collection_options[selected_display_name]
        
        # Show collection details
        if st.button("View Collection Details", key="view_coll_details"):
            # Get the collection info
            coll_info = collections[selected_collection]
            
            st.markdown(f"**Collection:** {selected_collection}")
            
            # Show files count and list without using expanders
            if "files" in coll_info and len(coll_info["files"]) > 0:
                st.markdown(f"**Files:** {len(coll_info['files'])} document(s)")
                
                # Display first 5 files directly
                files_to_show = min(5, len(coll_info["files"]))
                for idx, file in enumerate(coll_info["files"][:files_to_show]):
                    st.markdown(f"- {file}")
                    
                # If there are more files, provide option to view more
                remaining_files = len(coll_info["files"]) - files_to_show
                if remaining_files > 0:
                    display_option = st.radio(
                        f"Show all files ({remaining_files} more)",
                        ["Hide Additional Files", "Show All Files"],
                        key=f"show_all_{selected_collection}",
                        horizontal=True
                    )
                    if display_option == "Show All Files":
                        for idx, file in enumerate(coll_info["files"][files_to_show:]):
                            st.markdown(f"- {file}")
            else:
                st.markdown(f"**File:** {coll_info.get('file_name', 'Unknown')}")
                
            st.markdown(f"**Type:** {coll_info.get('file_type', 'Unknown')}")
            st.markdown(f"**Added:** {coll_info.get('added_date', 'Unknown')}")
            
            # Show description if available
            if "description" in coll_info and coll_info["description"]:
                st.markdown(f"**Description:** {coll_info['description']}")
            
            # Inspect button
            if st.button("Inspect Collection Data", key="inspect_from_dropdown"):
                inspect_collection(selected_collection)
    
    return

def process_text_input(input_text, collection_name, collection_description=None):
    """
    Process text input from a textbox and insert it into a ChromaDB collection.
    
    Args:
        input_text (str): The text entered by the user
        collection_name (str): Name of the collection to insert the text into
        collection_description (str, optional): Description of the collection
    
    Returns:
        bool: True if successful, False otherwise
    """
    if not init_loader():
        st.error("Document loader could not be initialized. Cannot process text.")
        return False
    
    logger.info(f"Processing direct text input for collection: {collection_name}")
    
    # Get current database
    current_db = st.session_state.current_database
    
    # Get existing collections for the current database
    existing_collections = list(st.session_state.database_collections.get(current_db, {}).keys())
    
    # Add timestamp to make each text entry unique
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    temp_filename = f"text_input_{timestamp}.txt"
    
    try:
        # Get user-specific storage path
        user_storage_path = get_user_storage_path(st.session_state.username)
        current_db_path = os.path.join(user_storage_path, current_db)
        
        # Format the text with proper paragraph breaks for better processing
        formatted_text = input_text.replace('\n\n', '\n-----\n').replace('\n', ' ').replace('-----', '\n\n')
        
        # Create a temporary file to save the text
        with tempfile.NamedTemporaryFile(mode="w+", suffix=".txt", delete=False) as temp_file:
            temp_file.write(formatted_text)
            temp_file_path = temp_file.name
        
        # Log details for debugging
        logger.info(f"Created temporary file for text: {temp_file_path}")
        
        # Set persistent directory for the loader
        st.session_state.loader.persist_directory = current_db_path
        
        # Try to create or insert the collection
        try:
            # Notify user about processing
            st.info("Processing text... This may take a moment.")
            
            # Process the text file
            success = st.session_state.loader.create_or_insert_collection(
                temp_file_path,
                collection_name,
                extension="txt",
                parser_type='local_parser'
            )
            
            if not success:
                st.error("Failed to process text. No chunks could be generated.")
                return False
                
        except Exception as embedding_error:
            logger.error(f"Error generating embeddings: {str(embedding_error)}")
            st.error(f"Failed to generate embeddings for the text: {str(embedding_error)}")
            return False
        
        # Store the collection in session state
        if "file_collections" not in st.session_state:
            st.session_state.file_collections = {}
        
        # Use the collection name as the key
        st.session_state.file_collections[collection_name] = collection_name
        
        # Store in database collections
        if "database_collections" not in st.session_state:
            st.session_state.database_collections = {}
            
        if current_db not in st.session_state.database_collections:
            st.session_state.database_collections[current_db] = {}
        
        # Create or update collection metadata
        if collection_name in st.session_state.database_collections[current_db]:
            # Update existing collection metadata
            collection_metadata = st.session_state.database_collections[current_db][collection_name]
            # Update the modified date
            collection_metadata["modified_date"] = time.strftime("%Y-%m-%d %H:%M:%S")
            # Track number of entries
            collection_metadata["entries"] = collection_metadata.get("entries", 1) + 1
        else:
            # Create new collection metadata
            collection_metadata = {
                "collection_name": collection_name,
                "file_name": temp_filename,
                "file_type": "txt",
                "added_date": time.strftime("%Y-%m-%d %H:%M:%S"),
                "entries": 1
            }
            
            # Add description if provided
            if collection_description:
                collection_metadata["description"] = collection_description
        
        # Store collection metadata
        st.session_state.database_collections[current_db][collection_name] = collection_metadata
        
        logger.info(f"Text processed successfully and added to collection: {collection_name}")
        return True
        
    except Exception as e:
        logger.error(f"Error processing text input: {str(e)}\n{traceback.format_exc()}")
        st.error(f"Error processing text: {str(e)}")
        return False
    finally:
        # Clean up the temporary file
        try:
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)
        except Exception as cleanup_error:
            logger.warning(f"Failed to clean up temporary file: {str(cleanup_error)}")

if __name__ == "__main__":
    main()