use anyhow::Result;
use rag_module::RagModule;
use rag_module::services::search_service::{EstateSearchOptions, ChatSearchOptions};
use rag_module::StartSessionOptions;

#[tokio::main]
async fn main() -> Result<()> {
    println!("🔐 TESTING: Decrypt-First Search Implementation");
    println!("=============================================\n");

    // === TEST 1: ESTATE SEARCH WITH FILTERS ===
    println!("🏢 TEST 1: Estate Search with Region Filter (Decrypt-First)");
    
    let mut estate_rag = RagModule::new("./example-enhanced-estate-data").await?;
    estate_rag.initialize().await?;
    estate_rag.set_user_context("estate_user_example").await?;
    
    // Check collection health
    let health = estate_rag.collection_manager.get_collections_health().await?;
    println!("📊 Estate Collection: {} points", health.aws_estate.points_count);
    
    // Test estate search WITH region filter (should now work)
    let estate_search_options = EstateSearchOptions {
        resource_types: None,
        account_ids: None,
        regions: Some(vec!["us-east-1".to_string()]), // This should now work!
        services: None,
        states: None,
        environment: None,
        application: None,
        synced_after: None,
        limit: Some(5),
        score_threshold: Some(0.1),
        include_metadata: true,
        use_anonymous_ids: false,
    };
    
    println!("\n🔍 Testing estate search WITH region filter (decrypt-first approach)...");
    let estate_results = estate_rag.search_service
        .search_estate_resources("AWS EC2 instances", estate_search_options, None)
        .await?;
    
    println!("✅ Estate Search Results (With Region Filter): {}", estate_results.len());
    for (i, result) in estate_results.iter().take(3).enumerate() {
        if let (Some(resource_type), Some(region), Some(account)) = (
            result.get("resource_type").and_then(|v| v.as_str()),
            result.get("region").and_then(|v| v.as_str()),
            result.get("account_id").and_then(|v| v.as_str()),
        ) {
            let score = result.get("score").and_then(|v| v.as_f64()).unwrap_or(0.0);
            println!("   {}. {} in {} ({}) - Score: {:.4}", 
                i + 1, resource_type, region, account, score);
        }
    }

    // Test with multiple filters
    let multi_filter_options = EstateSearchOptions {
        resource_types: Some(vec!["ec2-instance".to_string()]),
        account_ids: None,
        regions: Some(vec!["us-east-1".to_string()]),
        services: Some(vec!["ec2".to_string()]),
        states: None,
        environment: None,
        application: None,
        synced_after: None,
        limit: Some(5),
        score_threshold: Some(0.1),
        include_metadata: true,
        use_anonymous_ids: false,
    };
    
    println!("\n🔍 Testing estate search WITH multiple filters...");
    let multi_results = estate_rag.search_service
        .search_estate_resources("EC2 instances", multi_filter_options, None)
        .await?;
    
    println!("✅ Estate Search Results (Multiple Filters): {}", multi_results.len());
    
    // === TEST 2: CHAT SEARCH WITH FILTERS ===
    println!("\n\n💬 TEST 2: Chat Search with Context Filter (Decrypt-First)");
    
    let mut chat_rag = RagModule::new("./example-enhanced-chat-data").await?;
    chat_rag.initialize().await?;
    chat_rag.set_user_context("chat_user_example").await?;
    
    // Create a test session and add some conversation
    let session_options = StartSessionOptions {
        user_id: "test_user_decrypt".to_string(),
        chat_title: Some("Decrypt Test Session".to_string()),
        context_id: Some("decrypt_test_context".to_string()),
    };
    
    let session = chat_rag.start_session(session_options).await?;
    println!("✅ Created session: {} (context: {})", session.id, session.context_id);
    
    // Add test conversations
    let conversations = vec![
        ("What is machine learning?", "Machine learning is AI that learns from data."),
        ("How does neural network work?", "Neural networks use interconnected nodes to process information."),
        ("Explain deep learning", "Deep learning uses multi-layer neural networks for complex pattern recognition."),
    ];
    
    let conversation_count = conversations.len();
    for (prompt, response) in conversations {
        let _prompt_id = chat_rag.add_prompt(&session.id, prompt, "test_user_decrypt").await?;
        let _response_id = chat_rag.add_response(&session.id, response, "test_user_decrypt").await?;
    }
    
    println!("✅ Added {} conversation pairs", conversation_count);
    
    // Test chat search with context filter
    let chat_search_options = ChatSearchOptions {
        context_id: Some(session.context_id.clone()),
        role: None, // All messages
        from_timestamp: None,
        to_timestamp: None,
        from_message_index: None,
        to_message_index: None,
        limit: Some(10),
        include_metadata: true,
    };
    
    println!("\n🔍 Testing chat search WITH context filter (decrypt-first approach)...");
    let chat_results = chat_rag.search_service
        .search_chat_history(chat_search_options)
        .await?;
    
    println!("✅ Chat Search Results (With Context Filter): {}", chat_results.len());
    for (i, result) in chat_results.iter().take(3).enumerate() {
        if let (Some(role), Some(content)) = (
            result.get("role").and_then(|v| v.as_str()),
            result.get("content").and_then(|v| v.as_str()),
        ) {
            let preview = if content.len() > 50 {
                format!("{}...", &content[..50])
            } else {
                content.to_string()
            };
            println!("   {}. {}: {}", i + 1, role, preview);
        }
    }
    
    // Test chat search with role filter
    let role_filter_options = ChatSearchOptions {
        context_id: Some(session.context_id.clone()),
        role: Some("user".to_string()), // Only user messages
        from_timestamp: None,
        to_timestamp: None,
        from_message_index: None,
        to_message_index: None,
        limit: Some(10),
        include_metadata: true,
    };
    
    println!("\n🔍 Testing chat search WITH role filter...");
    let role_results = chat_rag.search_service
        .search_chat_history(role_filter_options)
        .await?;
    
    println!("✅ Chat Search Results (User Messages Only): {}", role_results.len());
    
    // === FINAL SUMMARY ===
    println!("\n\n🎉 DECRYPT-FIRST SEARCH VERIFICATION COMPLETED!");
    println!("===============================================");
    println!("✅ Estate Search with Region Filter: {} results", estate_results.len());
    println!("✅ Estate Search with Multiple Filters: {} results", multi_results.len());
    println!("✅ Chat Search with Context Filter: {} results", chat_results.len());
    println!("✅ Chat Search with Role Filter: {} results", role_results.len());
    
    if estate_results.len() > 0 && chat_results.len() > 0 {
        println!("\n🚀 SUCCESS: Both search functions now work with encrypted metadata!");
        println!("   - Decrypt-first approach successfully implemented");
        println!("   - Filters work on decrypted data");
        println!("   - Universal solution for all collections");
    } else {
        println!("\n⚠️ Some searches returned 0 results - check implementation");
    }
    
    Ok(())
}