//! Test: Store AWS Estate Data with Encryption
//! Uses the Rust RAG module to store your AWS infrastructure data

use anyhow::Result;
use rag_module::{RagModule, services::search_service::EstateSearchOptions};
use serde_json::json;
use std::fs;
use std::path::Path;

#[tokio::main]
async fn main() -> Result<()> {
    println!("🚀 AWS Estate Data Storage Test with Encryption\n");

    // Your AWS estate data - exact same structure as Node.js version
    let aws_estate_data = json!([
        {
            "account_id": "288761761556",
            "role_arn": null,
            "account_name": "",
            "scan_timestamp": null,
            "services": {
                "ec2": {
                    "instances": [
                        {
                            "instance_id": "i-08619fa3385e63bab",
                            "instance_type": "c4.large",
                            "state": "running",
                            "region": "us-east-1",
                            "launch_time": "2025-08-26T12:37:51+00:00",
                            "name": "escher-scan-stream-v2"
                        },
                        {
                            "instance_id": "i-0f6ce804a1a464e1d",
                            "instance_type": "c4.large",
                            "state": "running",
                            "region": "us-east-1",
                            "launch_time": "2025-08-23T06:45:12+00:00",
                            "name": "slack-ec2-stream"
                        },
                        {
                            "instance_id": "i-01697357246074533",
                            "instance_type": "t2.micro",
                            "state": "running",
                            "region": "us-east-1",
                            "launch_time": "2025-09-05T07:06:47+00:00",
                            "name": "space-argo-dev-inst-1"
                        },
                        {
                            "instance_id": "i-0b1e8b61d9b846a3e",
                            "instance_type": "t2.micro",
                            "state": "running",
                            "region": "us-east-1",
                            "launch_time": "2025-07-21T04:22:19+00:00",
                            "name": "escher-ec2"
                        },
                        {
                            "instance_id": "i-0a5ff673a52c7a1c7",
                            "instance_type": "c4.large",
                            "state": "running",
                            "region": "us-east-1",
                            "launch_time": "2025-07-28T05:05:24+00:00",
                            "name": "escher-stream-ec2"
                        },
                        {
                            "instance_id": "i-06a446d8aaec15d27",
                            "instance_type": "c4.large",
                            "state": "running",
                            "region": "us-east-1",
                            "launch_time": "2025-08-29T04:11:05+00:00",
                            "name": "finops-ec2-stream"
                        },
                        {
                            "instance_id": "i-07901ab76d1aca487",
                            "instance_type": "c1.xlarge",
                            "state": "running",
                            "region": "us-east-1",
                            "launch_time": "2025-09-29T10:42:10+00:00",
                            "name": "dev-benchmark"
                        }
                    ],
                    "volumes_count": 13,
                    "elastic_ips_total": 7,
                    "elastic_ips_attached": 3,
                    "elastic_ips_unattached": 4,
                    "elastic_ips_instance_attached": 0,
                    "elastic_ips_eni_attached": 3,
                    "vpcs_count": 7
                },
                "rds": {
                    "instances": [
                        {
                            "instance_id": "db-VXDRDM3CQ6OIKGF7HLGVPPO454",
                            "db_instance_identifier": "dev-eshop-mysql-rds",
                            "db_instance_class": "db.t3.micro",
                            "engine": "mysql",
                            "db_instance_status": "available",
                            "region": "us-east-1",
                            "name": "dev-eshop-mysql-rds"
                        },
                        {
                            "instance_id": "db-2XYHZFX3PDXT3CD3LGIK435ALY",
                            "db_instance_identifier": "dev-tesserade-pg-rds",
                            "db_instance_class": "db.m5.xlarge",
                            "engine": "postgres",
                            "db_instance_status": "available",
                            "region": "us-east-1",
                            "name": "dev-tesserade-pg-rds"
                        },
                        {
                            "instance_id": "db-QTKHKSX6HCJBUAIZSSTOLAHQMY",
                            "db_instance_identifier": "mysql-test-instance1",
                            "db_instance_class": "db.t3.micro",
                            "engine": "mysql",
                            "db_instance_status": "available",
                            "region": "us-east-1",
                            "name": "mysql-test-instance1"
                        }
                    ]
                },
                "s3": {
                    "total_buckets": 23,
                    "public_buckets": 0,
                    "latest_buckets_info": [
                        {
                            "name": "escher-bucket-test-dev",
                            "s3_identifier": "escher-bucket-test-dev",
                            "creation_date": "2025-10-03T05:05:37+00:00",
                            "region": "us-east-1",
                            "is_public": false
                        },
                        {
                            "name": "test-escher-prod",
                            "s3_identifier": "test-escher-prod",
                            "creation_date": "2025-10-01T08:44:14+00:00",
                            "region": "us-east-1",
                            "is_public": false
                        }
                    ]
                },
                "lambda": {
                    "functions": [
                        {
                            "name": "eShopping_User_Simulator",
                            "lambda_identifier": "eShopping_User_Simulator",
                            "runtime": "python3.13",
                            "memory_size": 128,
                            "timeout": 3,
                            "region": "us-east-1"
                        },
                        {
                            "name": "prod-escher-storage-service",
                            "lambda_identifier": "prod-escher-storage-service",
                            "runtime": "provided.al2",
                            "memory_size": 256,
                            "timeout": 30,
                            "region": "us-east-1"
                        }
                    ]
                }
            }
        },
        {
            "account_id": "547639878546",
            "role_arn": "arn:aws:iam::547639878546:role/escher_cross_role",
            "account_name": "",
            "scan_timestamp": null,
            "services": {
                "ec2": {
                    "instances": [
                        {
                            "instance_id": "i-0f9f85890fc371f78",
                            "instance_type": "t2.micro",
                            "state": "running",
                            "region": "us-east-1",
                            "launch_time": "2025-08-22T18:19:11+00:00",
                            "name": "i-0f9f85890fc371f78"
                        }
                    ]
                },
                "s3": {
                    "latest_buckets_info": [
                        {
                            "name": "aparna-dev-test",
                            "s3_identifier": "aparna-dev-test",
                            "creation_date": "2025-10-03T07:16:28+00:00",
                            "region": "us-east-1",
                            "is_public": false
                        },
                        {
                            "name": "tex-instru-test-bckt",
                            "s3_identifier": "tex-instru-test-bckt",
                            "creation_date": "2025-07-24T04:50:07+00:00",
                            "region": "us-east-1",
                            "is_public": false
                        }
                    ]
                }
            }
        }
    ]);

    store_aws_estate_data(aws_estate_data).await?;
    Ok(())
}

async fn store_aws_estate_data(aws_estate_data: serde_json::Value) -> Result<()> {
    // Initialize RAG module with encryption enabled
    let mut rag_module = RagModule::new("./example-enhanced-chat-data").await?;

    println!("🔧 Initializing RAG module with encryption...");
    rag_module.initialize().await?;
    println!("✅ RAG module initialized successfully!\n");

    println!("📦 Processing AWS Estate Data...");
    let mut total_resources_stored = 0;

    // Specify user ID for user-specific storage
    let user_id = "cudewvsbxakj1n"; // User ID for encryption and file organization

    // NOTE: For AWS estate data, context_id is provided by the user (not from session)
    // Example context_id could be: "aws-estate-main", "aws-estate-prod", etc.
    // The context helps organize estate data separately from chat sessions

    println!("Processing complete AWS estate data structure for user: {}...", user_id);

    let document_ids = rag_module.process_aws_estate(aws_estate_data.clone(), user_id).await?;
    total_resources_stored = document_ids.len();

    // Count and display resources
    if let Some(accounts) = aws_estate_data.as_array() {
        for account in accounts {
            if let Some(account_obj) = account.as_object() {
                let account_id = account_obj.get("account_id")
                    .and_then(|v| v.as_str())
                    .unwrap_or("unknown");

                println!("\n🏢 Processing Account: {}", account_id);

                if let Some(services) = account_obj.get("services").and_then(|s| s.as_object()) {
                    // Count EC2 instances
                    if let Some(ec2) = services.get("ec2").and_then(|e| e.as_object()) {
                        if let Some(instances) = ec2.get("instances").and_then(|i| i.as_array()) {
                            println!("  📱 Processing {} EC2 instances...", instances.len());
                            for instance in instances {
                                if let Some(instance_obj) = instance.as_object() {
                                    let instance_id = instance_obj.get("instance_id").and_then(|v| v.as_str()).unwrap_or("unknown");
                                    println!("    ✅ Processed EC2: {}", instance_id);
                                }
                            }
                        }
                    }

                    // Count RDS instances
                    if let Some(rds) = services.get("rds").and_then(|r| r.as_object()) {
                        if let Some(instances) = rds.get("instances").and_then(|i| i.as_array()) {
                            println!("  🗄️ Processing {} RDS instances...", instances.len());
                            for db in instances {
                                if let Some(db_obj) = db.as_object() {
                                    let db_instance_identifier = db_obj.get("db_instance_identifier").and_then(|v| v.as_str()).unwrap_or("unknown");
                                    println!("    ✅ Processed RDS: {}", db_instance_identifier);
                                }
                            }
                        }
                    }

                    // Count S3 buckets
                    if let Some(s3) = services.get("s3").and_then(|s| s.as_object()) {
                        if let Some(buckets) = s3.get("latest_buckets_info").and_then(|b| b.as_array()) {
                            println!("  📁 Processing {} S3 buckets...", buckets.len());
                            for bucket in buckets {
                                if let Some(bucket_obj) = bucket.as_object() {
                                    let name = bucket_obj.get("name").and_then(|v| v.as_str()).unwrap_or("unknown");
                                    println!("    ✅ Processed S3: {}", name);
                                }
                            }
                        }
                    }

                    // Count Lambda functions
                    if let Some(lambda) = services.get("lambda").and_then(|l| l.as_object()) {
                        if let Some(functions) = lambda.get("functions").and_then(|f| f.as_array()) {
                            println!("  ⚡ Processing {} Lambda functions...", functions.len());
                            for function in functions {
                                if let Some(lambda_obj) = function.as_object() {
                                    let name = lambda_obj.get("name").and_then(|v| v.as_str()).unwrap_or("unknown");
                                    println!("    ✅ Processed Lambda: {}", name);
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    println!("\n🎉 Successfully stored {} AWS resources with encryption!", total_resources_stored);

    // Test search functionality
    println!("\n🔍 Testing Search Functionality...");

    // Test 1: Simple search with very low threshold
    println!("Test 1: Basic search with low threshold...");
    let estate_search_options = EstateSearchOptions {
        resource_types: None,
        account_ids: None,
        regions: None,
        services: None,
        states: None,
        environment: None,
        application: None,
        synced_after: None,
        limit: Some(10),
        score_threshold: Some(0.0), // Zero threshold to get all results
        include_metadata: false,
        use_anonymous_ids: true,
    };

    // Generate embedding for the query to see if embedding service works
    println!("  Generating embedding for query 'EC2'...");
    let query_embedding = rag_module.embedding_service.generate_embedding("EC2 running instances").await?;
    println!("  Generated embedding with {} dimensions", query_embedding.len());

    let search_results = rag_module.search_service.search_estate_resources("EC2", estate_search_options, None).await?;
    println!("✅ Found {} results for \"EC2\" with zero threshold", search_results.len());

    // Show detailed decrypted content
    for (i, result) in search_results.iter().take(3).enumerate() {
        println!("\n📋 Result {}: ID={}", i + 1, result.get("id").unwrap_or(&serde_json::Value::Null));

        // Show metadata
        if let Some(service) = result.get("service") {
            println!("    🔧 Service: {}", service);
        }
        if let Some(account_id) = result.get("account_id") {
            println!("    🏢 Account ID: {}", account_id);
        }
        if let Some(account_name) = result.get("account_name") {
            println!("    📛 Account Name: {}", account_name);
        }

        // Show full decrypted content
        if let Some(content) = result.get("content") {
            if let Some(content_str) = content.as_str() {
                println!("    📄 Decrypted Content:");
                // Try to parse as JSON for pretty printing
                if let Ok(json_content) = serde_json::from_str::<serde_json::Value>(content_str) {
                    let pretty_json = serde_json::to_string_pretty(&json_content).unwrap_or_else(|_| content_str.to_string());
                    // Show first 500 chars of formatted content
                    let display_content = if pretty_json.len() > 500 {
                        format!("{}...\n        [Content truncated - {} total chars]", &pretty_json[..500], pretty_json.len())
                    } else {
                        pretty_json
                    };
                    println!("        {}", display_content.replace('\n', "\n        "));
                } else {
                    // Fallback to raw content if not valid JSON
                    let display_content = if content_str.len() > 200 {
                        format!("{}...", &content_str[..200])
                    } else {
                        content_str.to_string()
                    };
                    println!("        {}", display_content);
                }
            }
        }
        println!("    ─────────────────────────────────────────────────────");
    }

    // Test 2: Search with filters
    println!("Test 2: Search with EC2 filter...");
    let estate_search_options_filtered = EstateSearchOptions {
        resource_types: Some(vec!["instance".to_string()]),
        account_ids: None,
        regions: None,
        services: Some(vec!["ec2".to_string()]),
        states: Some(vec!["running".to_string()]),
        environment: None,
        application: None,
        synced_after: None,
        limit: Some(5),
        score_threshold: Some(0.01), // Very low threshold
        include_metadata: false,
        use_anonymous_ids: true,
    };
    let filtered_results = rag_module.search_service.search_estate_resources("instance", estate_search_options_filtered, None).await?;
    println!("✅ Found {} results for filtered EC2 instance search", filtered_results.len());

    // Test 3: Check encrypted documents and vector database
    println!("Test 3: Check encrypted documents and vector database...");
    let all_encrypted_docs = rag_module.get_encrypted_aws_estate_documents().await?;
    println!("✅ Total encrypted estate documents: {}", all_encrypted_docs.len());

    // Test direct vector store search (bypassing search service)
    println!("Test 4: Direct vector database search...");
    let search_options = rag_module::types::SearchOptions {
        limit: Some(10),
        score_threshold: Some(0.0), // Zero threshold
        with_payload: Some(true),
        filter: None,
        ..Default::default()
    };
    let direct_results = rag_module.vector_store.search("aws_estate", query_embedding, search_options).await?;
    println!("✅ Direct vector search found {} results", direct_results.len());

    // Show direct search results
    for (i, result) in direct_results.iter().take(2).enumerate() {
        println!("  Direct Result {}: ID={}, Score={:.4}", i + 1, result.id, result.score);
        if let Some(payload) = &result.payload {
            if let Some(service) = payload.get("service") {
                println!("    Service: {}", service);
            }
            if let Some(account_id) = payload.get("account_id") {
                println!("    Account: {}", account_id);
            }
        }
    }

    // Show stats (approximate - Rust version doesn't have exact getStats method)
    println!("\n📊 RAG Module Statistics:");
    let encrypted_docs = rag_module.get_encrypted_aws_estate_documents().await?;
    println!("   📋 Total Estate Documents: {}", encrypted_docs.len());
    println!("   🔒 Encryption Enabled: true");
    println!("   📁 Base Path: ./test-estate-data");

    // Check file structure
    println!("\n📁 Checking stored file structure...");
    let data_path = "./test-estate-data/qdrant-data";

    if Path::new(data_path).exists() {
        println!("✅ Qdrant data directory created");

        // Check for estate documents file
        let estate_docs_path = format!("{}/aws_estate-metadata.json", data_path);
        if Path::new(&estate_docs_path).exists() {
            println!("✅ AWS estate metadata file created");

            // Show preview of encrypted file
            if let Ok(content) = fs::read_to_string(&estate_docs_path) {
                let preview = if content.len() > 100 {
                    format!("{}...", &content[..100])
                } else {
                    content
                };
                println!("   🔒 Estate metadata preview (encrypted): {}", preview);
            }
        }

        // Check segments directory
        let segments_path = format!("{}/segments", data_path);
        if Path::new(&segments_path).exists() {
            println!("✅ Vector segments directory created");
        }
    }

    println!("\n🏁 Test completed");
    Ok(())
}