//! 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;

// 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>,
}

#[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();

        // Check for QDRANT_URL environment variable to override config
        if let Ok(qdrant_url) = std::env::var("QDRANT_URL") {
            info!("🔧 QDRANT_URL environment variable detected: {}", qdrant_url);
            config.vector_store.backend = "qdrant-server".to_string();
            config.vector_store.connection.url = qdrant_url.clone();

            // Save updated config to file so users can see the change
            config_manager.update_config(config.clone()).await?;
            info!("✅ Config file updated: backend set to 'qdrant-server'");
        }

        // 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" => {
                info!("💾 Using Qdrant embedded mode");
                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?
        );
        
        // 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);
        
        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);
        
        // 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
    }

    /// 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);
            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,
            };
            let results = self.search_service.search_estate_resources(collection_type, query, estate_options, None, user_id).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 {
            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)
    }
    
    /// 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> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        self.set_user_context(user_id).await?;

        // Extract content field from the JSON
        let content_str = if let Some(content) = estate_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"));
        };

        // Generate embedding from the content
        let embedding = self.embedding_service.generate_embedding(&content_str).await?;

        // Generate a unique ID for the document
        let doc_id = format!("{}-{}", collection_name, Uuid::new_v4());

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

        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 with content and metadata
        let document = types::Document::new(doc_id, content_str)
            .with_embedding(embedding)
            .with_metadata(metadata);

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

        // 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(),
            },
        };

        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>) -> 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,
        };

        // 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
        });

        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 (SERVER MODE ONLY - REQUIRED)
        if let Some(server_store) = self.vector_store.as_any().downcast_ref::<db::QdrantServerVectorStore>() {
            server_store.set_user_context(user_id).await;
        } else {
            return Err(anyhow::anyhow!("Server mode required - only QdrantServerVectorStore is supported. Set QDRANT_URL environment variable."));
        }

        // 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>) -> 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,
        };

        // 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
        });

        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 (SERVER MODE ONLY - REQUIRED)
        if let Some(server_store) = self.vector_store.as_any().downcast_ref::<db::QdrantServerVectorStore>() {
            server_store.set_user_context(user_id).await;
        } else {
            return Err(anyhow::anyhow!("Server mode required - only QdrantServerVectorStore is supported. Set QDRANT_URL environment variable."));
        }

        // 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) -> Result<Vec<ContextInfo>> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

        // Load all messages
        let all_messages = self.load_all_messages().await?;

        // Group by context_id to get unique contexts
        let mut context_map: HashMap<String, ContextInfo> = HashMap::new();

        for message in all_messages {
            context_map
                .entry(message.context_id.clone())
                .and_modify(|existing| {
                    // Update with most recent timestamp
                    if message.timestamp > existing.updated_at {
                        existing.updated_at = message.timestamp;
                    }
                    // Update chat_title if it's not "Chat" and current one is
                    if existing.chat_title == "Chat" && message.chat_title != "Chat" {
                        existing.chat_title = message.chat_title.clone();
                    }
                })
                .or_insert(ContextInfo {
                    context_id: message.context_id.clone(),
                    chat_title: message.chat_title.clone(),
                    created_at: message.timestamp,
                    updated_at: message.timestamp,
                });
        }

        // Convert to Vec and sort by updated_at (most recent first)
        let mut contexts: Vec<ContextInfo> = context_map.into_values().collect();
        contexts.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));

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

    /// Get prompts and responses as JSON for a specific context_id
    ///
    /// Returns all conversation turns (prompt-response pairs) for the given context_id
    /// in a structured JSON format.
    ///
    /// # Arguments
    /// * `context_id` - The context ID to retrieve chat history for
    ///
    /// # Returns
    /// * `Vec<ConversationTurn>` - Array of conversation turns with prompts and responses
    pub async fn get_chat_history_json(&self, context_id: &str) -> Result<Vec<ConversationTurn>> {
        if !self.is_initialized().await {
            self.initialize().await?;
        }

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

        // Load messages for this context
        let messages = self.load_messages_by_context(context_id).await?;

        // Separate prompts and responses
        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
        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 conversation turns
        let pair_count = std::cmp::min(prompts.len(), responses.len());
        let mut conversation = Vec::new();

        for i in 0..pair_count {
            conversation.push(ConversationTurn {
                turn_number: i + 1,
                prompt: MessageJson {
                    message_id: prompts[i].id.clone(),
                    content: prompts[i].content.clone(),
                    timestamp: prompts[i].timestamp,
                    message_index: prompts[i].message_index,
                },
                response: MessageJson {
                    message_id: responses[i].id.clone(),
                    content: responses[i].content.clone(),
                    timestamp: responses[i].timestamp,
                    message_index: responses[i].message_index,
                },
            });
        }

        info!("Retrieved chat history for context: {} with {} conversation turns", context_id, conversation.len());

        Ok(conversation)
    }

    // 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();
        for line in contents.lines() {
            if !line.trim().is_empty() {
                match serde_json::from_str::<T>(line) {
                    Ok(item) => items.push(item),
                    Err(e) => error!("Failed to parse line: {}, error: {}", line, e),
                }
            }
        }

        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
        let 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        });

        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"));
        };

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

        // Generate a unique ID for the document
        let doc_id = format!("{}-{}", collection_name, Uuid::new_v4());

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

        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
            };

            // 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));
            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));
                    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 (SERVER MODE ONLY)
        // COMMENTED OUT: Embedded mode
        if let Some(embedded_store) = self.vector_store.as_any().downcast_ref::<db::EmbeddedQdrantVectorStore>() {
            embedded_store.set_user_context(user_id).await;
        }
        if let Some(server_store) = self.vector_store.as_any().downcast_ref::<db::QdrantServerVectorStore>() {
            server_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(())
    }
}

// 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();

      
        assert_eq!(session.user_id, "test_user");
        assert_eq!(session.chat_title, "Test Chat");
        assert!(matches!(session.status, SessionStatus::Active));

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

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

        // Get chat history
        let history = rag.get_session_chat_history(&session.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!");

        // End session
        let ended_session = rag.end_session(&session.id).await.unwrap();
        assert!(matches!(ended_session.status, SessionStatus::Ended));
        assert!(ended_session.end_time.is_some());
    }

    #[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);
    }
}