rag-module 0.6.7

//! Enterprise RAG Module - Rust Implementation
//! 
//! A high-performance, secure RAG module with chat context storage, vector search,
//! and session management. Supports AWS, Azure, GCP estate data management.

pub mod config;
pub mod db;
pub mod services;
pub mod types;
pub mod utils;

use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::RwLock;
use anyhow::Result;
use tracing::{info, error, debug, warn};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tokio::fs;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use uuid::Uuid;
use sha2::{Sha256, Digest};

// Encrypted document structures for chat_history_documents.json
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptedChatDocument {
    pub id: String,
    #[serde(rename = "vectorId")]
    pub vector_id: String,
    pub content: String, // Base64 encrypted content
    pub embedding: Vec<f32>,
    pub metadata: EncryptedDocumentMetadata,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptedDocumentMetadata {
    #[serde(rename = "_encrypted_metadata")]
    pub encrypted_metadata: String, // Base64 encrypted metadata
    #[serde(rename = "_encrypted_content")]
    pub encrypted_content: bool,
    pub created_at: String, // ISO string format
    pub updated_at: String, // ISO string format
}

// Re-exports for public API
pub use config::ConfigManager;
pub use db::{VectorStore, LocalFileVectorStore};
pub use services::*;
pub use types::*;


#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionMessage {
    pub id: String,
    pub context_id: String,
    pub message_index: u32,
    pub role: MessageRole,
    pub content: String,
    pub timestamp: DateTime<Utc>,
    pub chat_title: String,
    pub metadata: Option<serde_json::Value>,
    pub request_id: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MessageRole {
    User,
    Assistant,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryResponsePair {
    pub query: SessionMessage,
    pub response: SessionMessage,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionHistoryResult {
    pub context_id: String,
    pub chat_title: String,
    pub pairs: Vec<QueryResponsePair>,
    pub total_pairs: usize,
    pub is_empty: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextInfo {
    pub context_id: String,
    pub chat_title: String,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatHistoryJson {
    pub context_id: String,
    pub chat_title: String,
    pub total_messages: usize,
    pub conversation: Vec<ConversationTurn>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConversationTurn {
    pub turn_number: usize,
    pub prompt: MessageJson,
    pub response: MessageJson,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageJson {
    pub message_id: String,
    pub content: String,
    pub timestamp: DateTime<Utc>,
    pub message_index: u32,
}



#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateResult {
    pub created: usize,
    pub failed: Vec<String>,
}

/// Main RAG Module - Business Architecture Implementation
/// Dual collection approach: chat_history (1D dummy) + aws_estate (1024D BGE-M3)
/// Provides unified chat history and estate resource management
#[derive(Clone)]
pub struct RagModule {
    base_path: PathBuf,
    initialized: Arc<RwLock<bool>>,

    // Core services
    pub config_manager: Arc<config::ConfigManager>,
    pub embedding_service: Arc<services::EmbeddingService>,
    pub vector_store: Arc<dyn db::VectorStore + Send + Sync>,
    pub document_service: Arc<services::DocumentService>,
    pub search_service: Arc<services::SearchService>,
    pub encryption_service: Arc<services::EncryptionService>,
    pub security_service: Arc<services::SecurityService>,
    pub mapping_service: Arc<services::MappingService>,
    pub sync_service: Arc<services::SyncService>,
    pub indexing_service: Arc<services::IndexingService>,
    pub kb_service: Arc<services::KBService>,

    // Business Architecture Services
    pub iam_service: Arc<services::IAMService>,
    pub operation_context_service: Arc<services::OperationContextService>,
    pub data_filtering_service: Arc<services::DataFilteringService>,
    pub collection_manager: Arc<services::CollectionManager>,
    pub aws_estate_service: Arc<services::AwsEstateService>,

    // Chat services - removed, functionality integrated into main RagModule

    // Session management
    messages_file: PathBuf,
    qdrant_data_path: PathBuf, // Base path for user-specific chat document folders
    aws_estate_file: PathBuf,
    message_counters: Arc<RwLock<HashMap<String, u32>>>,
}

impl RagModule {
    /// Create a new RAG module instance
    pub async fn new(base_path: impl Into<PathBuf>) -> Result<Self> {
        let base_path = base_path.into();
        
        // Initialize configuration
        let mut config_manager = config::ConfigManager::new(&base_path).await?;
        let mut config = config_manager.get_config();
        let mut config_updated = false;

        // Check for QDRANT_URL environment variable to set server sync URL
        if let Ok(qdrant_url) = std::env::var("QDRANT_URL") {
            info!("🔧 QDRANT_URL environment variable detected: {}", qdrant_url);

            // Use embedded mode with custom server sync URL
            config.vector_store.backend = "qdrant-embedded".to_string();
            config.vector_store.server_sync_url = Some(qdrant_url.clone());
            config.vector_store.enable_server_sync = true;
            config_updated = true;

            info!("✅ Using qdrant-embedded with server sync to: {}", qdrant_url);
        } else {
            // Auto-upgrade old configs from pure server mode to dual mode
            if config.vector_store.backend == "qdrant-server" {
                info!("🔄 Auto-upgrading from 'qdrant-server' to 'qdrant-embedded' with server sync");

                // Switch to embedded mode
                config.vector_store.backend = "qdrant-embedded".to_string();

                // Use existing server URL as sync URL
                if config.vector_store.server_sync_url.is_none() {
                    config.vector_store.server_sync_url = Some(config.vector_store.connection.url.clone());
                }

                config.vector_store.enable_server_sync = true;
                config_updated = true;

                info!("   ✅ Upgraded to DUAL mode: Embedded (local) + Server sync");
                info!("   📍 Local: ./qdrant-data");
                info!("   📡 Sync: {}", config.vector_store.server_sync_url.as_ref().unwrap());
            }
        }

        // Save config if updated
        if config_updated {
            config_manager.update_config(config.clone()).await?;
            info!("💾 Config saved");
        }

        // Now wrap in Arc after updating
        let config_manager = Arc::new(config_manager);
        
        // Check for RAG_DISABLE_ENCRYPTION environment variable to disable encryption for server-side operations
        if let Ok(val) = std::env::var("RAG_DISABLE_ENCRYPTION") {
            info!("🔧 RAG_DISABLE_ENCRYPTION environment variable detected: {} - Disabling encryption for server-side operations", val);
            config.encryption.enable_content_encryption = false;
            config.encryption.enable_metadata_encryption = false;
            config.encryption.enable_embedding_encryption = false;
            info!("✅ Encryption disabled: content={}, metadata={}, embedding={}", 
                config.encryption.enable_content_encryption,
                config.encryption.enable_metadata_encryption, 
                config.encryption.enable_embedding_encryption);
        } else {
            warn!("❌ RAG_DISABLE_ENCRYPTION environment variable not found - encryption remains enabled");
        }
        
        // Initialize encryption service first (needed by other services)
        let mut encryption_service = services::EncryptionService::new(&config.encryption, base_path.to_str().unwrap()).await?;
        encryption_service.initialize().await?;
        let encryption_service = Arc::new(encryption_service);
        
        // Initialize vector store based on configuration
        let vector_store: Arc<dyn db::VectorStore + Send + Sync> = match config.vector_store.backend.as_str() {
            "qdrant-embedded" => {
                // Check if server sync is enabled
                if config.vector_store.enable_server_sync {
                    if let Some(server_url) = &config.vector_store.server_sync_url {
                        info!("💾🔄 Using Qdrant DUAL mode: Embedded (primary) + Server sync ({})", server_url);

                        // Create embedded store
                        let embedded = db::EmbeddedQdrantVectorStore::new(
                            &base_path,
                            encryption_service.clone(),
                        ).await?;

                        // Create server store for syncing
                        let server = db::QdrantServerVectorStore::new(
                            server_url,
                            config.vector_store.connection.api_key.clone(),
                            &base_path.join("qdrant-data"),
                            encryption_service.clone(),
                        ).await?;

                        // Create dual store (embedded primary, server backup)
                        Arc::new(db::DualVectorStore::new(
                            embedded,
                            server,
                            false, // Don't fail writes if server sync fails
                        ).await?)
                    } else {
                        warn!("⚠️ Server sync enabled but no serverSyncUrl configured. Using embedded only.");
                        info!("💾 Using Qdrant embedded mode");
                        Arc::new(db::EmbeddedQdrantVectorStore::new(
                            &base_path,
                            encryption_service.clone(),
                        ).await?)
                    }
                } else {
                    info!("💾 Using Qdrant embedded mode (server sync disabled)");
                    Arc::new(db::EmbeddedQdrantVectorStore::new(
                        &base_path,
                        encryption_service.clone(),
                    ).await?)
                }
            }
            "qdrant-server" => {
                info!("🌐 Using Qdrant server mode: {}", config.vector_store.connection.url);
                Arc::new(db::QdrantServerVectorStore::new(
                    &config.vector_store.connection.url,
                    config.vector_store.connection.api_key.clone(),
                    &base_path.join("qdrant-data"),
                    encryption_service.clone(),
                ).await?)
            }
            "local-files" => {
                info!("📁 Using local file storage mode");
                Arc::new(db::LocalFileVectorStore::new(
                    &base_path.join("vector-data"),
                    encryption_service.clone(),
                ).await?)
            }
            _ => return Err(anyhow::anyhow!(
                "Unsupported vector store backend: '{}'. Supported backends: 'qdrant-embedded', 'qdrant-server', 'local-files'. \
                Use 'qdrant-embedded' for embedded mode (default) or set QDRANT_URL environment variable to use server mode.",
                config.vector_store.backend
            )),
        };
        
        // Initialize embedding service with models path
        let models_path = base_path.join("models");
        let embedding_service = {
            let service = services::EmbeddingService::new(&config.embedding, &models_path).await?;
            service.initialize().await?;
            Arc::new(service)
        };
        
        // Set dimensions in vector store
        let dimensions = embedding_service.get_dimensions().await?;
        vector_store.set_dimensions(dimensions).await?;
        
        // Initialize other core services with enhanced initialization
        let security_service = Arc::new(
            services::SecurityService::new(config_manager.as_ref().clone()).await?
        );
        
        let mapping_service = Arc::new(
            services::MappingService::new(&config.privacy).await?
        );
        
        let sync_service = Arc::new(
            services::SyncService::new(
                encryption_service.clone(),
                config_manager.clone()
            ).await?
        );

        let indexing_service = Arc::new(
            services::IndexingService::new(
                embedding_service.clone(),
                encryption_service.clone(),
            ).await?
        );

        let document_service = Arc::new(
            services::DocumentService::new(
                vector_store.clone(),
                indexing_service.clone(),
                mapping_service.clone(),
                security_service.clone(),
                encryption_service.clone(),
            ).await?
        );

        let kb_service = Arc::new(
            services::KBService::new().await?
        );
        
        // Business architecture services
        let iam_service = Arc::new(
            services::IAMService::new(&config.iam).await?
        );
        
        let operation_context_service = Arc::new(
            services::OperationContextService::new().await?
        );

        let data_filtering_service = Arc::new(
            services::DataFilteringService::new().await?
        );

        let collection_manager = Arc::new(
            services::CollectionManager::new(vector_store.clone())
        );
        
        // Initialize AWS Estate Service with parsers
        let mut aws_estate_service = services::AwsEstateService::new(document_service.clone());
        
        // Register AWS service parsers
        aws_estate_service.register_parser(Box::new(services::aws_parsers::Ec2Parser::new()));
        aws_estate_service.register_parser(Box::new(services::aws_parsers::RdsParser::new()));
        aws_estate_service.register_parser(Box::new(services::aws_parsers::S3Parser::new()));
        aws_estate_service.register_parser(Box::new(services::aws_parsers::LambdaParser::new()));
        aws_estate_service.register_parser(Box::new(services::aws_parsers::EbsParser::new()));
        aws_estate_service.register_parser(Box::new(services::aws_parsers::VpcParser::new()));
        aws_estate_service.register_parser(Box::new(services::aws_parsers::IamParser::new()));
        
        let aws_estate_service = Arc::new(aws_estate_service);

        // Initialize search service with enhanced services
        let mut search_service_instance = services::SearchService::new_with_services(
            vector_store.clone(),
            embedding_service.clone(),
            Some(security_service.clone()),
            Some(encryption_service.clone()),
            Some(mapping_service.clone()),
            None, // IAM service will be set later
        ).await?;

        // Set base path for local file access
        search_service_instance.set_base_path(base_path.clone());
        let search_service = Arc::new(search_service_instance);

        // Chat services - functionality now integrated into RagModule

   // Initialize message storage
        let storage_path = base_path.join("sessions");
        tokio::fs::create_dir_all(&storage_path).await?;
        let messages_file = storage_path.join("messages.jsonl");

        // Initialize paths for user-specific encrypted documents storage
        let qdrant_data_path = base_path.join("qdrant-data");
        // Don't create directory here - it's already created by vector store initialization
        let aws_estate_file = qdrant_data_path.join("aws_estate-metadata.json");

        Ok(Self {
            base_path,
            initialized: Arc::new(RwLock::new(false)),
            config_manager,
            embedding_service,
            vector_store,
            document_service,
            search_service,
            encryption_service,
            security_service,
            mapping_service,
            sync_service,
            indexing_service,
            kb_service,
            iam_service,
            operation_context_service,
            data_filtering_service,
            collection_manager,
            aws_estate_service,
            messages_file,
            qdrant_data_path,
            aws_estate_file,
            message_counters: Arc::new(RwLock::new(HashMap::new())),
            // estate_encryption_disabled,
        })
    }
    
    /// Initialize the RAG module
    pub async fn initialize(&self) -> Result<()> {
        let mut initialized = self.initialized.write().await;
        if *initialized {
            return Ok(());
        }

        info!("Initializing RAG Module at path: {:?}", self.base_path);

        // Create directory structure
        self.create_directory_structure().await?;

        // Initialize vector store
        self.vector_store.initialize().await?;
        
        // Initialize encryption service first (if not already initialized)
        if !self.encryption_service.is_initialized() {
            // Note: encryption service should already be initialized in new()
            warn!("Encryption service not initialized, initializing now");
        }
        
        // Initialize embedding service (required for estate collections)
        self.embedding_service.initialize().await?;
        
        // Initialize other services in dependency order
        self.document_service.initialize().await?;
        self.search_service.initialize().await?;
        self.security_service.initialize().await?;
        self.mapping_service.initialize().await?;
        self.sync_service.initialize().await?;
        self.indexing_service.initialize().await?;
        self.kb_service.initialize().await?;
        self.iam_service.initialize().await?;
        self.operation_context_service.initialize().await?;
        self.data_filtering_service.initialize().await?;
        self.collection_manager.initialize().await?;
        
        *initialized = true;
        info!("RAG Module initialized successfully");
        Ok(())
    }
    
    /// Check if the module is initialized
    pub async fn is_initialized(&self) -> bool {
        *self.initialized.read().await
    }

    /// Reset the RAG module state without full shutdown
    /// This clears embedding cache, resets GPU fallback state, and reloads models
    /// Useful between profile scans to prevent state corruption
    pub async fn reset_state(&self) -> Result<()> {
        info!("🔄 Resetting RAG Module state...");

        // Reset embedding service (clears cache, reloads model, resets force_cpu flag)
        self.embedding_service.reset_state().await?;

        // Clear document cache to force reload from disk (fixes stale cache after profile switches)
        info!("🗑️  Clearing vector store document cache...");
        self.vector_store.clear_document_cache().await?;

        info!("✅ RAG Module state reset complete");
        Ok(())
    }

    /// Create directory structure for the RAG module
    async fn create_directory_structure(&self) -> Result<()> {
        // Directories to create
        let directories = [
            "config",
            "data",
            "models",
            "cache",
            "sync",
            "keys",
            "qdrant-data", // Include qdrant-data directory
        ];

        // Create base directory first
        tokio::fs::create_dir_all(&self.base_path).await?;

        // Create all subdirectories
        for dir in directories.iter() {
            let dir_path = self.base_path.join(dir);
            tokio::fs::create_dir_all(&dir_path).await?;
            debug!("Created directory: {:?}", dir_path);
        }

        // Create .gitignore file to exclude sensitive data
        let gitignore_content = r#"# RAG Module - Exclude sensitive data
keys/
sync/
cache/
*.key
*.encrypted
config.local.yaml
"#;

        let gitignore_path = self.base_path.join(".gitignore");
        tokio::fs::write(&gitignore_path, gitignore_content).await?;
        debug!("Created .gitignore file: {:?}", gitignore_path);

        info!("Directory structure created successfully");
        Ok(())
    }
    
    /// Add document to the RAG system
    pub async fn add_document(&self, collection_type: &str, document: Document) -> Result<String> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }
        
        debug!("Adding document to collection: {}", collection_type);
        self.document_service.add_document(collection_type, document).await
    }
    
    /// Search documents with user context
    pub async fn search(&self, collection_type: &str, query: &str, user_id: &str, options: SearchOptions) -> Result<Vec<SearchResult>> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }
        
        debug!("Searching in collection: {} with query: {} for user: {}", collection_type, query, user_id);
        
        // Route to appropriate search method based on collection type
        if collection_type == "chat_history" || collection_type == "chat" {
            let chat_options = services::search_service::ChatSearchOptions {
                context_id: None,
                role: None,
                from_timestamp: None,
                to_timestamp: None,
                from_message_index: None,
                to_message_index: None,
                limit: options.limit,
                include_metadata: false,
                user_id: Some(user_id.to_string()),
            };
            let results = self.search_service.search_chat_history(chat_options).await?;
            // Convert JsonValue results to SearchResult for compatibility
            Ok(results.into_iter().enumerate().map(|(i, result)| {
                let payload = if let serde_json::Value::Object(map) = result {
                    Some(map.into_iter().collect())
                } else {
                    None
                };
                types::SearchResult {
                    id: payload.as_ref()
                        .and_then(|p: &std::collections::HashMap<String, serde_json::Value>| p.get("id"))
                        .and_then(|v| v.as_str())
                        .unwrap_or(&i.to_string()).to_string(),
                    score: payload.as_ref()
                        .and_then(|p: &std::collections::HashMap<String, serde_json::Value>| p.get("score"))
                        .and_then(|v| v.as_f64())
                        .unwrap_or(0.0) as f32,
                    document: None,
                    payload,
                }
            }).collect())
        } else if collection_type == "escher_library" || collection_type.starts_with("tenant_") {
            info!("📚 ROUTE: Playbook Search ({})", collection_type);
            info!("🔍 Calling search_service.search_playbooks()...");

            let results = self.search_service.search_playbooks(collection_type, query, &options).await?;
            info!("✅ Playbook search completed: {} results", results.len());
            Ok(results)
        } else if collection_type == "aws_estate" || collection_type.ends_with("_estate") {
            // Support dynamic estate collections: aws_estate, azure_estate, gcp_estate, etc.
            info!("🏗️ ROUTE: Estate Search ({})", collection_type);
            // Keep reference to parameters for later SQLite field extraction
            let requested_parameters = options.parameters.clone();

            let estate_options = services::search_service::EstateSearchOptions {
                resource_types: None,
                account_ids: None,
                regions: None,
                services: None,
                states: None,
                environment: None,
                application: None,
                synced_after: None,
                limit: options.limit,
                score_threshold: options.score_threshold,
                include_metadata: false,
                use_anonymous_ids: true,
                parameters: options.parameters, // Will return all fields by default
            };
            let results = self.search_service.search_estate_resources(collection_type, query, estate_options, None, user_id).await?;

            // Convert JsonValue results to SearchResult and extract fields from encrypted_metadata
            Ok(results.into_iter().enumerate().map(|(i, result)| {
                let mut payload: Option<std::collections::HashMap<String, serde_json::Value>> = if let serde_json::Value::Object(map) = result {
                    Some(map.into_iter().collect())
                } else {
                    None
                };

                // Parse and extract fields from Qdrant's _encrypted_metadata field (v0.5.9 format)
                if let Some(ref mut p) = payload {
                    // Get _encrypted_metadata field from Qdrant payload
                    if let Some(encrypted_metadata_str) = p.get("_encrypted_metadata").and_then(|v| v.as_str()) {
                        // Parse the JSON string from Qdrant
                        match serde_json::from_str::<serde_json::Value>(encrypted_metadata_str) {
                            Ok(metadata_json) => {
                                if let Some(metadata_obj) = metadata_json.as_object() {
                                    // Extract fields based on parameters
                                    if let Some(ref params) = requested_parameters {
                                        // Only extract requested fields
                                        info!("📋 Extracting {} requested fields from Qdrant metadata", params.len());
                                        for (field_name, _) in params {
                                            if let Some(field_value) = metadata_obj.get(field_name) {
                                                p.insert(field_name.clone(), field_value.clone());
                                                debug!("  ✅ Extracted field '{}' from Qdrant", field_name);
                                            } else {
                                                debug!("  ⚠️ Field '{}' not found in Qdrant metadata", field_name);
                                            }
                                        }
                                    } else {
                                        // No parameters specified - extract all fields from Qdrant
                                        info!("📋 Extracting all {} fields from Qdrant metadata", metadata_obj.len());
                                        for (key, value) in metadata_obj {
                                            p.insert(key.clone(), value.clone());
                                        }
                                    }
                                    debug!("✅ Extracted metadata from Qdrant for result {}", i);
                                } else {
                                    warn!("⚠️ Qdrant metadata for result {} is not a JSON object", i);
                                }
                            },
                            Err(e) => {
                                warn!("⚠️ Failed to parse Qdrant metadata for result {}: {}", i, e);
                            }
                        }
                    } else {
                        debug!("⚠️ No _encrypted_metadata field found in Qdrant payload for result {}", i);
                    }
                }

                types::SearchResult {
                    id: payload.as_ref()
                        .and_then(|p: &std::collections::HashMap<String, serde_json::Value>| p.get("id"))
                        .and_then(|v: &serde_json::Value| v.as_str())
                        .unwrap_or(&i.to_string()).to_string(),
                    score: payload.as_ref()
                        .and_then(|p: &std::collections::HashMap<String, serde_json::Value>| p.get("score"))
                        .and_then(|v: &serde_json::Value| v.as_f64())
                        .unwrap_or(0.0) as f32,
                    document: None,
                    payload,
                }
            }).collect())
        } else {
            Err(anyhow::anyhow!("Unsupported collection type: {}", collection_type))
        }
    }
    
    /// Get document by ID
    pub async fn get_document(&self, collection_type: &str, id: &str) -> Result<Option<Document>> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }
        
        self.document_service.get_document(collection_type, id).await
    }
    
    /// Process AWS estate data
    pub async fn process_aws_estate(&self, data: serde_json::Value, user_id: &str, collection_name: &str) -> Result<Vec<String>> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        info!("Processing AWS estate data for user: {}", user_id);
        let document_ids = self.iam_service.process_estate_data(data.clone()).await?;

        // Save AWS estate data as encrypted documents for the specific user
        self.save_aws_estate_documents(&data, &document_ids, user_id, collection_name).await?;

        Ok(document_ids)
    }

    /// Generate deterministic document ID from ARN (prevents duplicates)
    ///
    /// Uses SHA256 hash of the ARN to create a stable, collision-resistant ID.
    /// Same ARN always produces same ID, enabling proper deduplication.
    fn generate_doc_id_from_arn(collection_name: &str, metadata_obj: &serde_json::Map<String, serde_json::Value>) -> String {
        if let Some(arn) = metadata_obj.get("arnId") {
            if let Some(arn_str) = arn.as_str() {
                // Hash the ARN to create deterministic ID
                let mut hasher = Sha256::new();
                hasher.update(arn_str.as_bytes());
                let hash_result = hasher.finalize();
                let hash_hex = format!("{:x}", hash_result);

                // Use first 16 chars of hash (64 bits - very low collision probability)
                let doc_id = format!("{}-{}", collection_name, &hash_hex[..16]);
                info!("🆔 Generated deterministic doc_id from ARN: {} -> {}", arn_str, doc_id);
                return doc_id;
            } else {
                warn!("⚠️ arnId field exists but is not a string, using UUID fallback");
            }
        } else {
            warn!("⚠️ arnId field missing from metadata, using UUID fallback");
        }

        // Fallback to UUID if arnId is missing or invalid
        let fallback_id = format!("{}-{}", collection_name, uuid::Uuid::new_v4());
        warn!("🔄 Using fallback UUID doc_id: {}", fallback_id);
        fallback_id
    }

    /// Ingest AWS estate data - simplified to just store string as one point with embedding
    pub async fn ingest_aws_estate(&self, estate_data: serde_json::Value, user_id: &str, collection_name:&str) -> Result<services::AwsEstateIngestResult> {
        let start_time = std::time::Instant::now();
        info!("🚀 Starting AWS estate ingestion - user_id: {}, collection: {}", user_id, collection_name);

        // Step 1: Initialize
        if !self.is_initialized().await {
            info!("📦 RAG module not initialized, initializing...");
            match self.initialize().await {
                Ok(_) => info!("✅ RAG module initialized successfully"),
                Err(e) => {
                    error!("❌ Failed to initialize RAG module: {}", e);
                    return Err(e);
                }
            }
        }

        // Step 2: Set user context
        info!("👤 Setting user context for user_id: {}", user_id);
        if let Err(e) = self.set_user_context(user_id).await {
            error!("❌ Failed to set user context: {}", e);
            return Err(e);
        }

        // Step 3: Extract content field
        info!("📄 Extracting content from estate data...");
        let content_str = match estate_data.get("content") {
            Some(content) => {
                match content.as_str() {
                    Some(s) => {
                        let content_len = s.len();
                        info!("✅ Content extracted successfully (length: {} chars)", content_len);
                        s.to_string()
                    },
                    None => {
                        error!("❌ Content field exists but is not a string: {:?}", content);
                        return Err(anyhow::anyhow!("content field must be a string"));
                    }
                }
            },
            None => {
                error!("❌ Content field missing from estate_data. Available fields: {:?}",
                    estate_data.as_object().map(|o| o.keys().collect::<Vec<_>>()));
                return Err(anyhow::anyhow!("content field is required in the JSON object"));
            }
        };

        // Step 4: Generate embedding
        info!("🧮 Generating embedding for content...");
        let embedding = match self.embedding_service.generate_embedding(&content_str).await {
            Ok(emb) => {
                info!("✅ Embedding generated successfully (dimensions: {})", emb.len());
                emb
            },
            Err(e) => {
                error!("❌ Failed to generate embedding: {}", e);
                return Err(e);
            }
        };

        // Step 5: Create metadata first for canonical ID generation
        info!("📋 Processing metadata...");
        let mut metadata_obj = match estate_data.as_object() {
            Some(obj) => {
                info!("✅ Metadata object extracted ({} fields)", obj.len());
                let mut metadata_obj = obj.clone();
                metadata_obj.remove("content");
                metadata_obj
            },
            None => {
                error!("❌ estate_data is not a JSON object");
                return Err(anyhow::anyhow!("estate_data must be a JSON object"));
            }
        };

        // Step 6: Generate deterministic document ID from ARN (enables deduplication)
        let doc_id = Self::generate_doc_id_from_arn(collection_name, &metadata_obj);

        let metadata_str = match serde_json::to_string(&metadata_obj) {
            Ok(s) => {
                info!("✅ Metadata serialized successfully (length: {} bytes)", s.len());
                s
            },
            Err(e) => {
                error!("❌ Failed to serialize metadata: {}", e);
                return Err(anyhow::anyhow!("Failed to serialize metadata: {}", e));
            }
        };

        // Step 6.5: Create metadata for Qdrant with filterable fields + full metadata
        // Store full metadata in _encrypted_metadata field (prevents vector store from wrapping it)
        let mut metadata = indexmap::IndexMap::new();
        metadata.insert("_encrypted_metadata".to_string(), serde_json::Value::String(metadata_str));

        // Add filterable fields as separate top-level fields for Qdrant pre-filtering
        if let Some(val) = metadata_obj.get("profile") {
            metadata.insert("profile".to_string(), val.clone());
        }
        if let Some(val) = metadata_obj.get("service") {
            metadata.insert("service".to_string(), val.clone());
        }

        info!("📋 Qdrant metadata ready (hybrid: _encrypted_metadata + filterable fields [profile, service])");

        // Step 7: Create and store document in Qdrant (with v0.5.9 metadata format)
        info!("💾 Creating and storing document in Qdrant collection '{}'...", collection_name);
        let document = types::Document::new(doc_id.clone(), content_str)
            .with_embedding(embedding)
            .with_metadata(metadata);

        match self.vector_store.add_document(collection_name, document).await {
            Ok(_) => {
                let elapsed = start_time.elapsed();
                info!("✅ Document stored successfully (doc_id: {}) in {:?}", doc_id, elapsed);
            },
            Err(e) => {
                error!("❌ Failed to store document in Qdrant: {}", e);
                return Err(e);
            }
        }

        // Return simple result
        let result = services::AwsEstateIngestResult {
            total_accounts: 1,
            total_services: 1,
            total_resources: 1,
            parsed_resources: 1,
            failed_resources: 0,
            supported_services: vec![],
            unsupported_services: vec![],
            create_result: CreateResult {
                created: 1,
                failed: Vec::new(),
            },
        };

        let total_time = start_time.elapsed();
        info!("🎉 AWS estate ingestion completed successfully in {:?}", total_time);
        Ok(result)
    }

    /// Ingest AWS estate data in batch - optimized for processing multiple documents at once
    pub async fn ingest_aws_estate_batch(&self, estate_data_batch: Vec<serde_json::Value>, user_id: &str, collection_name: &str) -> Result<services::AwsEstateIngestResult> {
        let start_time = std::time::Instant::now();
        let batch_size = estate_data_batch.len();
        info!("🚀 Starting batch AWS estate ingestion - user_id: {}, collection: {}, batch_size: {}",
            user_id, collection_name, batch_size);

        // Step 1: Initialize
        if !self.is_initialized().await {
            info!("📦 RAG module not initialized, initializing...");
            match self.initialize().await {
                Ok(_) => info!("✅ RAG module initialized successfully"),
                Err(e) => {
                    error!("❌ Failed to initialize RAG module: {}", e);
                    return Err(e);
                }
            }
        }

        // Step 2: Set user context
        info!("👤 Setting user context for user_id: {}", user_id);
        if let Err(e) = self.set_user_context(user_id).await {
            error!("❌ Failed to set user context: {}", e);
            return Err(e);
        }

        let mut documents = Vec::new();
        let mut content_strings = Vec::new();
        let mut failed_resources = 0;
        let mut failed_errors = Vec::new();

        // Step 3: Process each document in the batch
        info!("📄 Processing {} documents in batch...", batch_size);
        for (idx, estate_data) in estate_data_batch.iter().enumerate() {
            debug!("Processing document {}/{}", idx + 1, batch_size);
            // Extract content field from the JSON
            let content_str = match estate_data.get("content") {
                Some(content) => {
                    match content.as_str() {
                        Some(s) => s.to_string(),
                        None => {
                            failed_resources += 1;
                            let error_msg = format!("Document {}/{}: content field must be a string", idx + 1, batch_size);
                            error!("❌ {}", error_msg);
                            failed_errors.push(error_msg);
                            continue;
                        }
                    }
                },
                None => {
                    failed_resources += 1;
                    let error_msg = format!("Document {}/{}: content field is required", idx + 1, batch_size);
                    error!("❌ {}", error_msg);
                    failed_errors.push(error_msg);
                    continue;
                }
            };

            // Create metadata with all fields except "content" first
            let metadata_obj = match estate_data.as_object() {
                Some(obj) => {
                    let mut metadata_obj = obj.clone();
                    metadata_obj.remove("content");
                    metadata_obj
                },
                None => {
                    failed_resources += 1;
                    let error_msg = format!("Document {}/{}: estate_data must be a JSON object", idx + 1, batch_size);
                    error!("❌ {}", error_msg);
                    failed_errors.push(error_msg);
                    continue;
                }
            };

            // Generate deterministic document ID from ARN (enables deduplication)
            let doc_id = Self::generate_doc_id_from_arn(collection_name, &metadata_obj);
            debug!("🆔 Document {}/{}: Generated doc_id: {}", idx + 1, batch_size, doc_id);

            let metadata_str = match serde_json::to_string(&metadata_obj) {
                Ok(s) => s,
                Err(e) => {
                    failed_resources += 1;
                    let error_msg = format!("Document {}/{}: Failed to serialize metadata: {}", idx + 1, batch_size, e);
                    error!("❌ {}", error_msg);
                    failed_errors.push(error_msg);
                    continue;
                }
            };

            // Create metadata for Qdrant with filterable fields + full metadata
            // Store full metadata in _encrypted_metadata field (prevents vector store from wrapping it)
            let mut metadata = indexmap::IndexMap::new();
            metadata.insert("_encrypted_metadata".to_string(), serde_json::Value::String(metadata_str));

            // Add filterable fields as separate top-level fields for Qdrant pre-filtering
            if let Some(val) = metadata_obj.get("profile") {
                metadata.insert("profile".to_string(), val.clone());
            }
            if let Some(val) = metadata_obj.get("service") {
                metadata.insert("service".to_string(), val.clone());
            }

            // Store document data for batch processing (with hybrid metadata format)
            documents.push((doc_id, content_str.clone(), metadata));
            content_strings.push(content_str);
        }

        let processed_count = documents.len();
        info!("✅ Successfully processed {}/{} documents", processed_count, batch_size);

        if documents.is_empty() {
            error!("❌ No valid documents to process. All {} documents failed.", batch_size);
            error!("Errors encountered: {:?}", failed_errors);
            return Ok(services::AwsEstateIngestResult {
                total_accounts: 0,
                total_services: 0,
                total_resources: batch_size,
                parsed_resources: 0,
                failed_resources,
                supported_services: vec![],
                unsupported_services: vec![],
                create_result: CreateResult {
                    created: 0,
                    failed: failed_errors,
                },
            });
        }

        // Step 4: Generate embeddings for all content strings in batch
        info!("🧮 Generating embeddings for {} documents...", processed_count);
        let embedding_start = std::time::Instant::now();
        let embeddings = match self.embedding_service.generate_embeddings_batch(&content_strings).await {
            Ok(embs) => {
                let embedding_duration = embedding_start.elapsed();
                info!("✅ Batch embeddings generated successfully ({} embeddings in {:?})", embs.len(), embedding_duration);
                embs
            },
            Err(e) => {
                // Fallback to individual embedding generation if batch fails
                warn!("⚠️ Batch embedding generation failed, falling back to individual generation: {}", e);
                let mut individual_embeddings = Vec::new();
                for (idx, content) in content_strings.iter().enumerate() {
                    match self.embedding_service.generate_embedding(content).await {
                        Ok(embedding) => {
                            debug!("✅ Generated embedding for document {}/{}", idx + 1, content_strings.len());
                            individual_embeddings.push(embedding);
                        },
                        Err(e) => {
                            failed_resources += 1;
                            let error_msg = format!("Document {}/{}: Failed to generate embedding: {}", idx + 1, content_strings.len(), e);
                            error!("❌ {}", error_msg);
                            failed_errors.push(error_msg);
                            continue;
                        }
                    }
                }
                let embedding_duration = embedding_start.elapsed();
                info!("✅ Individual embeddings generated ({} successful, {} failed in {:?})",
                    individual_embeddings.len(), content_strings.len() - individual_embeddings.len(), embedding_duration);
                individual_embeddings
            }
        };

        // Step 5: Create final documents with embeddings
        info!("📦 Creating final documents with embeddings...");
        let mut final_documents = Vec::new();
        for (i, (doc_id, content_str, metadata)) in documents.into_iter().enumerate() {
            if i < embeddings.len() {
                let document = types::Document::new(doc_id, content_str)
                    .with_embedding(embeddings[i].clone())
                    .with_metadata(metadata);
                final_documents.push(document);
            } else {
                failed_resources += 1;
                let error_msg = format!("Document {}: Missing embedding (embedding count mismatch)", i + 1);
                error!("❌ {}", error_msg);
                failed_errors.push(error_msg);
            }
        }
        info!("✅ Created {} final documents ready for storage", final_documents.len());

        // Step 5.5: Disable optimizer to prevent data loss during bulk insert
        info!("🛑 Disabling optimizer before bulk insert (prevents segment merge corruption)");
        if let Err(e) = self.vector_store.disable_optimizer(collection_name).await {
            warn!("⚠️ Failed to disable optimizer (non-critical): {}", e);
        }

        // Step 6: Add documents to vector store using batch insertion
        info!("💾 Storing {} documents to collection '{}'...", final_documents.len(), collection_name);
        let storage_start = std::time::Instant::now();
        let created = match self.vector_store.add_documents(collection_name, final_documents.clone()).await {
            Ok(document_ids) => {
                let count = document_ids.len();
                let storage_duration = storage_start.elapsed();
                info!("✅ Successfully stored {} documents to vector store in {:?}", count, storage_duration);
                count
            }
            Err(e) => {
                let storage_duration = storage_start.elapsed();
                error!("❌ Batch insertion failed after {:?}: {}", storage_duration, e);
                error!("   Collection: {}", collection_name);
                error!("   Attempted to store: {} documents", final_documents.len());
                error!("   Error details: {}", e);
                failed_resources += final_documents.len();
                failed_errors.push(format!("Batch insertion to vector store failed: {}", e));
                0
            }
        };

        // Step 6.5: TEMPORARY FIX - Re-enable optimizer after bulk insert
        info!("✅ Re-enabling optimizer after bulk insert");
        if let Err(e) = self.vector_store.enable_optimizer(collection_name).await {
            warn!("⚠️ Failed to re-enable optimizer (non-critical): {}", e);
        }

        // Return batch result
        let total_duration = start_time.elapsed();
        info!("🎉 Batch AWS estate ingestion completed in {:?}", total_duration);
        info!("   Total resources: {}", estate_data_batch.len());
        info!("   Successfully created: {}", created);
        info!("   Failed: {}", failed_resources);

        if failed_resources > 0 {
            warn!("⚠️ {} documents failed during batch ingestion:", failed_resources);
            for (idx, error) in failed_errors.iter().enumerate() {
                warn!("   {}. {}", idx + 1, error);
            }
        }

        let result = services::AwsEstateIngestResult {
            total_accounts: 1,
            total_services: 1,
            total_resources: estate_data_batch.len(),
            parsed_resources: created,
            failed_resources,
            supported_services: vec![],
            unsupported_services: vec![],
            create_result: CreateResult {
                created,
                failed: failed_errors,
            },
        };

        Ok(result)
    }
    
    // Chat functionality now integrated into session management methods:
    // - Use add_prompt() and add_response() instead of add_chat_message()
    // - Use get_session_chat_history() instead of get_chat_history()

    // Session Management Methods

    /// Start a new chat session

    /// Add a prompt to the current session
    pub async fn add_prompt(&self, context_id: &str, prompt: &str, user_id: &str, chat_title: Option<&str>, request_id: &str) -> Result<String> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        if context_id.is_empty() || prompt.is_empty() {
            return Err(anyhow::anyhow!("context_id and prompt are required"));
        }

        let message_index = self.get_next_message_index(context_id).await;
        let title = chat_title.unwrap_or("Chat").to_string();

        let chat_message = SessionMessage {
            id: Uuid::new_v4().to_string(),
            context_id: context_id.to_string(),
            message_index,
            role: MessageRole::User,
            content: prompt.to_string(),
            timestamp: Utc::now(),
            chat_title: title.clone(),
            metadata: None,
            request_id: Some(request_id.to_string()),
        };

        // Create document content JSON that will be encrypted by the vector store
        let document_content = serde_json::json!({
            "message_id": chat_message.id,
            "i": context_id,
            "r": "0",
            "c": prompt,
            "t": chat_message.timestamp,
            "m": message_index,
            "ct": title,
            "ri": request_id  // Client-provided request_id
        });

        let content_str = serde_json::to_string(&document_content)?;

        // Create metadata that will be encrypted by the vector store
        let metadata = serde_json::json!({
            "i": context_id,
            "m": message_index,
            "r": "0",
            "t": chat_message.timestamp.to_rfc3339(),
            "type": "chat_prompt",
            "ct": title
        });

        // Convert metadata JSON to IndexMap for Document
        let mut metadata_map = indexmap::IndexMap::new();
        if let Some(obj) = metadata.as_object() {
            for (k, v) in obj {
                metadata_map.insert(k.clone(), v.clone());
            }
        }

        // Create document for Qdrant storage - vector store will handle encryption
        let document = types::Document::new_with_vector_id(
            chat_message.id.clone(), // Use chat message ID as document ID
            Uuid::new_v4().to_string(), // Generate separate vectorId
            content_str, // Plain content - will be encrypted by vector store
        )
        .with_metadata(metadata_map)
        .with_embedding(vec![0.1]); // Use 0.1 dummy embedding for chat

        // Set user context for vector store
        if let Some(server_store) = self.vector_store.as_any().downcast_ref::<db::QdrantServerVectorStore>() {
            server_store.set_user_context(user_id).await;
        } else if let Some(dual_store) = self.vector_store.as_any().downcast_ref::<db::DualVectorStore>() {
            dual_store.set_user_context(user_id).await;
        } else if let Some(embedded_store) = self.vector_store.as_any().downcast_ref::<db::EmbeddedQdrantVectorStore>() {
            embedded_store.set_user_context(user_id).await;
        } else {
            return Err(anyhow::anyhow!("Unsupported vector store type for chat operations"));
        }

        // Store in vector store like Node.js: await this.vectorStore.insertToCollection('chat', chatDoc)
        self.vector_store.add_document("chat_history", document).await?;

        // Store message in file
        self.store_message(&chat_message).await?;

        debug!("Added prompt for context: {}", context_id);
        Ok(chat_message.id)
    }

    /// Add a response to a chat context
    ///
    /// # Arguments
    /// * `context_id` - The context ID for this conversation
    /// * `raw_response` - The assistant's response text (may include streaming artifacts)
    /// * `user_id` - The user ID
    /// * `chat_title` - Optional chat title (defaults to "Chat")
    pub async fn add_response(&self, context_id: &str, raw_response: &str, user_id: &str, chat_title: Option<&str>, request_id: &str) -> Result<String> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        if context_id.is_empty() || raw_response.is_empty() {
            return Err(anyhow::anyhow!("context_id and rawResponse are required"));
        }

        let clean_content = self.extract_content_from_streaming_response(raw_response);
        let message_index = self.get_next_message_index(context_id).await;
        let title = chat_title.unwrap_or("Chat").to_string();

        let chat_message = SessionMessage {
            id: Uuid::new_v4().to_string(),
            context_id: context_id.to_string(),
            message_index,
            role: MessageRole::Assistant,
            content: clean_content.clone(),
            timestamp: Utc::now(),
            chat_title: title.clone(),
            metadata: None,
            request_id: Some(request_id.to_string()),
        };

        // Create document content JSON that will be encrypted by the vector store
        let document_content = serde_json::json!({
            "message_id": chat_message.id,
            "i": context_id,
            "r": "1",
            "c": clean_content,
            "t": chat_message.timestamp,
            "m": message_index,
            "ct": title,
            "ri": request_id  // Same request_id as the prompt
        });

        let content_str = serde_json::to_string(&document_content)?;

        // Create metadata that will be encrypted by the vector store
        let metadata = serde_json::json!({
            "i": context_id,
            "messageIndex": message_index,
            "r": "1",
            "t": chat_message.timestamp.to_rfc3339(),
            "type": "chat_response",
            "ct": title
        });

        // Convert metadata JSON to IndexMap for Document
        let mut metadata_map = indexmap::IndexMap::new();
        if let Some(obj) = metadata.as_object() {
            for (k, v) in obj {
                metadata_map.insert(k.clone(), v.clone());
            }
        }

        // Create document for Qdrant storage - vector store will handle encryption
        let document = types::Document::new_with_vector_id(
            chat_message.id.clone(), // Use chat message ID as document ID
            Uuid::new_v4().to_string(), // Generate separate vectorId
            content_str, // Plain content - will be encrypted by vector store
        )
        .with_metadata(metadata_map)
        .with_embedding(vec![0.1]); // Use 0.1 dummy embedding for chat

        // Set user context for vector store
        if let Some(server_store) = self.vector_store.as_any().downcast_ref::<db::QdrantServerVectorStore>() {
            server_store.set_user_context(user_id).await;
        } else if let Some(dual_store) = self.vector_store.as_any().downcast_ref::<db::DualVectorStore>() {
            dual_store.set_user_context(user_id).await;
        } else if let Some(embedded_store) = self.vector_store.as_any().downcast_ref::<db::EmbeddedQdrantVectorStore>() {
            embedded_store.set_user_context(user_id).await;
        } else {
            return Err(anyhow::anyhow!("Unsupported vector store type for chat operations"));
        }

        // Store in vector store like Node.js: await this.vectorStore.insertToCollection('chat', chatDoc)
        self.vector_store.add_document("chat_history", document).await?;

        // Store message in file
        self.store_message(&chat_message).await?;

        debug!("Added response for context: {}", context_id);
        Ok(chat_message.id)
    }

    /// Get session chat history
    pub async fn get_session_chat_history(&self, context_id: &str) -> Result<SessionHistoryResult> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        if context_id.is_empty() {
            return Err(anyhow::anyhow!("contextId is required"));
        }

        let messages = self.load_messages_by_context(context_id).await?;

        let mut prompts: Vec<SessionMessage> = messages
            .iter()
            .filter(|m| matches!(m.role, MessageRole::User))
            .cloned()
            .collect();

        let mut responses: Vec<SessionMessage> = messages
            .iter()
            .filter(|m| matches!(m.role, MessageRole::Assistant))
            .cloned()
            .collect();

        // Sort by message index or timestamp
        prompts.sort_by(|a, b| a.message_index.cmp(&b.message_index));
        responses.sort_by(|a, b| a.message_index.cmp(&b.message_index));

        // Create pairs
        let mut pairs = Vec::new();
        let pair_count = std::cmp::min(prompts.len(), responses.len());

        for i in 0..pair_count {
            pairs.push(QueryResponsePair {
                query: prompts[i].clone(),
                response: responses[i].clone(),
            });
        }

        Ok(SessionHistoryResult {
            context_id: context_id.to_string(),
            chat_title: "Retrieved Chat History".to_string(),
            pairs,
            total_pairs: pair_count,
            is_empty: pair_count == 0,
        })
    }


    /// Get query-response pairs for context retrieval
    pub async fn get_query_response_pairs(
        &self,
        context_id: &str,
        max_pairs: Option<usize>
    ) -> Result<SessionHistoryResult> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        if context_id.is_empty() {
            return Err(anyhow::anyhow!("contextId is required"));
        }

        let messages = self.load_messages_by_context(context_id).await?;

        let mut prompts: Vec<SessionMessage> = messages
            .iter()
            .filter(|m| matches!(m.role, MessageRole::User))
            .cloned()
            .collect();

        let mut responses: Vec<SessionMessage> = messages
            .iter()
            .filter(|m| matches!(m.role, MessageRole::Assistant))
            .cloned()
            .collect();

        prompts.sort_by(|a, b| a.message_index.cmp(&b.message_index));
        responses.sort_by(|a, b| a.message_index.cmp(&b.message_index));

        let pair_count = std::cmp::min(prompts.len(), responses.len());
        let actual_max_pairs = max_pairs.unwrap_or(pair_count).min(pair_count);

        let mut pairs = Vec::new();
        let start_index = pair_count.saturating_sub(actual_max_pairs);

        for i in start_index..pair_count {
            pairs.push(QueryResponsePair {
                query: prompts[i].clone(),
                response: responses[i].clone(),
            });
        }

        Ok(SessionHistoryResult {
            context_id: context_id.to_string(),
            chat_title: format!("Chat History ({} pairs)", pairs.len()),
            pairs,
            total_pairs: pair_count,
            is_empty: pair_count == 0,
        })
    }

    pub async fn get_all_contexts(&self, user_id: &str) -> Result<Vec<ContextInfo>> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        // Get all decrypted chat history
        let all_contexts_data = self.get_decrypted_chat_history(user_id).await?;

        // Convert to ContextInfo format
        let mut contexts: Vec<ContextInfo> = all_contexts_data.iter().filter_map(|ctx| {
            let context_id = ctx.get("i")?.as_str()?.to_string();
            let chat_title = ctx.get("ct")?.as_str().unwrap_or("Chat").to_string();
            let timestamp_str = ctx.get("t")?.as_str()?;

            // Parse timestamp
            let timestamp = chrono::DateTime::parse_from_rfc3339(timestamp_str)
                .ok()?
                .with_timezone(&chrono::Utc);

            Some(ContextInfo {
                context_id,
                chat_title,
                created_at: timestamp,
                updated_at: timestamp,
            })
        }).collect();

        // Sort by updated_at in descending order (most recent first)
        contexts.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));

        info!("Found {} unique contexts", contexts.len());
        Ok(contexts)
    }

    /// Get prompts and responses as JSON for a specific context_id
    ///
    /// Returns chat history for the given context_id in a structured format.
    ///
    /// # Arguments
    /// * `user_id` - The user ID
    /// * `context_id` - The context ID to retrieve chat history for
    ///
    /// # Returns
    /// * `serde_json::Value` - Chat history in format: {i, a, ct, m: [{r, c}], t, l}
    pub async fn get_chat_history_json(&self, user_id: &str, context_id: &str) -> Result<serde_json::Value> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        if context_id.is_empty() {
            return Err(anyhow::anyhow!("context_id is required"));
        }

        // Get all decrypted chat history for the user
        let all_contexts = self.get_decrypted_chat_history(user_id).await?;

        // Find the specific context
        let context_data = all_contexts.iter()
            .find(|ctx| {
                ctx.get("i")
                    .and_then(|v| v.as_str())
                    .map(|id| id == context_id)
                    .unwrap_or(false)
            });

        match context_data {
            Some(data) => {
                info!("Retrieved chat history for context: {}", context_id);
                Ok(data.clone())
            }
            None => {
                Err(anyhow::anyhow!("No messages found for context_id: {}", context_id))
            }
        }
    }

    /// Delete all documents and messages for a specific context_id from local chat history
    ///
    /// This function removes all chat history data associated with the given context_id from both:
    /// 1. The encrypted chat documents file (chat_history-documents.json)
    /// 2. The messages file (messages.jsonl)
    ///
    /// # Arguments
    /// * `context_id` - The context ID to delete
    /// * `user_id` - The user ID whose chat history to modify
    ///
    /// # Returns
    /// * `Result<usize>` - Number of documents deleted, or error if deletion failed
    pub async fn delete_context_from_chat_history(&self, context_id: &str, user_id: &str) -> Result<usize> {
        let start_time = std::time::Instant::now();
        info!("🗑️  Starting context deletion - context_id: {}, user_id: {}", context_id, user_id);

        if context_id.is_empty() {
            error!("❌ Context ID cannot be empty");
            return Err(anyhow::anyhow!("context_id cannot be empty"));
        }

        if user_id.is_empty() {
            error!("❌ User ID cannot be empty");
            return Err(anyhow::anyhow!("user_id cannot be empty"));
        }

        let mut total_deleted = 0;

        // Step 1: Delete from chat_history-documents.json
        info!("📄 Step 1: Processing chat_history-documents.json");
        let user_documents_file = self.get_user_encrypted_documents_path(user_id);

        if !user_documents_file.exists() {
            warn!("⚠️  Chat history file not found at: {}", user_documents_file.display());
            info!("   No documents to delete for user {}", user_id);
        } else {
            info!("✅ Found chat history file at: {}", user_documents_file.display());

            // Load encrypted documents
            let encrypted_docs = match self.get_encrypted_chat_documents(user_id).await {
                Ok(docs) => {
                    info!("✅ Loaded {} encrypted documents", docs.len());
                    docs
                }
                Err(e) => {
                    error!("❌ Failed to load encrypted documents: {}", e);
                    return Err(anyhow::anyhow!("Failed to load chat documents: {}", e));
                }
            };

            let original_count = encrypted_docs.len();
            info!("📊 Original document count: {}", original_count);

            // Filter documents by context_id
            let mut filtered_docs = Vec::new();
            let mut deleted_count = 0;

            for (idx, doc) in encrypted_docs.iter().enumerate() {
                // Decrypt content to check context_id
                let decrypted_content = match self.encryption_service.decrypt_content(&doc.content).await {
                    Ok(content) => content,
                    Err(e) => {
                        warn!("⚠️  Document {}/{}: Failed to decrypt content: {}", idx + 1, original_count, e);
                        warn!("   Skipping document (id: {})", doc.id);
                        // Keep the document if we can't decrypt (safer to not delete unknown data)
                        filtered_docs.push(doc.clone());
                        continue;
                    }
                };

                // Parse decrypted content as JSON
                let content_json = match serde_json::from_str::<serde_json::Value>(&decrypted_content) {
                    Ok(json) => json,
                    Err(e) => {
                        warn!("⚠️  Document {}/{}: Failed to parse decrypted content: {}", idx + 1, original_count, e);
                        warn!("   Skipping document (id: {})", doc.id);
                        // Keep the document if we can't parse (safer to not delete unknown data)
                        filtered_docs.push(doc.clone());
                        continue;
                    }
                };

                // Extract context_id from content (stored as "i" in compact format)
                let doc_context_id = content_json
                    .get("i")
                    .and_then(|v| v.as_str())
                    .unwrap_or("");

                if doc_context_id == context_id {
                    deleted_count += 1;
                    info!("🗑️  Document {}/{}: Marked for deletion (context_id: {}, doc_id: {})",
                        idx + 1, original_count, doc_context_id, doc.id);
                } else {
                    filtered_docs.push(doc.clone());
                }
            }

            if deleted_count == 0 {
                info!("ℹ️  No documents found matching context_id: {}", context_id);
            } else {
                info!("📊 Deletion summary:");
                info!("   - Documents to delete: {}", deleted_count);
                info!("   - Documents to keep: {}", filtered_docs.len());

                // Save filtered documents back to file
                info!("💾 Saving filtered documents to file...");

                let wrapper = serde_json::json!({
                    "documents": filtered_docs,
                    "count": filtered_docs.len(),
                    "lastModified": Utc::now().to_rfc3339(),
                });

                match serde_json::to_string_pretty(&wrapper) {
                    Ok(documents_content) => {
                        match fs::write(&user_documents_file, documents_content).await {
                            Ok(_) => {
                                let elapsed = start_time.elapsed();
                                info!("✅ Successfully saved {} documents (deleted {} documents) in {:?}",
                                    filtered_docs.len(), deleted_count, elapsed);
                                total_deleted += deleted_count;
                            }
                            Err(e) => {
                                error!("❌ Failed to write filtered documents to file: {}", e);
                                error!("   File path: {}", user_documents_file.display());
                                return Err(anyhow::anyhow!("Failed to save filtered documents: {}", e));
                            }
                        }
                    }
                    Err(e) => {
                        error!("❌ Failed to serialize filtered documents: {}", e);
                        return Err(anyhow::anyhow!("Failed to serialize documents: {}", e));
                    }
                }
            }
        }

        // Step 2: Delete from messages.jsonl
        info!("📄 Step 2: Processing messages.jsonl");
        if !self.messages_file.exists() {
            warn!("⚠️  Messages file not found at: {}", self.messages_file.display());
            info!("   No messages to delete");
        } else {
            info!("✅ Found messages file at: {}", self.messages_file.display());

            match fs::read_to_string(&self.messages_file).await {
                Ok(content) => {
                    let lines: Vec<&str> = content.lines().collect();
                    let original_message_count = lines.len();
                    info!("📊 Original message count: {}", original_message_count);

                    let mut filtered_lines = Vec::new();
                    let mut deleted_message_count = 0;

                    for (idx, line) in lines.iter().enumerate() {
                        if line.trim().is_empty() {
                            continue;
                        }

                        match serde_json::from_str::<serde_json::Value>(line) {
                            Ok(message) => {
                                let msg_context_id = message
                                    .get("context_id")
                                    .and_then(|v| v.as_str())
                                    .unwrap_or("");

                                if msg_context_id == context_id {
                                    deleted_message_count += 1;
                                    debug!("🗑️  Message {}/{}: Marked for deletion (context_id: {})",
                                        idx + 1, original_message_count, msg_context_id);
                                } else {
                                    filtered_lines.push(line.to_string());
                                }
                            }
                            Err(e) => {
                                warn!("⚠️  Message {}/{}: Failed to parse JSON: {}", idx + 1, original_message_count, e);
                                warn!("   Keeping message (safer to not delete unparseable data)");
                                filtered_lines.push(line.to_string());
                            }
                        }
                    }

                    if deleted_message_count == 0 {
                        info!("ℹ️  No messages found matching context_id: {}", context_id);
                    } else {
                        info!("📊 Message deletion summary:");
                        info!("   - Messages to delete: {}", deleted_message_count);
                        info!("   - Messages to keep: {}", filtered_lines.len());

                        // Save filtered messages back to file
                        info!("💾 Saving filtered messages to file...");
                        let new_content = filtered_lines.join("\n");
                        let new_content_with_newline = if !new_content.is_empty() {
                            format!("{}\n", new_content)
                        } else {
                            new_content
                        };

                        match fs::write(&self.messages_file, new_content_with_newline).await {
                            Ok(_) => {
                                info!("✅ Successfully saved {} messages (deleted {} messages)",
                                    filtered_lines.len(), deleted_message_count);
                            }
                            Err(e) => {
                                error!("❌ Failed to write filtered messages to file: {}", e);
                                error!("   File path: {}", self.messages_file.display());
                                return Err(anyhow::anyhow!("Failed to save filtered messages: {}", e));
                            }
                        }
                    }
                }
                Err(e) => {
                    error!("❌ Failed to read messages file: {}", e);
                    error!("   File path: {}", self.messages_file.display());
                    return Err(anyhow::anyhow!("Failed to read messages file: {}", e));
                }
            }
        }

        // Step 3: Clean up in-memory counter for this context
        info!("🧹 Step 3: Cleaning up in-memory message counter");
        {
            let mut counters = self.message_counters.write().await;
            if counters.remove(context_id).is_some() {
                info!("✅ Removed in-memory counter for context_id: {}", context_id);
            } else {
                info!("ℹ️  No in-memory counter found for context_id: {}", context_id);
            }
        }

        let total_time = start_time.elapsed();
        info!("🎉 Context deletion completed successfully in {:?}", total_time);
        info!("📊 Final summary:");
        info!("   - Context ID: {}", context_id);
        info!("   - User ID: {}", user_id);
        info!("   - Total documents deleted: {}", total_deleted);
        info!("   - Duration: {:?}", total_time);

        Ok(total_deleted)
    }

    // Private helper methods

    async fn get_next_message_index(&self, context_id: &str) -> u32 {
        let mut counters = self.message_counters.write().await;

        // Check if we already have a counter for this context_id
        if let Some(&current_index) = counters.get(context_id) {
            // Use the in-memory counter if it exists
            let next_index = current_index + 1;
            counters.insert(context_id.to_string(), next_index);
            return current_index;
        }

        // If no in-memory counter exists, check the local messages file for max index
        // This handles conversation continuation across app restarts
        let max_index_result = self.get_max_index_from_local_file(context_id).await;

        match max_index_result {
            Ok(0) => {
                // No existing messages found - start from 0
                counters.insert(context_id.to_string(), 1);
                0
            },
            Ok(max_index) => {
                // Existing messages found - continue from max_index + 1
                counters.insert(context_id.to_string(), max_index + 2);
                max_index + 1
            },
            Err(_) => {
                // Error reading file - start from 0
                counters.insert(context_id.to_string(), 1);
                0
            }
        }
    }

    /// Get the maximum message index from local messages.jsonl file
    async fn get_max_index_from_local_file(&self, context_id: &str) -> Result<u32> {
        use tokio::fs;

        if !self.messages_file.exists() {
            return Ok(0);
        }

        let content = fs::read_to_string(&self.messages_file).await?;
        let mut max_index = 0u32;

        for line in content.lines() {
            if line.trim().is_empty() {
                continue;
            }

            if let Ok(msg) = serde_json::from_str::<serde_json::Value>(line) {
                // Check if this message belongs to our context_id
                if let Some(msg_context) = msg.get("context_id").and_then(|v| v.as_str()) {
                    if msg_context == context_id {
                        // Get the message_index
                        if let Some(idx) = msg.get("message_index").and_then(|v| v.as_u64()) {
                            max_index = max_index.max(idx as u32);
                        }
                    }
                }
            }
        }

        Ok(max_index)
    }

    /// Get the maximum message index for a context_id from the database
    async fn get_max_message_index_from_db(&self, context_id: &str, user_id: &str) -> Result<u32> {
        // Use the search service to query all messages for this context_id
        let chat_options = services::search_service::ChatSearchOptions {
            context_id: Some(context_id.to_string()),
            role: None,
            from_timestamp: None,
            to_timestamp: None,
            from_message_index: None,
            to_message_index: None,
            limit: Some(1000), // Get enough messages to find the max
            include_metadata: false,
            user_id: Some(user_id.to_string()),
        };

        let results = self.search_service.search_chat_history(chat_options).await?;

        // Find the maximum message index from the results
        let mut max_index = 0u32;

        for result in results {
            if let Some(messages) = result.get("m").and_then(|v| v.as_array()) {
                // The result is a conversation object with messages array
                // The max index is the number of messages (since they're 0-indexed)
                max_index = max_index.max(messages.len() as u32);
            } else if let Some(index) = result.get("m").and_then(|v| v.as_i64()) {
                // Individual message with index field
                max_index = max_index.max(index as u32);
            } else if let Some(index) = result.get("message_index").and_then(|v| v.as_i64()) {
                // Alternative field name
                max_index = max_index.max(index as u32);
            }
        }

        Ok(max_index)
    }

    fn extract_content_from_streaming_response(&self, raw_response: &str) -> String {
        // Simple implementation - in practice, you'd want more sophisticated parsing
        raw_response
            .lines()
            .filter(|line| !line.trim().is_empty())
            .filter(|line| !line.starts_with("data:"))
            .collect::<Vec<_>>()
            .join("\n")
            .trim()
            .to_string()
    }

 

    async fn store_message(&self, message: &SessionMessage) -> Result<()> {
        let json_line = serde_json::to_string(message)? + "\n";
        let mut file = fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&self.messages_file)
            .await?;
        file.write_all(json_line.as_bytes()).await?;
        file.flush().await?;
        Ok(())
    }


    async fn load_messages_by_context(&self, context_id: &str) -> Result<Vec<SessionMessage>> {
        let all_messages = self.load_all_messages().await?;
        Ok(all_messages
            .into_iter()
            .filter(|m| m.context_id == context_id)
            .collect())
    }

    async fn load_all_messages(&self) -> Result<Vec<SessionMessage>> {
        self.load_jsonl_file(&self.messages_file).await
    }

    async fn load_jsonl_file<T: for<'de> Deserialize<'de>>(&self, file_path: &PathBuf) -> Result<Vec<T>> {
        if !file_path.exists() {
            return Ok(Vec::new());
        }

        let mut file = fs::File::open(file_path).await?;
        let mut contents = String::new();
        file.read_to_string(&mut contents).await?;

        let mut items = Vec::new();
        let mut failed_count = 0;
        for (line_num, line) in contents.lines().enumerate() {
            if !line.trim().is_empty() {
                match serde_json::from_str::<T>(line) {
                    Ok(item) => items.push(item),
                    Err(e) => {
                        failed_count += 1;
                        error!("Failed to parse line {}: {}, error: {}", line_num + 1, line, e);
                    }
                }
            }
        }

        if failed_count > 0 {
            error!("Total failed to parse: {} lines from {:?}", failed_count, file_path);
        }

        Ok(items)
    }

    /// Save a chat message as an encrypted document
    async fn save_encrypted_chat_document(&self, message: &SessionMessage, user_id: &str) -> Result<()> {
        // Create document content with optional request_id
        let mut document_content = serde_json::json!({
            "message_id": message.id,
            "i": message.context_id,
            "r": match message.role {
                MessageRole::User => "0",
                MessageRole::Assistant => "1"
            },
            "c": message.content,
            "t": message.timestamp,
            "m": message.message_index,
            "ct": message.chat_title
        });

        // Add request_id if present
        if let Some(ref ri) = message.request_id {
            document_content["ri"] = serde_json::Value::String(ri.clone());
        }

        let content_str = serde_json::to_string(&document_content)?;

        // Use simple dummy embedding for chat history (no need for real embeddings)
        let embedding = vec![0.1];

        // Encrypt content (base64 format)
        let encrypted_content = self.encryption_service.encrypt_content(&content_str).await?;

        // Create metadata
        let metadata = serde_json::json!({
            "i": message.context_id,
            "r": match message.role {
                MessageRole::User => "0",
                MessageRole::Assistant => "1"
            },
            "m": message.message_index,
            "t": message.timestamp,
            "ct": message.chat_title
        });

        // Encrypt metadata (base64 format)
        let encrypted_metadata_obj = self.encryption_service.encrypt_metadata(&metadata).await?;
        let encrypted_metadata_str = encrypted_metadata_obj
            .get("_encrypted_metadata")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("Failed to extract encrypted metadata"))?
            .to_string();

        // Create encrypted document
        let encrypted_doc = EncryptedChatDocument {
            id: Uuid::new_v4().to_string(),
            vector_id: Uuid::new_v4().to_string(),
            content: encrypted_content,
            embedding, // Use the generated embedding
            metadata: EncryptedDocumentMetadata {
                encrypted_metadata: encrypted_metadata_str,
                encrypted_content: true,
                created_at: message.timestamp.to_rfc3339(),
                updated_at: Utc::now().to_rfc3339(),
            },
        };

        // Get user-specific file path
        let user_documents_file = self.get_user_encrypted_documents_path(user_id);

        // Load existing documents, add new one, and save
        let mut existing_docs: Vec<EncryptedChatDocument> = if user_documents_file.exists() {
            let content = fs::read_to_string(&user_documents_file).await?;
            if !content.trim().is_empty() {
                // Try to parse as new wrapper format first
                if let Ok(wrapper) = serde_json::from_str::<serde_json::Value>(&content) {
                    if let Some(documents) = wrapper.get("documents") {
                        serde_json::from_value(documents.clone()).unwrap_or_default()
                    } else {
                        // Fallback: try to parse as old array format
                        serde_json::from_str(&content).unwrap_or_default()
                    }
                } else {
                    Vec::new()
                }
            } else {
                Vec::new()
            }
        } else {
            Vec::new()
        };

        existing_docs.push(encrypted_doc);

        // Ensure the user directory exists before saving
        let user_directory = self.get_user_directory_path(user_id);
        fs::create_dir_all(&user_directory).await?;

        // Create wrapper struct with specific field ordering
        #[derive(serde::Serialize)]
        struct DocumentsWrapper<'a> {
            documents: &'a Vec<EncryptedChatDocument>,
            count: usize,
            #[serde(rename = "lastModified")]
            last_modified: String,
        }

        let wrapper = DocumentsWrapper {
            documents: &existing_docs,
            count: existing_docs.len(),
            last_modified: Utc::now().to_rfc3339(),
        };

        // Save updated documents with wrapper
        let documents_content = serde_json::to_string_pretty(&wrapper)?;
        fs::write(&user_documents_file, documents_content).await?;

        Ok(())
    }

    /// Generic method to save encrypted documents for any collection with user-specific files
    async fn save_encrypted_document_to_collection(&self, collection_name: &str, content: &serde_json::Value, metadata: &serde_json::Value, user_id: &str, embedding: Option<Vec<f32>>) -> Result<String> {
        // Encrypt content (base64 format)
        let content_str = serde_json::to_string(content)?;
        let encrypted_content = self.encryption_service.encrypt_content(&content_str).await?;

        // Encrypt metadata (base64 format)
        let encrypted_metadata_obj = self.encryption_service.encrypt_metadata(metadata).await?;
        let encrypted_metadata_str = encrypted_metadata_obj
            .get("_encrypted_metadata")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("Failed to extract encrypted metadata"))?
            .to_string();

        // Use provided embedding or generate default based on collection type
        let final_embedding = if let Some(emb) = embedding {
            emb
        } else if collection_name == "chat_history" {
            vec![0.1] // Dummy embedding for chat history
        } else {
            // Generate real embedding for other collections
            self.embedding_service.generate_embedding(&content_str).await
                .unwrap_or_else(|_| {
                    // Fallback to random embedding
                    use rand::Rng;
                    let mut rng = rand::thread_rng();
                    (0..1024).map(|_| rng.gen::<f64>() as f32).collect()
                })
        };

        // Create encrypted document
        let document_id = Uuid::new_v4().to_string();
        let encrypted_doc = EncryptedChatDocument {
            id: document_id.clone(),
            vector_id: Uuid::new_v4().to_string(),
            content: encrypted_content,
            embedding: final_embedding,
            metadata: EncryptedDocumentMetadata {
                encrypted_metadata: encrypted_metadata_str,
                encrypted_content: true,
                created_at: Utc::now().to_rfc3339(),
                updated_at: Utc::now().to_rfc3339(),
            },
        };

        // Get user-specific file path for this collection
        let user_documents_file = self.get_user_collection_documents_path(user_id, collection_name);

        // Load existing documents, add new one, and save
        let mut existing_docs: Vec<EncryptedChatDocument> = if user_documents_file.exists() {
            let content = fs::read_to_string(&user_documents_file).await?;
            if !content.trim().is_empty() {
                // Try to parse as new wrapper format first
                if let Ok(wrapper) = serde_json::from_str::<serde_json::Value>(&content) {
                    if let Some(documents) = wrapper.get("documents") {
                        serde_json::from_value(documents.clone()).unwrap_or_default()
                    } else {
                        // Fallback: try to parse as old array format
                        serde_json::from_str(&content).unwrap_or_default()
                    }
                } else {
                    Vec::new()
                }
            } else {
                Vec::new()
            }
        } else {
            Vec::new()
        };

        // Convert encrypted document to Document type for vector store
        // NOTE: We don't manually save to JSON file here - the vector store handles that internally
        let mut metadata_map = indexmap::IndexMap::new();
        metadata_map.insert("_encrypted_metadata".to_string(), serde_json::Value::String(encrypted_doc.metadata.encrypted_metadata));
        metadata_map.insert("_encrypted_content".to_string(), serde_json::Value::Bool(true));
        metadata_map.insert("created_at".to_string(), serde_json::Value::String(encrypted_doc.metadata.created_at));
        metadata_map.insert("updated_at".to_string(), serde_json::Value::String(encrypted_doc.metadata.updated_at));

        let vector_doc = types::Document::new_with_vector_id(
            encrypted_doc.id,
            encrypted_doc.vector_id,
            encrypted_doc.content
        )
        .with_metadata(metadata_map)
        .with_embedding(encrypted_doc.embedding);

        // Set user context before adding to vector store
        if let Some(embedded_store) = self.vector_store.as_any().downcast_ref::<db::EmbeddedQdrantVectorStore>() {
            embedded_store.set_user_context(user_id).await;
        }
        else if let Some(server_store) = self.vector_store.as_any().downcast_ref::<db::QdrantServerVectorStore>() {
            server_store.set_user_context(user_id).await;
        }

        // Add to vector store - this will handle BOTH vector store insertion AND JSON file saving
        self.vector_store.add_document(collection_name, vector_doc).await?;

        Ok(document_id)
    }

    /// Get user-specific collection documents file path
    pub fn get_user_collection_documents_path(&self, user_id: &str, collection_name: &str) -> std::path::PathBuf {
        self.get_user_directory_path(user_id).join(format!("{}-documents.json", collection_name))
    }

    /// Generic method for any new collection to add documents - they will automatically follow the same pattern
    pub async fn add_document_to_collection(&self, collection_name: &str, content: &serde_json::Value, metadata: &serde_json::Value, user_id: &str) -> Result<String> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        // Use the generic encrypted document saving method
        let document_id = self.save_encrypted_document_to_collection(collection_name, content, metadata, user_id, None).await?;

        info!("Added document {} to collection {} for user {}", document_id, collection_name, user_id);

        Ok(document_id)
    }

    /// Get documents from any collection for a specific user (returns encrypted documents)
    pub async fn get_collection_documents(&self, collection_name: &str, user_id: &str) -> Result<Vec<EncryptedChatDocument>> {
        let user_documents_file = self.get_user_collection_documents_path(user_id, collection_name);

        if !user_documents_file.exists() {
            return Ok(Vec::new());
        }

        let content = fs::read_to_string(&user_documents_file).await?;
        if content.trim().is_empty() {
            return Ok(Vec::new());
        }

        // Try to parse as wrapper format first
        if let Ok(wrapper) = serde_json::from_str::<serde_json::Value>(&content) {
            if let Some(documents) = wrapper.get("documents") {
                let docs: Vec<EncryptedChatDocument> = serde_json::from_value(documents.clone())?;
                return Ok(docs);
            }
        }

        // Fallback: try to parse as old array format
        let docs: Vec<EncryptedChatDocument> = serde_json::from_str(&content).unwrap_or_default();
        Ok(docs)
    }

    /// Update existing chat history documents JSON files from Qdrant points with wrapper structure
    /// This method syncs documents for a specific user by organizing them into user-specific folders
    /// NOTE: This method is not used in normal operation - documents are saved directly during add_prompt/add_response
    #[allow(dead_code)]
    pub async fn sync_chat_history_from_qdrant(&self, user_id: &str) -> Result<()> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        // Get user-specific file path
        let user_documents_file = self.get_user_encrypted_documents_path(user_id);

        // Load existing documents from the user-specific JSON file
        let mut existing_docs: Vec<EncryptedChatDocument> = if user_documents_file.exists() {
            let content = fs::read_to_string(&user_documents_file).await?;
            if !content.trim().is_empty() {
                // Try to parse as wrapper format first
                if let Ok(wrapper) = serde_json::from_str::<serde_json::Value>(&content) {
                    if let Some(documents) = wrapper.get("documents") {
                        serde_json::from_value(documents.clone()).unwrap_or_default()
                    } else {
                        // Fallback: try to parse as old array format
                        serde_json::from_str(&content).unwrap_or_default()
                    }
                } else {
                    Vec::new()
                }
            } else {
                Vec::new()
            }
        } else {
            Vec::new()
        };

        // Get documents from the chat_history collection in Qdrant filtered by user_id
        let qdrant_documents = self.vector_store.list_documents("chat_history", None, None).await?;

        // Filter documents by user_id from encrypted metadata (include ALL contexts for this user)
        for doc in qdrant_documents {
            // Decrypt metadata to check if it belongs to our user
            let mut belongs_to_user = false;
            if let Some(encrypted_metadata) = doc.metadata.get("_encrypted_metadata").and_then(|v| v.as_str()) {
                if let Ok(decrypted_metadata) = self.encryption_service.decrypt_metadata(&serde_json::json!({"_encrypted_metadata": encrypted_metadata})).await {
                    // Check if this document belongs to our user (we include ALL contexts for this user)
                    if let Some(doc_session_id) = decrypted_metadata.get("sessionId").and_then(|v| v.as_str()) {
                        // For now, we'll include all documents since filtering by user would require
                        // more complex logic. In practice, you'd filter by user_id here.
                        belongs_to_user = true; // Include all documents for this sync
                    }
                }
            }

            if !belongs_to_user {
                continue; // Skip documents that don't belong to this user
            }

            let encrypted_metadata_str = doc.metadata.get("_encrypted_metadata")
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_string();

            let encrypted_content_flag = doc.metadata.get("_encrypted_content")
                .and_then(|v| v.as_bool())
                .unwrap_or(false);

            let created_at = doc.metadata.get("created_at")
                .and_then(|v| v.as_str())
                .unwrap_or(&doc.created_at.to_rfc3339())
                .to_string();

            let updated_at = doc.metadata.get("updated_at")
                .and_then(|v| v.as_str())
                .unwrap_or(&doc.updated_at.to_rfc3339())
                .to_string();

            let encrypted_doc = EncryptedChatDocument {
                id: doc.id.clone(),
                vector_id: doc.vector_id.clone(),
                content: doc.content.clone(),
                embedding: doc.embedding.unwrap_or_else(|| vec![0.1]),
                metadata: EncryptedDocumentMetadata {
                    encrypted_metadata: encrypted_metadata_str,
                    encrypted_content: encrypted_content_flag,
                    created_at,
                    updated_at,
                },
            };

            // Update existing document or add new one
            if let Some(existing_index) = existing_docs.iter().position(|d| d.id == doc.id) {
                existing_docs[existing_index] = encrypted_doc;
            } else {
                existing_docs.push(encrypted_doc);
            }
        }

        // Create wrapper struct with the requested structure
        #[derive(serde::Serialize)]
        struct ChatHistoryWrapper<'a> {
            documents: &'a Vec<EncryptedChatDocument>,
            count: usize,
            #[serde(rename = "lastModified")]
            last_modified: String,
        }

        let wrapper = ChatHistoryWrapper {
            documents: &existing_docs,
            count: existing_docs.len(),
            last_modified: Utc::now().to_rfc3339(),
        };

        // Ensure the user directory exists before saving
        let user_directory = self.get_user_directory_path(user_id);
        fs::create_dir_all(&user_directory).await?;

        // Save the updated wrapper structure to the user-specific JSON file
        let content = serde_json::to_string_pretty(&wrapper)?;
        fs::write(&user_documents_file, content).await?;

        info!("Updated chat history JSON file for user {} with {} documents from Qdrant", user_id, existing_docs.len());
        Ok(())
    }
    
    /// Save AWS estate data - store the string as one Qdrant point
    async fn save_aws_estate_documents(&self, data: &serde_json::Value, _document_ids: &[String], user_id: &str, collection_name: &str) -> Result<()> {
        // Extract content field from the JSON
        let content_str = if let Some(content) = data.get("content") {
            content.as_str()
                .ok_or_else(|| anyhow::anyhow!("content field must be a string"))?
                .to_string()
        } else {
            return Err(anyhow::anyhow!("content field is required in the JSON object"));
        };

        // Create metadata with all fields except "content" first
        let mut metadata_obj = data.as_object()
            .ok_or_else(|| anyhow::anyhow!("data must be a JSON object"))?
            .clone();
        metadata_obj.remove("content");

        // Generate embedding using existing embedding service
        let embedding = self.embedding_service.generate_embedding(&content_str).await?;

        // Generate deterministic document ID from ARN (enables deduplication)
        let doc_id = Self::generate_doc_id_from_arn(collection_name, &metadata_obj);

        let metadata_str = serde_json::to_string(&metadata_obj)?;

        // Create IndexMap for metadata
        let mut metadata = indexmap::IndexMap::new();
        metadata.insert("_encrypted_metadata".to_string(), serde_json::Value::String(metadata_str));

        // Create document and store as one point in Qdrant
        let document = types::Document::new(doc_id, content_str)
            .with_embedding(embedding)
            .with_metadata(metadata);

        self.vector_store.add_document(collection_name, document).await?;

        Ok(())
    }

    /// Load and return encrypted chat documents for a specific user (without decryption)
    pub async fn get_encrypted_chat_documents(&self, user_id: &str) -> Result<Vec<EncryptedChatDocument>> {
        let user_documents_file = self.get_user_encrypted_documents_path(user_id);

        if user_documents_file.exists() {
            let content = fs::read_to_string(&user_documents_file).await?;
            if !content.trim().is_empty() {
                // Try to parse as new wrapper format first
                if let Ok(wrapper) = serde_json::from_str::<serde_json::Value>(&content) {
                    if let Some(documents) = wrapper.get("documents") {
                        let docs: Vec<EncryptedChatDocument> = serde_json::from_value(documents.clone())?;
                        Ok(docs)
                    } else {
                        // Fallback: try to parse as old array format
                        let docs: Vec<EncryptedChatDocument> = serde_json::from_str(&content)?;
                        Ok(docs)
                    }
                } else {
                    Ok(Vec::new())
                }
            } else {
                Ok(Vec::new())
            }
        } else {
            Ok(Vec::new())
        }
    }

    /// Load and return encrypted AWS estate documents (without decryption)
    pub async fn get_encrypted_aws_estate_documents(&self) -> Result<Vec<EncryptedChatDocument>> {
        if self.aws_estate_file.exists() {
            let content = fs::read_to_string(&self.aws_estate_file).await?;
            if !content.trim().is_empty() {
                // Try to parse as new wrapper format first
                if let Ok(wrapper) = serde_json::from_str::<serde_json::Value>(&content) {
                    if let Some(documents) = wrapper.get("documents") {
                        let docs: Vec<EncryptedChatDocument> = serde_json::from_value(documents.clone())?;
                        Ok(docs)
                    } else {
                        // Fallback: try to parse as old array format
                        let docs: Vec<EncryptedChatDocument> = serde_json::from_str(&content)?;
                        Ok(docs)
                    }
                } else {
                    Ok(Vec::new())
                }
            } else {
                Ok(Vec::new())
            }
        } else {
            Ok(Vec::new())
        }
    }

    /// Get the path to user-specific encrypted documents file (deprecated - use context-specific version)
    pub fn get_user_encrypted_documents_path(&self, user_id: &str) -> PathBuf {
        self.qdrant_data_path
            .join(user_id)
            .join("chat_history-documents.json")
    }

    /// Get the user directory path (for creating directories - deprecated - use context-specific version)
    fn get_user_directory_path(&self, user_id: &str) -> PathBuf {
        self.qdrant_data_path.join(user_id)
    }

    /// Get decrypted chat history for a user organized by context_id with prompts and responses
    pub async fn get_decrypted_chat_history(&self, user_id: &str) -> Result<Vec<serde_json::Value>> {
        let encrypted_docs = match self.get_encrypted_chat_documents(user_id).await {
            Ok(docs) => docs,
            Err(e) => {
                // If we can't get documents (file doesn't exist, empty, etc.), return empty result
                println!("Warning: Could not load encrypted documents for user {}: {}", user_id, e);
                return Ok(Vec::new());
            }
        };

        if encrypted_docs.is_empty() {
            return Ok(Vec::new());
        }

        use indexmap::IndexMap;
        let mut contexts: IndexMap<String, serde_json::Value> = IndexMap::new();

        for doc in encrypted_docs {
            // Decrypt the content
            let decrypted_content = match self.encryption_service.decrypt_content(&doc.content).await {
                Ok(content) => content,
                Err(e) => {
                    println!("Warning: Failed to decrypt content for document {}: {}", doc.id, e);
                    continue;
                }
            };

            // Try to parse as JSON first, if that fails, treat as plain text
            let (message_content, content_json) = match serde_json::from_str::<serde_json::Value>(&decrypted_content) {
                Ok(json) => {
                    // Try to get the "c" field (short form) or "content" field
                    let content = json.get("c")
                        .or_else(|| json.get("content"))
                        .and_then(|v| v.as_str())
                        .map(|s| s.to_string())
                        .unwrap_or_else(|| {
                            // If neither field exists, this might not be a message object
                            // Return empty string or the raw content
                            String::new()
                        });
                    (content, Some(json))
                },
                Err(_) => {
                    // If it's not JSON, use the raw decrypted content as the message
                    (decrypted_content, None)
                }
            };

            // Decrypt the metadata
            let metadata_json = serde_json::json!({"_encrypted_metadata": doc.metadata.encrypted_metadata});

            let decrypted_metadata = match self.encryption_service.decrypt_metadata(&metadata_json).await {
                Ok(metadata) => metadata,
                Err(e) => {
                    println!("Warning: Failed to decrypt metadata for document {}: {}", doc.id, e);
                    continue;
                }
            };

            // Extract context_id - try content JSON first, then metadata
            let context_id = if let Some(ref json) = content_json {
                json.get("i")
                    .or_else(|| json.get("i"))
                    .and_then(|v| v.as_str())
                    .or_else(|| {
                        decrypted_metadata.get("i")
                            .or_else(|| decrypted_metadata.get("i"))
                            .and_then(|v| v.as_str())
                    })
                    .unwrap_or("unknown")
                    .to_string()
            } else {
                decrypted_metadata.get("i")
                    .or_else(|| decrypted_metadata.get("i"))
                    .and_then(|v| v.as_str())
                    .unwrap_or("unknown")
                    .to_string()
            };

            // Extract role - try content JSON first, then metadata
            let role = if let Some(ref json) = content_json {
                json.get("r")
                    .and_then(|v| v.as_str())
                    .or_else(|| {
                        decrypted_metadata.get("r")
                            .and_then(|v| v.as_str())
                    })
                    .unwrap_or("unknown")
            } else {
                decrypted_metadata.get("r")
                    .and_then(|v| v.as_str())
                    .unwrap_or("unknown")
            };


            // Extract timestamp - try content JSON first, then metadata
            let timestamp = if let Some(ref json) = content_json {
                json.get("t")
                    .and_then(|v| v.as_str())
                    .or_else(|| {
                        decrypted_metadata.get("t")
                            .or_else(|| decrypted_metadata.get("t"))
                            .or_else(|| decrypted_metadata.get("t"))
                            .and_then(|v| v.as_str())
                    })
                    .unwrap_or("")
                    .to_string()
            } else {
                decrypted_metadata.get("t")
                    .or_else(|| decrypted_metadata.get("t"))
                    .or_else(|| decrypted_metadata.get("t"))
                    .and_then(|v| v.as_str())
                    .unwrap_or("")
                    .to_string()
            };

            // Extract chat_title - try content JSON first, then metadata
            let chat_title = if let Some(ref json) = content_json {
                json.get("ct")
                    .or_else(|| json.get("ct"))
                    .and_then(|v| v.as_str())
                    .or_else(|| {
                        decrypted_metadata.get("ct")
                            .or_else(|| decrypted_metadata.get("ct"))
                            .and_then(|v| v.as_str())
                    })
                    .map(|s| s.to_string())
            } else {
                decrypted_metadata.get("ct")
                    .or_else(|| decrypted_metadata.get("ct"))
                    .and_then(|v| v.as_str())
                    .map(|s| s.to_string())
            };

            // Initialize context if it doesn't exist
            if !contexts.contains_key(&context_id) {
                let mut context_obj = IndexMap::new();
                context_obj.insert("i".to_string(), serde_json::Value::String(context_id.clone()));
                context_obj.insert("a".to_string(), serde_json::Value::String(user_id.to_string()));
                context_obj.insert("ct".to_string(), chat_title.clone().map(serde_json::Value::String).unwrap_or(serde_json::Value::Null));
                context_obj.insert("m".to_string(), serde_json::Value::Array(Vec::new()));
                context_obj.insert("t".to_string(), serde_json::Value::String("".to_string())); // Will be updated later with max index timestamp
                context_obj.insert("l".to_string(), serde_json::Value::Number(0.into()));
                contexts.insert(context_id.clone(), serde_json::Value::Object(context_obj.into_iter().collect()));
            } else if chat_title.is_some() {
                // Update chat_title if it's present in this message and not already set
                if let Some(context_entry) = contexts.get_mut(&context_id) {
                    if context_entry.get("ct").and_then(|v| v.as_str()).is_none() {
                        if let Some(obj) = context_entry.as_object_mut() {
                            obj.insert("ct".to_string(), serde_json::Value::String(chat_title.clone().unwrap()));
                        }
                    }
                }
            }

            // Extract message index from content JSON for proper sorting
            let message_index = if let Some(ref json) = content_json {
                json.get("m")
                    .or_else(|| json.get("message_index"))
                    .and_then(|v| v.as_u64())
                    .unwrap_or(0)
            } else {
                0
            };

            // Extract request_id (ri) from content JSON
            let request_id = if let Some(ref json) = content_json {
                json.get("ri")
                    .and_then(|v| v.as_str())
                    .map(|s| s.to_string())
            } else {
                None
            };

            // Create message object with role, content, timestamp, and index for sorting
            let mut message_obj = IndexMap::new();
            message_obj.insert("r".to_string(), serde_json::Value::String(role.to_string()));
            message_obj.insert("c".to_string(), serde_json::Value::String(message_content));

            // Add request_id if present
            if let Some(ri) = request_id {
                message_obj.insert("ri".to_string(), serde_json::Value::String(ri));
            }

            message_obj.insert("_t".to_string(), serde_json::Value::String(timestamp.clone())); // Temporary field for sorting
            message_obj.insert("_i".to_string(), serde_json::Value::Number(message_index.into())); // Temporary field for sorting

            // Add to messages array
            let context_entry = contexts.get_mut(&context_id).unwrap();
            if let Some(messages) = context_entry.get_mut("m").and_then(|m| m.as_array_mut()) {
                messages.push(serde_json::Value::Object(message_obj.into_iter().collect()));
            }
        }

        // Sort messages within each context and update timestamp/count
        for (_context_id, context_data) in contexts.iter_mut() {
            if let Some(messages) = context_data.get_mut("m").and_then(|m| m.as_array_mut()) {
                // Sort messages by index and timestamp for proper ordering
                messages.sort_by(|a, b| {
                    let a_idx = a.get("_i").and_then(|i| i.as_u64()).unwrap_or(0);
                    let b_idx = b.get("_i").and_then(|i| i.as_u64()).unwrap_or(0);
                    match a_idx.cmp(&b_idx) {
                        std::cmp::Ordering::Equal => {
                            let a_time = a.get("_t").and_then(|t| t.as_str()).unwrap_or("");
                            let b_time = b.get("_t").and_then(|t| t.as_str()).unwrap_or("");
                            a_time.cmp(b_time)
                        },
                        other => other
                    }
                });

                // Find the message with the highest index to get its timestamp
                let latest_timestamp = messages.iter()
                    .max_by_key(|msg| msg.get("_i").and_then(|i| i.as_u64()).unwrap_or(0))
                    .and_then(|msg| msg.get("_t").and_then(|t| t.as_str()))
                    .map(|s| s.to_string());

                // Remove temporary fields (_t and _i) from messages
                let cleaned_messages: Vec<serde_json::Value> = messages.iter().map(|message| {
                    let mut clean_msg = IndexMap::new();
                    clean_msg.insert("r".to_string(), message.get("r").cloned().unwrap_or(serde_json::Value::Null));
                    clean_msg.insert("c".to_string(), message.get("c").cloned().unwrap_or(serde_json::Value::Null));
                    // Preserve request_id (ri) field if present
                    if let Some(ri) = message.get("ri") {
                        clean_msg.insert("ri".to_string(), ri.clone());
                    }
                    serde_json::Value::Object(clean_msg.into_iter().collect())
                }).collect();

                // Replace messages array with cleaned format
                *messages = cleaned_messages;

                // Update message count and latest timestamp
                let message_count = messages.len();
                if let Some(obj) = context_data.as_object_mut() {
                    obj.insert("l".to_string(), serde_json::Value::Number(message_count.into()));
                    // Update timestamp to the latest message's timestamp
                    if let Some(ts) = latest_timestamp {
                        obj.insert("t".to_string(), serde_json::Value::String(ts));
                    }
                }
            }
        }

        // Convert to final format - return as array of contexts
        let result: Vec<serde_json::Value> = contexts.into_iter()
            .map(|(_key, value)| value)
            .collect();

        Ok(result)
    }

    /// Get the path to the AWS estate file
    pub fn get_aws_estate_path(&self) -> &std::path::Path {
        &self.aws_estate_file
    }

    /// Clear session storage (for testing)
    /// Clear message storage (for testing)
    pub async fn clear_message_storage(&self) -> Result<()> {
        if self.messages_file.exists() {
            fs::remove_file(&self.messages_file).await?;
        }

        {
            let mut counters = self.message_counters.write().await;
            counters.clear();
        }

        info!("Cleared all session storage files");
        Ok(())
    }

    /// Shutdown the RAG module
    pub async fn shutdown(&self) -> Result<()> {
        info!("Shutting down RAG Module");
        
        // Shutdown services in reverse order
        // self.sync_service.shutdown().await?;
        self.vector_store.shutdown().await?;
        
        let mut initialized = self.initialized.write().await;
        *initialized = false;
        
        info!("RAG Module shutdown complete");
        Ok(())
    }

    /// Create documents in bulk
    pub async fn create(&self, documents: Vec<Document>, collection_name: Option<&str>) -> Result<CreateResult> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        let mut created = 0;
        let mut failed = Vec::new();
        let default_collection = collection_name.unwrap_or("core_estate");

        for document in documents {
            // Determine collection type from document metadata or use provided/default collection
            let collection_type = document.metadata.get("collection_type")
                .and_then(|v| v.as_str())
                .unwrap_or(default_collection)
                .to_string();

            match self.document_service.add_document(&collection_type, document).await {
                Ok(_) => created += 1,
                Err(e) => {
                    failed.push(format!("Failed to create document: {}", e));
                }
            }
        }

        Ok(CreateResult { created, failed })
    }

    /// Get document count with optional filter
    pub async fn get_document_count(&self, collection_name: Option<&str>, filter: Option<serde_json::Value>) -> Result<usize> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        let collection = collection_name.unwrap_or("core_estate");
        // Use estate collection, convert JSON filter to SearchFilter if provided
        let search_filter = None; // TODO: Convert JSON filter to SearchFilter
        self.document_service.get_document_count(collection, search_filter).await
    }

    /// List documents with options
    pub async fn list_documents(&self, collection_name: Option<&str>, options: Option<serde_json::Value>) -> Result<Vec<Document>> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        let collection = collection_name.unwrap_or("core_estate");
        // Extract limit from options
        let limit = options.as_ref()
            .and_then(|o| o.get("limit"))
            .and_then(|l| l.as_u64())
            .map(|l| l as usize);

        let search_filter = None; // TODO: Convert JSON options to SearchFilter
        self.document_service.list_documents(collection, limit, search_filter).await
    }

    /// Set user context for document operations
    pub async fn set_user_context(&self, user_id: &str) -> Result<()> {
        // Set user context on the vector store
        if let Some(embedded_store) = self.vector_store.as_any().downcast_ref::<db::EmbeddedQdrantVectorStore>() {
            embedded_store.set_user_context(user_id).await;
        } else if let Some(server_store) = self.vector_store.as_any().downcast_ref::<db::QdrantServerVectorStore>() {
            server_store.set_user_context(user_id).await;
        } else if let Some(dual_store) = self.vector_store.as_any().downcast_ref::<db::DualVectorStore>() {
            dual_store.set_user_context(user_id).await;
        }

        // Also set user context on document service if needed
        // (currently document service doesn't have user context, but this is for future extensibility)

        Ok(())
    }

    /// Delete a collection and remove all associated files
    /// 
    /// This method will:
    /// 1. Delete the collection from the vector store
    /// 2. Remove all four generated files:
    ///    - {collection}-documents.json
    ///    - {collection}-vectors.bin  
    ///    - {collection}-vector-index.json
    ///    - {collection}-metadata.json
    /// 
    /// # Arguments
    /// * `collection_name` - The name of the collection to delete
    /// * `user_id` - The user ID (for user-specific data paths)
    /// 
    /// # Returns
    /// * `Result<CollectionDeleteResult>` - Summary of deletion operation
    /// 
    /// # Example
    /// ```rust
    /// let result = rag.delete_collection("aws_estate", "user123").await?;
    /// println!("Deleted collection with {} files removed", result.files_removed);
    /// ```
    pub async fn delete_collection(&self, collection_name: &str, user_id: &str) -> Result<CollectionDeleteResult> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        info!("🗑️ Deleting collection '{}' for user '{}'", collection_name, user_id);

        let mut result = CollectionDeleteResult {
            collection_name: collection_name.to_string(),
            user_id: user_id.to_string(),
            collection_deleted: false,
            files_removed: 0,
            removed_files: Vec::new(),
            errors: Vec::new(),
        };

        // Set user context
        self.set_user_context(user_id).await?;

        // Step 1: Delete the collection from vector store
        match self.vector_store.delete_collection(collection_name).await {
            Ok(_) => {
                result.collection_deleted = true;
                info!("✅ Collection '{}' deleted from vector store", collection_name);
            }
            Err(e) => {
                let error_msg = format!("Failed to delete collection from vector store: {}", e);
                error!("{}", error_msg);
                result.errors.push(error_msg);
            }
        }

        // Step 2: Remove the four generated files
        let user_data_path = self.base_path.join("qdrant-data").join(user_id);
        
        let files_to_remove = vec![
            format!("{}-documents.json", collection_name),
            format!("{}-vectors.bin", collection_name),
            format!("{}-vector-index.json", collection_name),
            format!("{}-metadata.json", collection_name),
        ];

        for file_name in files_to_remove {
            let file_path = user_data_path.join(&file_name);
            
            if file_path.exists() {
                match fs::remove_file(&file_path).await {
                    Ok(_) => {
                        result.files_removed += 1;
                        result.removed_files.push(file_name.clone());
                        info!("✅ Removed file: {}", file_name);
                    }
                    Err(e) => {
                        let error_msg = format!("Failed to remove file '{}': {}", file_name, e);
                        error!("{}", error_msg);
                        result.errors.push(error_msg);
                    }
                }
            } else {
                debug!("File '{}' does not exist, skipping", file_name);
            }
        }

        // Step 3: Try to remove user directory if it's empty
        if let Err(e) = fs::remove_dir(&user_data_path).await {
            // This is expected to fail if directory is not empty, so we don't treat it as an error
            debug!("Could not remove user directory (likely not empty): {}", e);
        }

        if result.errors.is_empty() {
            info!("🎉 Successfully deleted collection '{}': {} files removed", 
                  collection_name, result.files_removed);
        } else {
            warn!("⚠️ Collection '{}' deletion completed with {} errors", 
                  collection_name, result.errors.len());
        }

        Ok(result)
    }
}

// Implement Send + Sync for thread safety
unsafe impl Send for RagModule {}
unsafe impl Sync for RagModule {}

/// Factory function to create a RAG module
pub async fn create_rag_module(base_path: impl Into<PathBuf>) -> Result<RagModule> {
    RagModule::new(base_path).await
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;
    
    #[tokio::test]
    async fn test_rag_module_creation() {
        let temp_dir = TempDir::new().unwrap();
        let rag = RagModule::new(temp_dir.path()).await.unwrap();
        assert!(!rag.is_initialized().await);
    }
    
    #[tokio::test]
    async fn test_rag_module_initialization() {
        let temp_dir = TempDir::new().unwrap();
        let rag = RagModule::new(temp_dir.path()).await.unwrap();
        rag.initialize().await.unwrap();
        assert!(rag.is_initialized().await);
    }

    #[tokio::test]
    async fn test_session_management() {
        let temp_dir = TempDir::new().unwrap();
        let rag = RagModule::new(temp_dir.path()).await.unwrap();
        rag.initialize().await.unwrap();

        let context_id = "test_context";
        let user_id = "test_user";
        let chat_title = Some("Test Chat");

        // Add prompt and response
        let prompt_id = rag.add_prompt(context_id, "What is the weather?", user_id, chat_title).await.unwrap();
        let response_id = rag.add_response(context_id, "It's sunny today!", user_id, chat_title).await.unwrap();

        assert!(!prompt_id.is_empty());
        assert!(!response_id.is_empty());

        // Get chat history
        let history = rag.get_session_chat_history(context_id).await.unwrap();
        assert_eq!(history.total_pairs, 1);
        assert!(!history.is_empty);
        assert_eq!(history.pairs[0].query.content, "What is the weather?");
        assert_eq!(history.pairs[0].response.content, "It's sunny today!");
    }

    #[tokio::test]
    async fn test_query_response_pairs() {
        let temp_dir = TempDir::new().unwrap();
        let rag = RagModule::new(temp_dir.path()).await.unwrap();
        rag.initialize().await.unwrap();

        let context_id = "test_context_2";
        let user_id = "test_user";
        let chat_title = Some("Test Chat");

        rag.add_prompt(context_id, "Hello", user_id, chat_title).await.unwrap();
        rag.add_response(context_id, "Hi there!", user_id, chat_title).await.unwrap();

        rag.add_prompt(context_id, "How are you?", user_id, chat_title).await.unwrap();
        rag.add_response(context_id, "I'm doing well!", user_id, chat_title).await.unwrap();

        // Test getting limited pairs
        let pairs = rag.get_query_response_pairs(context_id, Some(1)).await.unwrap();
        assert_eq!(pairs.total_pairs, 2);
        assert_eq!(pairs.pairs.len(), 1); // Only the most recent pair
        assert_eq!(pairs.pairs[0].query.content, "How are you?");
        assert_eq!(pairs.pairs[0].response.content, "I'm doing well!");

        // Test getting all pairs
        let all_pairs = rag.get_query_response_pairs(context_id, None).await.unwrap();
        assert_eq!(all_pairs.total_pairs, 2);
        assert_eq!(all_pairs.pairs.len(), 2);
    }

    #[tokio::test]
    async fn test_batch_aws_estate_ingestion() {
        let temp_dir = TempDir::new().unwrap();
        let rag = RagModule::new(temp_dir.path()).await.unwrap();
        rag.initialize().await.unwrap();

        let user_id = "test_user_batch";
        let collection_name = "test_aws_estate";

        // Create test batch data (5 documents)
        let batch_data = vec![
            serde_json::json!({
                "content": "EC2 instance i-1234567890abcdef0 running in us-west-2",
                "resource_type": "ec2_instance",
                "instance_id": "i-1234567890abcdef0",
                "region": "us-west-2",
                "state": "running"
            }),
            serde_json::json!({
                "content": "S3 bucket my-test-bucket with 1000 objects",
                "resource_type": "s3_bucket",
                "bucket_name": "my-test-bucket",
                "region": "us-east-1",
                "object_count": 1000
            }),
            serde_json::json!({
                "content": "RDS database mydb-prod running MySQL 8.0",
                "resource_type": "rds_instance",
                "db_identifier": "mydb-prod",
                "engine": "mysql",
                "version": "8.0"
            }),
            serde_json::json!({
                "content": "Lambda function process-orders with Python 3.9 runtime",
                "resource_type": "lambda_function",
                "function_name": "process-orders",
                "runtime": "python3.9",
                "timeout": 300
            }),
            serde_json::json!({
                "content": "VPC vpc-12345678 with 3 subnets in us-west-2",
                "resource_type": "vpc",
                "vpc_id": "vpc-12345678",
                "subnet_count": 3,
                "region": "us-west-2"
            })
        ];

        // Test batch ingestion
        let result = rag.ingest_aws_estate_batch(batch_data.clone(), user_id, collection_name).await.unwrap();

        // Verify results
        assert_eq!(result.total_resources, 5);
        assert_eq!(result.parsed_resources, 5);
        assert_eq!(result.failed_resources, 0);
        assert_eq!(result.create_result.created, 5);
        assert!(result.create_result.failed.is_empty());

        println!("✅ Batch ingestion test completed: {} documents processed", result.parsed_resources);
    }

    #[tokio::test]
    async fn test_batch_ingestion_with_invalid_data() {
        let temp_dir = TempDir::new().unwrap();
        let rag = RagModule::new(temp_dir.path()).await.unwrap();
        rag.initialize().await.unwrap();

        let user_id = "test_user_invalid";
        let collection_name = "test_invalid_estate";

        // Create test data with some invalid documents (missing content field)
        let batch_data = vec![
            serde_json::json!({
                "content": "Valid EC2 instance",
                "resource_type": "ec2_instance",
                "instance_id": "i-valid123"
            }),
            serde_json::json!({
                // Missing content field - should fail
                "resource_type": "s3_bucket",
                "bucket_name": "invalid-bucket"
            }),
            serde_json::json!({
                "content": "Valid RDS instance",
                "resource_type": "rds_instance",
                "db_identifier": "valid-db"
            }),
            serde_json::json!({
                "content": 12345, // Invalid content type - should fail
                "resource_type": "lambda_function"
            })
        ];

        let result = rag.ingest_aws_estate_batch(batch_data, user_id, collection_name).await.unwrap();

        // Should have processed 2 valid documents and failed on 2 invalid ones
        assert_eq!(result.total_resources, 4);
        assert_eq!(result.parsed_resources, 2);
        assert_eq!(result.failed_resources, 2);
        assert_eq!(result.create_result.created, 2);
        assert_eq!(result.create_result.failed.len(), 2);

        println!("✅ Invalid data test completed: {} valid, {} failed", result.parsed_resources, result.failed_resources);
    }

    #[tokio::test]
    async fn test_empty_batch_ingestion() {
        let temp_dir = TempDir::new().unwrap();
        let rag = RagModule::new(temp_dir.path()).await.unwrap();
        rag.initialize().await.unwrap();

        let user_id = "test_user_empty";
        let collection_name = "test_empty_estate";

        // Test with empty batch
        let batch_data: Vec<serde_json::Value> = vec![];
        let result = rag.ingest_aws_estate_batch(batch_data, user_id, collection_name).await.unwrap();

        assert_eq!(result.total_resources, 0);
        assert_eq!(result.parsed_resources, 0);
        assert_eq!(result.failed_resources, 0);
        assert_eq!(result.create_result.created, 0);
        assert!(result.create_result.failed.is_empty());

        println!("✅ Empty batch test completed");
    }

    #[tokio::test]
    async fn test_large_batch_ingestion() {
        let temp_dir = TempDir::new().unwrap();
        let rag = RagModule::new(temp_dir.path()).await.unwrap();
        rag.initialize().await.unwrap();

        let user_id = "test_user_large";
        let collection_name = "test_large_estate";

        // Create a larger batch (32 documents as requested)
        let mut batch_data = Vec::new();
        for i in 0..32 {
            batch_data.push(serde_json::json!({
                "content": format!("AWS resource {} - EC2 instance with automated deployment", i),
                "resource_type": "ec2_instance",
                "instance_id": format!("i-{:016x}", i),
                "region": if i % 2 == 0 { "us-west-2" } else { "us-east-1" },
                "state": "running",
                "tags": {
                    "Environment": if i % 3 == 0 { "prod" } else { "dev" },
                    "Application": format!("app-{}", i % 5)
                }
            }));
        }

        let start_time = std::time::Instant::now();
        let result = rag.ingest_aws_estate_batch(batch_data, user_id, collection_name).await.unwrap();
        let duration = start_time.elapsed();

        // Verify all documents were processed successfully
        assert_eq!(result.total_resources, 32);
        assert_eq!(result.parsed_resources, 32);
        assert_eq!(result.failed_resources, 0);
        assert_eq!(result.create_result.created, 32);
        assert!(result.create_result.failed.is_empty());

        println!("✅ Large batch test completed: {} documents in {:?}", result.parsed_resources, duration);
    }

    #[tokio::test]
    async fn test_single_vs_batch_comparison() {
        let temp_dir = TempDir::new().unwrap();
        let rag = RagModule::new(temp_dir.path()).await.unwrap();
        rag.initialize().await.unwrap();

        let user_id = "test_user_comparison";

        // Create test data
        let test_documents = vec![
            serde_json::json!({
                "content": "Test EC2 instance for comparison",
                "resource_type": "ec2_instance",
                "instance_id": "i-comparison1"
            }),
            serde_json::json!({
                "content": "Test S3 bucket for comparison",
                "resource_type": "s3_bucket",
                "bucket_name": "comparison-bucket"
            }),
            serde_json::json!({
                "content": "Test RDS instance for comparison",
                "resource_type": "rds_instance",
                "db_identifier": "comparison-db"
            })
        ];

        // Test single ingestion (for comparison)
        let start_single = std::time::Instant::now();
        for (i, doc) in test_documents.iter().enumerate() {
            let collection_name = format!("test_single_{}", i);
            let _result = rag.ingest_aws_estate(doc.clone(), user_id, &collection_name).await.unwrap();
        }
        let single_duration = start_single.elapsed();

        // Test batch ingestion
        let start_batch = std::time::Instant::now();
        let batch_result = rag.ingest_aws_estate_batch(test_documents.clone(), user_id, "test_batch").await.unwrap();
        let batch_duration = start_batch.elapsed();

        // Verify batch worked correctly
        assert_eq!(batch_result.parsed_resources, 3);
        assert_eq!(batch_result.failed_resources, 0);

        println!("✅ Performance comparison:");
        println!("   Single ingestion (3 docs): {:?}", single_duration);
        println!("   Batch ingestion (3 docs): {:?}", batch_duration);
        println!("   Batch is {}x faster", single_duration.as_millis() as f64 / batch_duration.as_millis() as f64);
    }
}