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

#[tokio::main]
async fn main() -> Result<()> {
    println!("🔍 VERIFICATION: Testing Both Search Functions with Real Data");
    println!("=========================================================\n");

    // === TEST 1: ESTATE SEARCH WITH EXISTING DATA ===
    println!("🏢 TEST 1: Estate Search with Existing AWS Data");
    println!("Using existing estate data from previous examples...");
    
    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 what data exists
    let estate_count = estate_rag.get_document_count(None).await?;
    println!("✅ Estate documents available: {}", estate_count);
    
    let health = estate_rag.collection_manager.get_collections_health().await?;
    println!("✅ Estate collection points: {}", health.aws_estate.points_count);
    
    // Test multiple estate searches with different queries
    let estate_search_configs = vec![
        ("running EC2 instances", EstateSearchOptions {
            resource_types: Some(vec!["ec2-instance".to_string()]),
            account_ids: None,
            regions: None,
            services: None,
            states: Some(vec!["running".to_string()]),
            environment: None,
            application: None,
            synced_after: None,
            limit: Some(5),
            score_threshold: Some(0.1), // Lower threshold
            include_metadata: true,
            use_anonymous_ids: false,
        }),
        ("RDS databases", EstateSearchOptions {
            resource_types: Some(vec!["rds-instance".to_string()]),
            account_ids: None,
            regions: None,
            services: None,
            states: None,
            environment: None,
            application: None,
            synced_after: None,
            limit: Some(5),
            score_threshold: Some(0.1), // Lower threshold
            include_metadata: true,
            use_anonymous_ids: false,
        }),
        ("AWS resources in us-east-1", EstateSearchOptions {
            resource_types: None,
            account_ids: None,
            regions: Some(vec!["us-east-1".to_string()]),
            services: None,
            states: None,
            environment: None,
            application: None,
            synced_after: None,
            limit: Some(10),
            score_threshold: Some(0.05), // Very low threshold to see more results
            include_metadata: true,
            use_anonymous_ids: false,
        }),
    ];
    
    for (query, options) in estate_search_configs {
        println!("\n🔍 Estate Search Query: '{}'", query);
        let results = estate_rag.search_service
            .search_estate_resources(query, options, None)
            .await?;
        
        println!("   ✅ Results found: {}", results.len());
        for (i, result) in results.iter().take(3).enumerate() {
            if let (Some(resource_type), Some(account_id), Some(region)) = (
                result.get("resource_type").and_then(|v| v.as_str()),
                result.get("account_id").and_then(|v| v.as_str()),
                result.get("region").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_id, score);
            }
        }
    }

    // === TEST 2: CHAT SEARCH WITH NEW DATA ===
    println!("\n\n💬 TEST 2: Chat Search with Fresh Conversation Data");
    println!("Creating new chat data for testing...");
    
    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_search".to_string(),
        chat_title: Some("Search Test Session".to_string()),
        context_id: Some("search_test_context".to_string()),
    };
    
    let session = chat_rag.start_session(session_options).await?;
    println!("✅ Created session: {} (context: {})", session.id, session.context_id);
    
    // Add multiple conversation pairs
    let conversations = vec![
        ("What is machine learning?", "Machine learning is a subset of artificial intelligence that enables computers to learn and make decisions from data without being explicitly programmed."),
        ("How does deep learning work?", "Deep learning uses artificial neural networks with multiple layers to model and understand complex patterns in data, mimicking how the human brain processes information."),
        ("What are the types of machine learning?", "There are three main types: supervised learning (learning with labeled data), unsupervised learning (finding patterns in unlabeled data), and reinforcement learning (learning through interaction and rewards)."),
    ];
    
    for (prompt, response) in conversations {
        let prompt_id = chat_rag.add_prompt(&session.id, prompt, "test_user_search").await?;
        let response_id = chat_rag.add_response(&session.id, response, "test_user_search").await?;
        println!("   ✅ Added conversation pair: {} -> {}", prompt_id, response_id);
    }
    
    // Test chat history searches
    let chat_search_configs = vec![
        ("All messages in context", ChatSearchOptions {
            context_id: Some(session.context_id.clone()),
            role: None,
            from_timestamp: None,
            to_timestamp: None,
            from_message_index: None,
            to_message_index: None,
            limit: Some(10),
            include_metadata: true,
        }),
        ("User messages only", ChatSearchOptions {
            context_id: Some(session.context_id.clone()),
            role: Some("user".to_string()),
            from_timestamp: None,
            to_timestamp: None,
            from_message_index: None,
            to_message_index: None,
            limit: Some(10),
            include_metadata: true,
        }),
        ("Assistant messages only", ChatSearchOptions {
            context_id: Some(session.context_id.clone()),
            role: Some("assistant".to_string()),
            from_timestamp: None,
            to_timestamp: None,
            from_message_index: None,
            to_message_index: None,
            limit: Some(10),
            include_metadata: true,
        }),
    ];
    
    for (description, options) in chat_search_configs {
        println!("\n🔍 Chat Search: '{}'", description);
        let results = chat_rag.search_service
            .search_chat_history(options)
            .await?;
        
        println!("   ✅ Results found: {}", results.len());
        for (i, result) in 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() > 60 {
                    format!("{}...", &content[..60])
                } else {
                    content.to_string()
                };
                println!("   {}. {}: {}", i + 1, role, preview);
            }
        }
    }

    // === TEST 3: COMPARISON WITH DIRECT VECTOR SEARCH ===
    println!("\n\n🔧 TEST 3: Comparing Enhanced Search vs Direct Vector Search");
    
    // Test direct vector search for estate
    println!("\n Estate - Direct Vector Search:");
    let query_embedding = estate_rag.embedding_service.generate_embedding("EC2 instances").await?;
    
    use rag_module::types::SearchOptions;
    let direct_search_options = SearchOptions {
        limit: Some(5),
        score_threshold: Some(0.1),
        filter: None,
        with_payload: Some(true),
        ..Default::default()
    };
    
    let direct_results = estate_rag.vector_store.search("aws_estate", query_embedding, direct_search_options).await?;
    println!("   ✅ Direct vector search results: {}", direct_results.len());
    for (i, result) in direct_results.iter().take(3).enumerate() {
        println!("   {}. ID: {}, Score: {:.4}", i + 1, result.id, result.score);
    }
    
    // Compare with enhanced search
    println!("\n Estate - Enhanced Search:");
    let enhanced_results = estate_rag.search_service
        .search_estate_resources("EC2 instances", EstateSearchOptions {
            resource_types: None,
            account_ids: None,
            regions: None,
            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,
        }, None)
        .await?;
    
    println!("   ✅ Enhanced search results: {}", enhanced_results.len());
    for (i, result) in enhanced_results.iter().take(3).enumerate() {
        if let Some(score) = result.get("score").and_then(|v| v.as_f64()) {
            println!("   {}. Score: {:.4}", i + 1, score);
        }
    }

    // === FINAL SUMMARY ===
    println!("\n\n🎉 SEARCH VERIFICATION COMPLETED!");
    println!("=====================================");
    println!("✅ Estate Search Function: Working with {} documents", estate_count);
    println!("✅ Chat Search Function: Working with conversation data");
    println!("✅ Both search functions are properly indexing and retrieving data");
    println!("✅ Enhanced search filters are working as designed");

    Ok(())
}