//! Comprehensive test to verify AWS estate service flow
//! Tests document creation count, storage, and retrieval with all service parsers

use serde_json::json;
use rag_module::RagModule;
use std::path::Path;
use tokio::fs;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    println!("🧪 Starting Comprehensive AWS Estate Service Test");
    
    // Clean up previous test data
    let test_dir = "./comprehensive_test_data";
    if Path::new(test_dir).exists() {
        fs::remove_dir_all(test_dir).await?;
    }
    
    // Initialize the RAG module
    let rag = RagModule::new(test_dir).await?;
    rag.initialize().await?;
    
    // Set user context
    rag.set_user_context("test-user-123").await?;
    
    println!("✅ RAG Module initialized");
    
    // Comprehensive test data with all services
    let aws_estate_data = json!([
        {
            "account_id": "123456789012",
            "role_arn": null,
            "account_name": "test-account",
            "scan_timestamp": "2024-11-05T12:00:00Z",
            "services": {
                "ec2": {
                    "instances": [
                        {
                            "instance_id": "i-1234567890abcdef0",
                            "instance_type": "t3.medium",
                            "state": "running",
                            "region": "us-west-2",
                            "launch_time": "2024-11-01T10:00:00Z",
                            "name": "web-server-1",
                            "availability_zone": "us-west-2a",
                            "security_groups": ["sg-12345", "sg-67890"],
                            "permissions": [
                                "ec2:StartInstances",
                                "ec2:StopInstances",
                                "ec2:DescribeInstances",
                                "ec2:RebootInstances"
                            ]
                        },
                        {
                            "instance_id": "i-0987654321fedcba0",
                            "instance_type": "c5.large",
                            "state": "stopped",
                            "region": "us-west-2",
                            "launch_time": "2024-10-15T08:30:00Z",
                            "name": "batch-processor",
                            "availability_zone": "us-west-2b",
                            "permissions": [
                                "ec2:StartInstances",
                                "ec2:StopInstances"
                            ]
                        }
                    ],
                    "volumes_count": 15,
                    "total_size_gb": 500,
                    "attached_volumes": 10,
                    "unattached_volumes": 5,
                    "vpcs_count": 3,
                    "total_subnets": 12,
                    "internet_gateways_count": 2,
                    "elastic_ips_total": 5,
                    "elastic_ips_attached": 3
                },
                "rds": {
                    "instances": [
                        {
                            "instance_id": "db-instance-1",
                            "db_instance_identifier": "prod-database",
                            "region": "us-west-2",
                            "db_instance_class": "db.r5.large",
                            "engine": "postgresql",
                            "db_instance_status": "available",
                            "name": "production-db",
                            "multi_az": true,
                            "allocated_storage": 100,
                            "storage_type": "gp2",
                            "permissions": [
                                "rds:DescribeDBInstances",
                                "rds:RebootDBInstance"
                            ]
                        }
                    ],
                    "clusters_count": 2,
                    "snapshots_count": 10
                },
                "s3": {
                    "buckets": [
                        {
                            "bucket_name": "my-app-data-bucket",
                            "region": "us-west-2",
                            "creation_date": "2024-01-15T00:00:00Z",
                            "versioning": "Enabled",
                            "encryption": "AES256",
                            "public_access_blocked": true,
                            "object_count": 1250,
                            "size_bytes": 52428800,
                            "permissions": [
                                "s3:GetObject",
                                "s3:PutObject",
                                "s3:ListBucket"
                            ]
                        },
                        {
                            "bucket_name": "backup-storage-bucket",
                            "region": "us-east-1",
                            "creation_date": "2024-02-01T00:00:00Z",
                            "versioning": "Disabled",
                            "encryption": "Disabled",
                            "public_access_blocked": true,
                            "permissions": [
                                "s3:GetObject",
                                "s3:ListBucket"
                            ]
                        }
                    ],
                    "bucket_policies_count": 3,
                    "public_bucket_policies": 1
                },
                "lambda": {
                    "functions": [
                        {
                            "function_name": "data-processor",
                            "region": "us-west-2",
                            "runtime": "python3.9",
                            "timeout": 300,
                            "memory_size": 512,
                            "handler": "lambda_function.lambda_handler",
                            "last_modified": "2024-11-01T12:00:00Z",
                            "code_size": 1024000,
                            "role": "arn:aws:iam::123456789012:role/lambda-execution-role",
                            "environment": {
                                "Variables": {
                                    "ENV": "production",
                                    "DB_HOST": "prod-database.region.rds.amazonaws.com"
                                }
                            },
                            "permissions": [
                                "lambda:InvokeFunction",
                                "lambda:UpdateFunctionCode"
                            ]
                        }
                    ],
                    "layers_count": 5,
                    "event_sources_count": 3
                },
                "ebs": {
                    "volumes": [
                        {
                            "volume_id": "vol-1234567890abcdef0",
                            "region": "us-west-2",
                            "size": 100,
                            "volume_type": "gp3",
                            "state": "in-use",
                            "availability_zone": "us-west-2a",
                            "encrypted": true,
                            "create_time": "2024-11-01T10:00:00Z",
                            "attachment_state": "attached",
                            "attached_instance_id": "i-1234567890abcdef0",
                            "iops": 3000,
                            "throughput": 125,
                            "permissions": [
                                "ec2:AttachVolume",
                                "ec2:DetachVolume"
                            ]
                        }
                    ],
                    "snapshots_count": 25,
                    "snapshots_total_size_gb": 1000,
                    "encrypted_snapshots": 20,
                    "total_volumes": 15,
                    "encrypted_volumes": 12
                }
            }
        }
    ]);
    
    println!("📊 Test Data Summary:");
    println!("   • 1 AWS Account");
    println!("   • 5 AWS Services (EC2, RDS, S3, Lambda, EBS)");
    
    // Calculate expected document count
    let expected_documents = calculate_expected_documents(&aws_estate_data);
    println!("   • Expected Documents: {}", expected_documents);
    
    // Store the estate data
    println!("\n📥 Storing AWS estate data...");
    let store_result = rag.store_aws_estate_data(aws_estate_data).await?;
    
    println!("✅ Storage completed!");
    println!("   • Total accounts processed: {}", store_result.total_accounts);
    println!("   • Total services processed: {}", store_result.total_services);
    println!("   • Total resources found: {}", store_result.total_resources);
    println!("   • Successfully parsed: {}", store_result.parsed_resources);
    println!("   • Failed to parse: {}", store_result.failed_resources);
    println!("   • Documents created: {}", store_result.create_result.created);
    println!("   • Supported services: {:?}", store_result.supported_services);
    
    // Verify expected vs actual document count
    println!("\n📈 Document Count Analysis:");
    println!("   • Expected: {}", expected_documents);
    println!("   • Actually created: {}", store_result.create_result.created);
    
    if store_result.create_result.created == expected_documents {
        println!("   ✅ Document count MATCHES expectations!");
    } else {
        println!("   ❌ Document count MISMATCH!");
        println!("   • Difference: {}", 
            (expected_documents as i32) - (store_result.create_result.created as i32));
    }
    
    // Check local file storage
    println!("\n📁 Checking local file storage...");
    verify_local_storage(test_dir, "test-user-123").await?;
    
    // Test search functionality
    println!("\n🔍 Testing search functionality...");
    
    // Test EC2 search
    println!("\n   Testing: 'list my ec2 instances'");
    let ec2_results = rag.search_estate_resources(
        "list my ec2 instances", 
        None, 
        None, 
        None
    ).await?;
    
    println!("   📋 EC2 Search Results:");
    println!("      • Found {} results", ec2_results.len());
    
    for (i, result) in ec2_results.iter().enumerate() {
        if let Some(payload) = &result.payload {
            println!("      Result {}: Score: {:.4}", i + 1, result.score);
            println!("         • ID: {}", result.id);
            println!("         • Instance ID: {}", 
                payload.get("instance_id")
                    .and_then(|v| v.as_str())
                    .unwrap_or("N/A"));
            println!("         • IAM Permissions: {}", 
                payload.get("iam_permissions")
                    .map(|v| serde_json::to_string_pretty(v).unwrap_or("N/A".to_string()))
                    .unwrap_or("N/A".to_string()));
        }
    }
    
    // Test S3 search
    println!("\n   Testing: 'show s3 buckets'");
    let s3_results = rag.search_estate_resources(
        "show s3 buckets", 
        None, 
        None, 
        None
    ).await?;
    
    println!("   📋 S3 Search Results:");
    println!("      • Found {} results", s3_results.len());
    
    for (i, result) in s3_results.iter().enumerate() {
        if let Some(payload) = &result.payload {
            println!("      Result {}: Score: {:.4}", i + 1, result.score);
            println!("         • Bucket: {}", 
                payload.get("bucket_name")
                    .and_then(|v| v.as_str())
                    .unwrap_or("N/A"));
            println!("         • Encryption: {}", 
                payload.get("encryption")
                    .and_then(|v| v.as_str())
                    .unwrap_or("N/A"));
        }
    }
    
    // Test Lambda search
    println!("\n   Testing: 'lambda functions'");
    let lambda_results = rag.search_estate_resources(
        "lambda functions", 
        None, 
        None, 
        None
    ).await?;
    
    println!("   📋 Lambda Search Results:");
    println!("      • Found {} results", lambda_results.len());
    
    println!("\n🎉 Comprehensive test completed!");
    
    // Print final summary
    print_final_summary(&store_result, expected_documents, &ec2_results, &s3_results, &lambda_results);
    
    Ok(())
}

fn calculate_expected_documents(data: &serde_json::Value) -> usize {
    let mut count = 0;
    
    if let Some(accounts) = data.as_array() {
        for account in accounts {
            if let Some(services) = account["services"].as_object() {
                
                // EC2 documents
                if let Some(ec2) = services.get("ec2") {
                    // EC2 instances
                    if let Some(instances) = ec2["instances"].as_array() {
                        count += instances.len();
                    }
                    // EBS volumes summary
                    if ec2.get("volumes_count").is_some() {
                        count += 1;
                    }
                    // VPC summary
                    if ec2.get("vpcs_count").is_some() {
                        count += 1;
                    }
                }
                
                // RDS documents
                if let Some(rds) = services.get("rds") {
                    // RDS instances
                    if let Some(instances) = rds["instances"].as_array() {
                        count += instances.len();
                    }
                    // Clusters summary
                    if rds["clusters_count"].as_u64().unwrap_or(0) > 0 {
                        count += 1;
                    }
                    // Snapshots summary
                    if rds["snapshots_count"].as_u64().unwrap_or(0) > 0 {
                        count += 1;
                    }
                }
                
                // S3 documents
                if let Some(s3) = services.get("s3") {
                    // S3 buckets
                    if let Some(buckets) = s3["buckets"].as_array() {
                        count += buckets.len();
                    }
                    // Bucket policies summary
                    if s3["bucket_policies_count"].as_u64().unwrap_or(0) > 0 {
                        count += 1;
                    }
                }
                
                // Lambda documents
                if let Some(lambda) = services.get("lambda") {
                    // Lambda functions
                    if let Some(functions) = lambda["functions"].as_array() {
                        count += functions.len();
                    }
                    // Layers summary
                    if lambda["layers_count"].as_u64().unwrap_or(0) > 0 {
                        count += 1;
                    }
                    // Event sources summary
                    if lambda["event_sources_count"].as_u64().unwrap_or(0) > 0 {
                        count += 1;
                    }
                }
                
                // EBS documents
                if let Some(ebs) = services.get("ebs") {
                    // EBS volumes
                    if let Some(volumes) = ebs["volumes"].as_array() {
                        count += volumes.len();
                    }
                    // Snapshots summary
                    if ebs["snapshots_count"].as_u64().unwrap_or(0) > 0 {
                        count += 1;
                    }
                    // Encryption summary
                    if ebs.get("total_volumes").is_some() {
                        count += 1;
                    }
                }
            }
        }
    }
    
    count
}

async fn verify_local_storage(test_dir: &str, user_id: &str) -> anyhow::Result<()> {
    let qdrant_data_dir = Path::new(test_dir).join("qdrant-data").join(user_id);
    
    if !qdrant_data_dir.exists() {
        println!("   ❌ User directory not found: {:?}", qdrant_data_dir);
        return Ok(());
    }
    
    println!("   ✅ User directory exists: {:?}", qdrant_data_dir);
    
    // Check for aws_estate collection files
    let documents_file = qdrant_data_dir.join("aws_estate-documents.json");
    let metadata_file = qdrant_data_dir.join("aws_estate-metadata.json");
    let vectors_file = qdrant_data_dir.join("aws_estate-vectors.bin");
    let index_file = qdrant_data_dir.join("aws_estate-vector-index.json");
    
    println!("   📄 File existence check:");
    println!("      • Documents file: {}", documents_file.exists());
    println!("      • Metadata file: {}", metadata_file.exists());
    println!("      • Vectors file: {}", vectors_file.exists());
    println!("      • Index file: {}", index_file.exists());
    
    // Check document count in local file
    if documents_file.exists() {
        let content = fs::read_to_string(&documents_file).await?;
        if let Ok(wrapper) = serde_json::from_str::<serde_json::Value>(&content) {
            if let Some(docs) = wrapper["documents"].as_array() {
                println!("      • Documents in local file: {}", docs.len());
            }
            if let Some(count) = wrapper["count"].as_u64() {
                println!("      • Count field in wrapper: {}", count);
            }
        }
    }
    
    Ok(())
}

fn print_final_summary(
    store_result: &rag_module::types::AwsEstateIngestResult,
    expected_docs: usize,
    ec2_results: &[rag_module::types::SearchResult],
    s3_results: &[rag_module::types::SearchResult],
    lambda_results: &[rag_module::types::SearchResult],
) {
    println!("\n" + "=".repeat(60).as_str());
    println!("📊 COMPREHENSIVE TEST SUMMARY");
    println!("=".repeat(60));
    
    println!("\n🏗️  DOCUMENT CREATION:");
    println!("   • Expected documents: {}", expected_docs);
    println!("   • Actually created: {}", store_result.create_result.created);
    println!("   • Match: {}", if store_result.create_result.created == expected_docs { "✅ YES" } else { "❌ NO" });
    
    println!("\n🔧 SERVICE PARSING:");
    println!("   • Supported services: {}", store_result.supported_services.len());
    println!("   • Services: {:?}", store_result.supported_services);
    
    println!("\n🔍 SEARCH TESTING:");
    println!("   • EC2 search results: {} documents", ec2_results.len());
    println!("   • S3 search results: {} documents", s3_results.len());
    println!("   • Lambda search results: {} documents", lambda_results.len());
    
    println!("\n🎯 OVERALL STATUS:");
    if store_result.create_result.created == expected_docs && 
       !ec2_results.is_empty() && 
       !s3_results.is_empty() {
        println!("   ✅ ALL TESTS PASSED - System working correctly!");
    } else {
        println!("   ❌ SOME TESTS FAILED - Issues detected!");
    }
    
    println!("\n" + "=".repeat(60).as_str());
}