rag-module 0.6.7

use anyhow::Result;
use rag_module::{RagModule, StartSessionOptions};
use serde_json::json;

#[tokio::main]
async fn main() -> Result<()> {
    println!("🔄 Testing Unified Collection Pattern - All collections follow same user-specific encrypted saving pattern");

    // Initialize RAG module
    let mut rag_module = RagModule::new("./example-enhanced-chat-data").await?;
    rag_module.initialize().await?;

    let user_id = "cudewvsbxakj";

    // Test 1: Chat History Collection (already implemented)
    println!("\n📝 Test 1: Chat History Collection");
    let session_options = StartSessionOptions {
        user_id: user_id.to_string(),
        chat_title: Some("Unified Test".to_string()),
        context_id: None,
    };

    let session = rag_module.start_session(session_options).await?;
    rag_module.add_prompt(&session.id, "What is unified collection pattern?", user_id).await?;
    rag_module.add_response(&session.id, "It ensures all collections save data under user-specific encrypted files.", user_id).await?;

    let chat_docs = rag_module.get_collection_documents("chat_history", user_id).await?;
    println!("✅ Chat History: {} documents saved to ./test-unified-collections-data/qdrant-data/{}/chat_history-documents.json", chat_docs.len(), user_id);

    // Test 2: AWS Estate Collection (now uses unified pattern)
    println!("\n☁️ Test 2: AWS Estate Collection");
    let aws_data = json!([
        {
            "account_id": "123456789",
            "account_name": "Test Account",
            "services": {
                "s3": {
                    "service_metadata": {
                        "category": "storage"
                    },
                    "resources": ["bucket1", "bucket2"]
                }
            }
        }
    ]);

    rag_module.process_aws_estate(aws_data, user_id).await?;
    let aws_docs = rag_module.get_collection_documents("aws_estate", user_id).await?;
    println!("✅ AWS Estate: {} documents saved to ./test-unified-collections-data/qdrant-data/{}/aws_estate-documents.json", aws_docs.len(), user_id);

    // Test 3: Custom New Collection (demonstrates the pattern for future collections)
    println!("\n🆕 Test 3: Custom New Collection - 'user_preferences'");

    let preference_content = json!({
        "user_id": user_id,
        "preferences": {
            "theme": "dark",
            "notifications": true,
            "language": "en"
        },
        "updated_at": "2024-01-15T10:30:00Z"
    });

    let preference_metadata = json!({
        "collection_type": "user_preferences",
        "user_id": user_id,
        "category": "settings",
        "last_updated": "2024-01-15T10:30:00Z"
    });

    let pref_doc_id = rag_module.add_document_to_collection("user_preferences", &preference_content, &preference_metadata, user_id).await?;
    println!("✅ Added preference document: {}", pref_doc_id);

    let pref_docs = rag_module.get_collection_documents("user_preferences", user_id).await?;
    println!("✅ User Preferences: {} documents saved to ./test-unified-collections-data/qdrant-data/{}/user_preferences-documents.json", pref_docs.len(), user_id);

    // Test 4: Another Custom Collection - 'analytics_data'
    println!("\n📊 Test 4: Another Custom Collection - 'analytics_data'");

    let analytics_content = json!({
        "user_id": user_id,
        "session_data": {
            "page_views": 25,
            "time_spent": 1200,
            "actions": ["login", "search", "view_profile"]
        },
        "timestamp": "2024-01-15T11:45:00Z"
    });

    let analytics_metadata = json!({
        "collection_type": "analytics_data",
        "user_id": user_id,
        "category": "user_behavior",
        "session_id": session.id
    });

    let analytics_doc_id = rag_module.add_document_to_collection("analytics_data", &analytics_content, &analytics_metadata, user_id).await?;
    println!("✅ Added analytics document: {}", analytics_doc_id);

    let analytics_docs = rag_module.get_collection_documents("analytics_data", user_id).await?;
    println!("✅ Analytics Data: {} documents saved to ./test-unified-collections-data/qdrant-data/{}/analytics_data-documents.json", analytics_docs.len(), user_id);

    // Show the unified file structure
    println!("\n📁 Unified File Structure Created:");
    println!("./test-unified-collections-data/");
    println!("└── qdrant-data/");
    println!("    └── {}/", user_id);
    println!("        ├── chat_history-documents.json       ← Chat messages");
    println!("        ├── user_preferences-documents.json   ← User settings");
    println!("        ├── analytics_data-documents.json     ← User analytics");
    println!("        └── aws_estate-documents.json         ← AWS infrastructure");

    println!("\n🎉 All Collections Now Follow the Same Pattern:");
    println!("✅ User-specific directories: ./qdrant-data/{{user_id}}/");
    println!("✅ Consistent file naming: {{collection_name}}-documents.json");
    println!("✅ JSON wrapper format: {{documents, count, lastModified}}");
    println!("✅ Content + metadata encryption for security");
    println!("✅ Automatic embedding generation (dummy for chat, real for others)");

    println!("\n💡 For Future Collections:");
    println!("Just call: rag_module.add_document_to_collection(collection_name, content, metadata, user_id)");
    println!("And it automatically follows the same encrypted, user-specific pattern!");

    rag_module.end_session(&session.id).await?;

    Ok(())
}