#![cfg(feature = "disabled_tests")]

use rtx_preprocessing::{
    FaultToleranceConfig,
    // Loaders
    IntegratedDataLoader,
    IntegratedLoaderConfig,
    LoaderConfig,
    MinMaxScaler,
    PrefetchStrategy,
    PreprocessingError,
    RebalancingStrategy,
    // Core types
    Result,
    ShardingConfig,
    ShardingStrategy,
    // Transformers
    StandardScaler,
    Transformer,
    WorkerInfo,
    rtx_stubs::rtx_tensor::{Device, Tensor},
};
use std::fs::{File, create_dir_all};
use std::io::Write;
use std::path::{Path, PathBuf};
use tempfile::TempDir;

/// Helper function to create test dataset with known patterns
fn create_structured_test_dataset(
    dir: &Path,
    num_files: usize,
    samples_per_file: usize,
) -> Vec<PathBuf> {
    let mut file_paths = Vec::new();

    for file_idx in 0..num_files {
        let file_name = format!("dataset_{:03}.bin", file_idx);
        let file_path = dir.join(&file_name);
        let mut file = File::create(&file_path).unwrap();

        // Create data with a predictable pattern for testing transformers
        // Each file has a different base value to make sharding effects visible
        let base_value = file_idx as f32 * 10.0;

        for i in 0..samples_per_file {
            // Create a simple linear pattern: base + i + some variance
            let variance = (i % 7) as f32 * 0.1; // Small variance for testing
            let value = base_value + (i as f32) + variance;
            file.write_all(&value.to_le_bytes()).unwrap();
        }
        file.flush().unwrap();

        file_paths.push(file_path);
    }

    file_paths
}

#[test]
fn test_integrated_loader_initialization() {
    let temp_dir = TempDir::new().unwrap();
    let file_paths = create_structured_test_dataset(temp_dir.path(), 6, 200);

    let config = IntegratedLoaderConfig {
        mmap_config: LoaderConfig {
            page_size: 4096,
            prefetch_strategy: PrefetchStrategy::Sequential { window_size: 2 },
            use_shared_memory: true,
            enable_lazy_loading: true,
        },
        sharding_config: ShardingConfig {
            num_shards: 3,
            sharding_strategy: ShardingStrategy::RoundRobin,
            rebalancing_strategy: RebalancingStrategy::Static,
            fault_tolerance: FaultToleranceConfig::default(),
        },
        enable_cross_shard_prefetch: true,
        batch_size: 100,
    };

    let workers = vec![
        WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
        WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
        WorkerInfo::new(2, "worker2".to_string(), "127.0.0.1:8002".to_string()),
    ];

    let mut loader = IntegratedDataLoader::new(config, workers);

    // Test initialization
    let result = loader.initialize(&file_paths);
    assert!(result.is_ok());
    assert!(loader.is_initialized());

    // Verify distributed sharding worked
    let distributed_stats = loader.distributed_loader().get_statistics();
    assert_eq!(distributed_stats.total_shards, 3);
    assert_eq!(distributed_stats.total_files, 6);
    assert_eq!(distributed_stats.active_workers, 3);
}

#[test]
fn test_integrated_loader_with_transformers() {
    let temp_dir = TempDir::new().unwrap();
    let file_paths = create_structured_test_dataset(temp_dir.path(), 4, 100);

    let config = IntegratedLoaderConfig::default();
    let workers = vec![
        WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
        WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
    ];

    let mut loader = IntegratedDataLoader::new(config, workers);
    loader.initialize(&file_paths).unwrap();

    let device = Device::cuda(0).unwrap_or(Device::default());

    // Load data from worker 0
    let worker_batch = loader.load_worker_batch(0, 0, &device).unwrap();
    assert!(!worker_batch.is_empty());

    // Test with StandardScaler
    let tensor = &worker_batch[0];
    let mut scaler = StandardScaler::new();

    // Fit the scaler
    let fit_result = scaler.fit(tensor);
    assert!(fit_result.is_ok());
    assert!(scaler.is_fitted());

    // Transform the data
    let scaled_tensor = scaler.transform(tensor).unwrap();

    // Verify the transformation worked
    assert_eq!(scaled_tensor.shape().dims(), tensor.shape().dims());

    // The scaled data should have different values than original
    let original_data: Vec<f32> = tensor.iter().map(|&x| x as f32).take(10).collect();
    let scaled_data: Vec<f32> = scaled_tensor.iter().map(|&x| x as f32).take(10).collect();

    // Values should be different after scaling (unless variance is exactly 1.0)
    let has_different_values = original_data
        .iter()
        .zip(scaled_data.iter())
        .any(|(&orig, &scaled)| (orig - scaled).abs() > 1e-6);

    assert!(has_different_values);
}

#[test]
fn test_distributed_batch_loading_with_preprocessing() {
    let temp_dir = TempDir::new().unwrap();
    let file_paths = create_structured_test_dataset(temp_dir.path(), 8, 50);

    let config = IntegratedLoaderConfig {
        sharding_config: ShardingConfig {
            num_shards: 4,
            sharding_strategy: ShardingStrategy::SizeAware,
            rebalancing_strategy: RebalancingStrategy::Static,
            fault_tolerance: FaultToleranceConfig::default(),
        },
        batch_size: 25,
        ..Default::default()
    };

    let workers = vec![
        WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
        WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
    ];

    let mut loader = IntegratedDataLoader::new(config, workers.clone());
    loader.initialize(&file_paths).unwrap();

    let device = Device::cuda(0).unwrap_or(Device::default());

    // Load distributed batch
    let distributed_batch = loader.load_distributed_batch(0, &device).unwrap();

    // Should have data for both workers
    assert!(distributed_batch.contains_key(&0));
    assert!(distributed_batch.contains_key(&1));

    // Process each worker's data with different scalers
    let mut worker_scalers: std::collections::HashMap<u32, (StandardScaler, MinMaxScaler)> =
        std::collections::HashMap::new();

    for (&worker_id, tensors) in &distributed_batch {
        let mut std_scaler = StandardScaler::new();
        let mut minmax_scaler = MinMaxScaler::new();

        if let Some(tensor) = tensors.first() {
            // Fit both scalers
            std_scaler.fit(tensor).unwrap();
            minmax_scaler.fit(tensor).unwrap();

            worker_scalers.insert(worker_id, (std_scaler, minmax_scaler));
        }
    }

    // Verify that both workers have fitted scalers
    assert_eq!(worker_scalers.len(), 2);

    // Transform data for each worker
    for (&worker_id, tensors) in &distributed_batch {
        if let Some((std_scaler, minmax_scaler)) = worker_scalers.get(&worker_id) {
            for tensor in tensors {
                let std_transformed = std_scaler.transform(tensor).unwrap();
                let minmax_transformed = minmax_scaler.transform(tensor).unwrap();

                // Verify transformations
                assert_eq!(std_transformed.shape().dims(), tensor.shape().dims());
                assert_eq!(minmax_transformed.shape().dims(), tensor.shape().dims());
            }
        }
    }
}

#[test]
fn test_fault_tolerance_with_preprocessing_pipeline() {
    let temp_dir = TempDir::new().unwrap();
    let file_paths = create_structured_test_dataset(temp_dir.path(), 6, 100);

    let config = IntegratedLoaderConfig {
        sharding_config: ShardingConfig {
            num_shards: 3,
            sharding_strategy: ShardingStrategy::RoundRobin,
            rebalancing_strategy: RebalancingStrategy::Dynamic {
                imbalance_threshold: 0.3,
                rebalance_interval_secs: 5,
            },
            fault_tolerance: FaultToleranceConfig {
                enable_health_checks: true,
                health_check_interval_secs: 1,
                max_failures: 2,
                recovery_timeout_secs: 10,
            },
        },
        batch_size: 50,
        ..Default::default()
    };

    let workers = vec![
        WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
        WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
        WorkerInfo::new(2, "worker2".to_string(), "127.0.0.1:8002".to_string()),
    ];

    let mut loader = IntegratedDataLoader::new(config, workers);
    loader.initialize(&file_paths).unwrap();

    let device = Device::cuda(0).unwrap_or(Device::default());

    // Initial state: all workers healthy
    assert!(loader.distributed_loader().is_worker_healthy(0));
    assert!(loader.distributed_loader().is_worker_healthy(1));
    assert!(loader.distributed_loader().is_worker_healthy(2));

    // Simulate worker failure
    loader.distributed_loader().report_worker_failure(1);
    loader.distributed_loader().report_worker_failure(1); // Make worker 1 unhealthy

    assert!(!loader.distributed_loader().is_worker_healthy(1));

    // Load data should still work with remaining workers
    let distributed_batch = loader.load_distributed_batch(0, &device).unwrap();

    // Should have data from healthy workers only
    assert!(distributed_batch.contains_key(&0));
    assert!(distributed_batch.contains_key(&2));
    assert!(!distributed_batch.contains_key(&1)); // Failed worker should have no data

    // Test preprocessing still works
    for tensors in distributed_batch.values() {
        for tensor in tensors {
            let mut scaler = StandardScaler::new();
            scaler.fit(tensor).unwrap();
            let transformed = scaler.transform(tensor).unwrap();
            assert_eq!(transformed.shape().dims(), tensor.shape().dims());
        }
    }

    // Test rebalancing
    let rebalance_result = loader.rebalance_and_reload();
    assert!(rebalance_result.is_ok());

    // After rebalancing, healthy workers should have more shards
    let stats = loader.get_comprehensive_statistics();
    assert_eq!(stats.distributed_stats.active_workers, 2); // Only 2 healthy workers
    assert!(stats.distributed_stats.rebalancing_events > 0);
}

#[test]
fn test_performance_metrics_and_statistics() {
    let temp_dir = TempDir::new().unwrap();
    let file_paths = create_structured_test_dataset(temp_dir.path(), 10, 100);

    let config = IntegratedLoaderConfig {
        mmap_config: LoaderConfig {
            page_size: 2048,
            prefetch_strategy: PrefetchStrategy::Sequential { window_size: 3 },
            use_shared_memory: true,
            enable_lazy_loading: true,
        },
        sharding_config: ShardingConfig {
            num_shards: 5,
            sharding_strategy: ShardingStrategy::HashBased,
            rebalancing_strategy: RebalancingStrategy::Static,
            fault_tolerance: FaultToleranceConfig::default(),
        },
        enable_cross_shard_prefetch: true,
        batch_size: 200,
    };

    let workers = vec![
        WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
        WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
    ];

    let mut loader = IntegratedDataLoader::new(config, workers);
    loader.initialize(&file_paths).unwrap();

    let device = Device::cuda(0).unwrap_or(Device::default());

    // Perform multiple loads to generate statistics
    for batch_idx in 0..3 {
        let _batch = loader.load_distributed_batch(batch_idx, &device).unwrap();
    }

    // Report some performance metrics
    loader
        .distributed_loader()
        .report_worker_performance(0, 100);
    loader
        .distributed_loader()
        .report_worker_performance(1, 150);

    // Get comprehensive statistics
    let stats = loader.get_comprehensive_statistics();

    // Verify statistics
    assert_eq!(stats.distributed_stats.total_shards, 5);
    assert_eq!(stats.distributed_stats.total_files, 10);
    assert_eq!(stats.distributed_stats.active_workers, 2);
    assert!(stats.distributed_stats.average_worker_performance > 0.0);

    // Memory-mapped statistics
    assert!(stats.mmap_stats.total_reads > 0);
    assert!(stats.active_mmap_loaders > 0);

    // Test derived metrics
    let cache_hit_ratio = stats.cache_hit_ratio();
    assert!(cache_hit_ratio >= 0.0 && cache_hit_ratio <= 1.0);

    let efficiency_score = stats.efficiency_score();
    assert!(efficiency_score >= 0.0 && efficiency_score <= 1.0);

    // With no failures, efficiency should be high
    assert!(efficiency_score > 0.5);
}

#[test]
fn test_concurrent_access_with_transformers() {
    use std::sync::Arc;
    use std::thread;

    let temp_dir = TempDir::new().unwrap();
    let file_paths = create_structured_test_dataset(temp_dir.path(), 12, 50);

    let config = IntegratedLoaderConfig {
        sharding_config: ShardingConfig {
            num_shards: 6,
            sharding_strategy: ShardingStrategy::RoundRobin,
            rebalancing_strategy: RebalancingStrategy::Static,
            fault_tolerance: FaultToleranceConfig::default(),
        },
        batch_size: 25,
        ..Default::default()
    };

    let workers = vec![
        WorkerInfo::new(0, "worker0".to_string(), "127.0.0.1:8000".to_string()),
        WorkerInfo::new(1, "worker1".to_string(), "127.0.0.1:8001".to_string()),
        WorkerInfo::new(2, "worker2".to_string(), "127.0.0.1:8002".to_string()),
    ];

    let mut loader = IntegratedDataLoader::new(config, workers.clone());
    loader.initialize(&file_paths).unwrap();
    let loader = Arc::new(loader);

    let device = Device::cuda(0).unwrap_or(Device::default());
    let mut handles = vec![];

    // Spawn threads for concurrent data loading and processing
    for worker_id in 0..3 {
        let loader_clone = Arc::clone(&loader);
        let device_clone = device.clone();

        let handle = thread::spawn(move || {
            // Each thread loads data for its assigned worker and applies transformations
            let batch = loader_clone
                .load_worker_batch(worker_id, 0, &device_clone)
                .unwrap();

            let mut processed_count = 0;
            for tensor in batch {
                // Apply different transformers based on worker_id
                match worker_id {
                    0 => {
                        let mut scaler = StandardScaler::new();
                        scaler.fit(&tensor).unwrap();
                        let _transformed = scaler.transform(&tensor).unwrap();
                        processed_count += 1;
                    }
                    1 => {
                        let mut scaler = MinMaxScaler::new();
                        scaler.fit(&tensor).unwrap();
                        let _transformed = scaler.transform(&tensor).unwrap();
                        processed_count += 1;
                    }
                    2 => {
                        // Worker 2 uses both transformers in sequence
                        let mut std_scaler = StandardScaler::new();
                        let mut minmax_scaler = MinMaxScaler::new();

                        std_scaler.fit(&tensor).unwrap();
                        let std_transformed = std_scaler.transform(&tensor).unwrap();

                        minmax_scaler.fit(&std_transformed).unwrap();
                        let _final_transformed = minmax_scaler.transform(&std_transformed).unwrap();
                        processed_count += 1;
                    }
                    _ => unreachable!(),
                }
            }

            processed_count
        });

        handles.push(handle);
    }

    // Wait for all threads and collect results
    let mut total_processed = 0;
    for handle in handles {
        let count = handle.join().unwrap();
        total_processed += count;
    }

    // Verify that data was processed
    assert!(total_processed > 0);

    // Verify loader state is still consistent
    let final_stats = loader.get_comprehensive_statistics();
    assert_eq!(final_stats.distributed_stats.total_files, 12);
    assert_eq!(final_stats.distributed_stats.active_workers, 3);
}