//! Data to ML Pipeline Integration Tests
//!
//! This module tests the end-to-end flow from market data ingestion
//! to ML model prediction generation in the Foxhunt HFT system.
//!
//! Tests validate:
//! - Market data ingestion and preprocessing
//! - Feature engineering and data transformation
//! - ML model inference pipeline
//! - Data quality validation
//! - Performance and latency requirements

use anyhow::{Context, Result};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::{mpsc, RwLock, Mutex};
use tokio::time::{sleep, timeout};
use tracing::{debug, info, warn, error, instrument};
use uuid::Uuid;

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use rust_decimal::Decimal;

// Core types

/// Market data sample for testing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MarketDataSample {
    pub symbol: String,
    pub price: Decimal,
    pub volume: u64,
    pub bid: Decimal,
    pub ask: Decimal,
    pub timestamp: DateTime<Utc>,
    pub exchange: String,
}

/// ML prediction result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MLPrediction {
    pub symbol: String,
    pub signal_type: SignalType,
    pub confidence: f64,
    pub price_target: Option<Decimal>,
    pub time_horizon: Duration,
    pub timestamp: DateTime<Utc>,
    pub model_version: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SignalType {
    Buy,
    Sell,
    Hold,
    StrongBuy,
    StrongSell,
}

/// Feature vector for ML processing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureVector {
    pub symbol: String,
    pub features: Vec<f64>,
    pub feature_names: Vec<String>,
    pub timestamp: DateTime<Utc>,
    pub window_size: usize,
}

/// Data quality metrics
#[derive(Debug, Clone)]
pub struct DataQualityMetrics {
    pub completeness: f64,
    pub timeliness: f64,
    pub accuracy: f64,
    pub consistency: f64,
    pub total_samples: usize,
    pub invalid_samples: usize,
    pub processing_latency: Duration,
}

/// Mock market data service for testing
#[derive(Clone)]
pub struct MockMarketDataService {
    pub samples: Arc<RwLock<Vec<MarketDataSample>>>,
    pub is_running: Arc<RwLock<bool>>,
}

impl MockMarketDataService {
    pub fn new() -> Self {
        Self {
            samples: Arc::new(RwLock::new(Vec::new())),
            is_running: Arc::new(RwLock::new(false)),
        }
    }

    pub async fn start(&self) -> Result<()> {
        let mut is_running = self.is_running.write().await;
        *is_running = true;
        info!("Mock market data service started");
        Ok(())
    }

    pub async fn stop(&self) -> Result<()> {
        let mut is_running = self.is_running.write().await;
        *is_running = false;
        info!("Mock market data service stopped");
        Ok(())
    }

    pub async fn publish_sample(&self, sample: MarketDataSample) -> Result<()> {
        let mut samples = self.samples.write().await;
        samples.push(sample);
        debug!("Published market data sample");
        Ok(())
    }

    pub async fn fetch_latest_data(&self, symbol: &str, limit: usize) -> Result<Vec<MarketDataSample>> {
        let samples = self.samples.read().await;
        let filtered: Vec<MarketDataSample> = samples
            .iter()
            .filter(|s| s.symbol == symbol)
            .take(limit)
            .cloned()
            .collect();
        Ok(filtered)
    }
}

/// Mock ML service for testing
#[derive(Clone)]
pub struct MockMLService {
    pub predictions: Arc<RwLock<Vec<MLPrediction>>>,
    pub model_version: String,
    pub processing_latency: Duration,
}

impl MockMLService {
    pub fn new() -> Self {
        Self {
            predictions: Arc::new(RwLock::new(Vec::new())),
            model_version: "test-model-v1.0".to_owned(),
            processing_latency: Duration::from_millis(10),
        }
    }

    pub async fn preprocess_data(&self, samples: &[MarketDataSample]) -> Result<FeatureVector> {
        let start = Instant::now();
        
        // Simulate feature engineering
        let mut features = Vec::new();
        if !samples.is_empty() {
            let latest = &samples[0];
            
            // Price-based features
            features.push(latest.price.to_f64().unwrap_or(0.0));
            features.push(latest.volume as f64);
            features.push((latest.ask - latest.bid).to_f64().unwrap_or(0.0)); // spread
            
            // Technical indicators (simplified)
            if samples.len() >= 5 {
                let prices: Vec<f64> = samples.iter()
                    .map(|s| s.price.to_f64().unwrap_or(0.0))
                    .collect();
                
                // Moving average
                let ma = prices.iter().sum::<f64>() / prices.len() as f64;
                features.push(ma);
                
                // Price momentum
                let momentum = if prices.len() >= 2 {
                    prices[0] - prices[prices.len()-1]
                } else {
                    0.0
                };
                features.push(momentum);
            }
        }

        let feature_names = vec![
            "price".to_owned(),
            "volume".to_owned(),
            "spread".to_owned(),
            "moving_average".to_owned(),
            "momentum".to_owned(),
        ];

        sleep(self.processing_latency).await;

        Ok(FeatureVector {
            symbol: samples.first().map(|s| s.symbol.clone()).unwrap_or_default(),
            features,
            feature_names,
            timestamp: Utc::now(),
            window_size: samples.len(),
        })
    }

    pub async fn predict(&self, feature_vector: &FeatureVector) -> Result<MLPrediction> {
        let start = Instant::now();
        
        // Simulate ML inference
        let signal_type = if !feature_vector.features.is_empty() {
            let price = feature_vector.features[0];
            let momentum = feature_vector.features.get(4).copied().unwrap_or(0.0);
            
            match (momentum > 0.0, price > 100.0) {
                (true, true) => SignalType::Buy,
                (true, false) => SignalType::StrongBuy,
                (false, true) => SignalType::Sell,
                (false, false) => SignalType::Hold,
            }
        } else {
            SignalType::Hold
        };

        let confidence = fastrand::f64() * 0.3 + 0.7; // 0.7-1.0 range
        
        sleep(self.processing_latency).await;

        let prediction = MLPrediction {
            symbol: feature_vector.symbol.clone(),
            signal_type,
            confidence,
            price_target: None, // Could be derived from model
            time_horizon: Duration::from_secs(300), // 5 minutes
            timestamp: Utc::now(),
            model_version: self.model_version.clone(),
        };

        let mut predictions = self.predictions.write().await;
        predictions.push(prediction.clone());

        Ok(prediction)
    }
}

/// Data quality validator
pub struct DataQualityValidator;

impl DataQualityValidator {
    pub fn validate_market_data(&self, samples: &[MarketDataSample]) -> DataQualityMetrics {
        let start = Instant::now();
        
        let total_samples = samples.len();
        let mut invalid_samples = 0;
        
        for sample in samples {
            // Check for invalid prices
            if sample.price <= Decimal::ZERO || sample.bid <= Decimal::ZERO || sample.ask <= Decimal::ZERO {
                invalid_samples += 1;
                continue;
            }
            
            // Check for invalid spread (ask should be >= bid)
            if sample.ask < sample.bid {
                invalid_samples += 1;
                continue;
            }
            
            // Check for stale data (older than 5 minutes)
            let age = Utc::now().signed_duration_since(sample.timestamp);
            if age.num_minutes() > 5 {
                invalid_samples += 1;
                continue;
            }
        }
        
        let valid_samples = total_samples - invalid_samples;
        let completeness = if total_samples > 0 {
            valid_samples as f64 / total_samples as f64
        } else {
            0.0
        };
        
        DataQualityMetrics {
            completeness,
            timeliness: 0.95, // Simulated
            accuracy: 0.98,   // Simulated
            consistency: 0.97, // Simulated
            total_samples,
            invalid_samples,
            processing_latency: start.elapsed(),
        }
    }
}

/// Test configuration for symbols and their base prices
struct TestSymbolConfig {
    symbol: &'static str,
    base_price: f64,
    exchange: &'static str,
}

const TEST_SYMBOLS: &[TestSymbolConfig] = &[
    TestSymbolConfig { symbol: "SYM1", base_price: 100.0, exchange: "NASDAQ" },
    TestSymbolConfig { symbol: "SYM2", base_price: 200.0, exchange: "NYSE" },
    TestSymbolConfig { symbol: "SYM3", base_price: 50.0, exchange: "NASDAQ" },
    TestSymbolConfig { symbol: "SYM4", base_price: 300.0, exchange: "NYSE" },
    TestSymbolConfig { symbol: "SYM5", base_price: 150.0, exchange: "NASDAQ" },
];

/// Get test symbol configuration by index
fn get_test_symbol_config(index: usize) -> &'static TestSymbolConfig {
    &TEST_SYMBOLS[index % TEST_SYMBOLS.len()]
}

/// Generate sample market data for testing
fn generate_sample_market_data(symbol: &str, count: usize) -> Vec<MarketDataSample> {
    let mut samples = Vec::new();

    // Find matching test config or use default
    let config = TEST_SYMBOLS.iter()
        .find(|c| c.symbol == symbol)
        .unwrap_or(&TEST_SYMBOLS[0]);

    let base_price = config.base_price;

    for i in 0..count {
        let price_offset = (fastrand::f64() - 0.5) * (base_price * 0.1); // ±10% variation
        let price = Decimal::try_from(base_price + price_offset).unwrap_or(Decimal::new(base_price as i64, 0));
        let spread = Decimal::try_from(fastrand::f64() * 0.1 + 0.01).unwrap_or(Decimal::new(1, 2));

        let sample = MarketDataSample {
            symbol: symbol.to_string(),
            price,
            volume: fastrand::u64(1000..50000),
            bid: price - spread / Decimal::new(2, 0),
            ask: price + spread / Decimal::new(2, 0),
            timestamp: Utc::now() - Duration::seconds(i as i64),
            exchange: config.exchange.to_string(),
        };

        samples.push(sample);
    }

    samples
}

#[tokio::test]
async fn test_market_data_ingestion_pipeline() -> Result<()> {
    tracing_subscriber::fmt::init();
    info!("Starting market data ingestion pipeline test");
    
    let market_data_service = MockMarketDataService::new();
    let ml_service = MockMLService::new();
    let validator = DataQualityValidator;
    
    // Start services
    market_data_service.start().await?;
    
    // Generate and publish test data
    let test_config = get_test_symbol_config(0);
    let test_symbol = test_config.symbol;
    let samples = generate_sample_market_data(test_symbol, 100);
    
    for sample in &samples {
        market_data_service.publish_sample(sample.clone()).await?;
    }
    
    // Fetch data from market data service
    let fetched_data = market_data_service.fetch_latest_data(test_symbol, 20).await?;
    assert!(!fetched_data.is_empty(), "Should fetch market data");
    assert_eq!(fetched_data[0].symbol, test_symbol);
    
    info!("Market data ingestion: {} samples", fetched_data.len());
    
    // Validate data quality
    let quality_metrics = validator.validate_market_data(&fetched_data);
    assert!(quality_metrics.completeness > 0.8, "Data completeness should be > 80%");
    assert!(quality_metrics.processing_latency < Duration::from_millis(100));
    
    info!("Data quality validation passed: {:.2}% completeness", quality_metrics.completeness * 100.0);
    
    market_data_service.stop().await?;
    Ok(())
}

#[tokio::test]
async fn test_ml_feature_engineering() -> Result<()> {
    tracing_subscriber::fmt::init();
    info!("Starting ML feature engineering test");
    
    let ml_service = MockMLService::new();
    
    // Generate test data
    let test_config = get_test_symbol_config(1);
    let samples = generate_sample_market_data(test_config.symbol, 50);
    
    // Test feature engineering
    let start = Instant::now();
    let feature_vector = ml_service.preprocess_data(&samples).await?;
    let processing_time = start.elapsed();
    
    // Validate feature vector
    assert!(!feature_vector.features.is_empty(), "Features should not be empty");
    assert_eq!(feature_vector.features.len(), feature_vector.feature_names.len());
    assert_eq!(feature_vector.symbol, test_config.symbol);
    assert!(processing_time < Duration::from_millis(100), "Feature engineering should be fast");
    
    info!("Feature engineering completed: {} features in {:?}", 
          feature_vector.features.len(), processing_time);
    
    Ok(())
}

#[tokio::test]
async fn test_ml_prediction_generation() -> Result<()> {
    tracing_subscriber::fmt::init();
    info!("Starting ML prediction generation test");
    
    let ml_service = MockMLService::new();
    
    // Generate test data and features
    let test_config = get_test_symbol_config(2);
    let samples = generate_sample_market_data(test_config.symbol, 30);
    let feature_vector = ml_service.preprocess_data(&samples).await?;
    
    // Generate prediction
    let start = Instant::now();
    let prediction = ml_service.predict(&feature_vector).await?;
    let inference_time = start.elapsed();
    
    // Validate prediction
    assert_eq!(prediction.symbol, test_config.symbol);
    assert!(prediction.confidence >= 0.0 && prediction.confidence <= 1.0);
    assert!(inference_time < Duration::from_millis(50), "ML inference should be fast");
    assert!(!prediction.model_version.is_empty());
    
    info!("ML prediction generated: {:?} with {:.2}% confidence in {:?}", 
          prediction.signal_type, prediction.confidence * 100.0, inference_time);
    
    Ok(())
}

#[tokio::test]
async fn test_full_data_to_ml_pipeline() -> Result<()> {
    tracing_subscriber::fmt::init();
    info!("Starting full data-to-ML pipeline test");
    
    let market_data_service = MockMarketDataService::new();
    let ml_service = MockMLService::new();
    let validator = DataQualityValidator;
    
    // Start services
    market_data_service.start().await?;
    
    let test_config = get_test_symbol_config(3);
    let test_symbol = test_config.symbol;
    let pipeline_start = Instant::now();
    
    // Step 1: Market data ingestion
    let samples = generate_sample_market_data(test_symbol, 200);
    for sample in &samples {
        market_data_service.publish_sample(sample.clone()).await?;
    }
    
    // Step 2: Data fetching and validation
    let fetched_data = market_data_service.fetch_latest_data(test_symbol, 50).await?;
    let quality_metrics = validator.validate_market_data(&fetched_data);
    
    assert!(quality_metrics.completeness > 0.8);
    
    // Step 3: Feature engineering
    let feature_vector = ml_service.preprocess_data(&fetched_data).await?;
    
    // Step 4: ML prediction
    let prediction = ml_service.predict(&feature_vector).await?;
    
    let total_pipeline_time = pipeline_start.elapsed();
    
    // Validate end-to-end pipeline
    assert_eq!(prediction.symbol, test_symbol);
    assert!(total_pipeline_time < Duration::from_millis(500), 
            "Full pipeline should complete in <500ms");
    
    info!("Full pipeline completed in {:?}: {} samples -> {} features -> {:?} signal", 
          total_pipeline_time, 
          fetched_data.len(), 
          feature_vector.features.len(), 
          prediction.signal_type);
    
    market_data_service.stop().await?;
    Ok(())
}

#[tokio::test]
async fn test_pipeline_performance_under_load() -> Result<()> {
    tracing_subscriber::fmt::init();
    info!("Starting pipeline performance test under load");
    
    let market_data_service = MockMarketDataService::new();
    let ml_service = MockMLService::new();
    
    market_data_service.start().await?;
    
    let symbols: Vec<&str> = TEST_SYMBOLS.iter().map(|c| c.symbol).collect();
    let mut predictions = Vec::new();
    let start = Instant::now();
    
    // Process multiple symbols concurrently
    let handles: Vec<_> = symbols.into_iter().map(|symbol| {
        let market_data_service = market_data_service.clone();
        let ml_service = ml_service.clone();
        
        tokio::spawn(async move {
            // Generate and process data for each symbol
            let samples = generate_sample_market_data(symbol, 100);
            for sample in &samples {
                market_data_service.publish_sample(sample.clone()).await.ok();
            }
            
            let fetched_data = market_data_service.fetch_latest_data(symbol, 30).await?;
            let feature_vector = ml_service.preprocess_data(&fetched_data).await?;
            let prediction = ml_service.predict(&feature_vector).await?;
            
            Ok::<MLPrediction, anyhow::Error>(prediction)
        })
    }).collect();
    
    // Wait for all concurrent processing to complete
    for handle in handles {
        let prediction = handle.await??;
        predictions.push(prediction);
    }
    
    let total_time = start.elapsed();
    let throughput = predictions.len() as f64 / total_time.as_secs_f64();
    
    // Validate performance
    assert_eq!(predictions.len(), 5);
    assert!(total_time < Duration::from_secs(2), "Should process 5 symbols in <2s");
    assert!(throughput >= 2.0, "Should achieve >2 predictions/second");
    
    info!("Performance test completed: {} predictions in {:?} ({:.2} predictions/sec)", 
          predictions.len(), total_time, throughput);
    
    market_data_service.stop().await?;
    Ok(())
}

#[tokio::test]
async fn test_data_quality_validation() -> Result<()> {
    tracing_subscriber::fmt::init();
    info!("Starting data quality validation test");
    
    let validator = DataQualityValidator;
    
    // Test with good data
    let test_config = get_test_symbol_config(0);
    let good_samples = generate_sample_market_data(test_config.symbol, 100);
    let good_metrics = validator.validate_market_data(&good_samples);
    
    assert!(good_metrics.completeness > 0.95);
    assert_eq!(good_metrics.total_samples, 100);
    assert_eq!(good_metrics.invalid_samples, 0);
    
    // Test with some bad data
    let test_config = get_test_symbol_config(0);
    let mut mixed_samples = generate_sample_market_data(test_config.symbol, 50);
    
    // Add some invalid samples
    mixed_samples.push(MarketDataSample {
        symbol: test_config.symbol.to_string(),
        price: Decimal::ZERO, // Invalid price
        volume: 1000,
        bid: Decimal::new(100, 0),
        ask: Decimal::new(99, 0), // Invalid spread (ask < bid)
        timestamp: Utc::now(),
        exchange: test_config.exchange.to_string(),
    });
    
    mixed_samples.push(MarketDataSample {
        symbol: test_config.symbol.to_string(),
        price: Decimal::new(150, 0),
        volume: 1000,
        bid: Decimal::new(150, 0),
        ask: Decimal::new(149, 0), // Invalid spread
        timestamp: Utc::now() - Duration::minutes(10), // Stale data
        exchange: test_config.exchange.to_string(),
    });
    
    let mixed_metrics = validator.validate_market_data(&mixed_samples);
    
    assert!(mixed_metrics.completeness < 1.0);
    assert!(mixed_metrics.invalid_samples > 0);
    assert_eq!(mixed_metrics.total_samples, 52);
    
    info!("Data quality validation: {:.2}% completeness with {} invalid samples", 
          mixed_metrics.completeness * 100.0, mixed_metrics.invalid_samples);
    
    Ok(())
}