foxhunt/testing/integration/chaos/ml_training_chaos.rs

//! ML Training Service Chaos Engineering Tests
//!
//! Specialized chaos tests for MLTrainingService resilience, checkpoint recovery,
//! and training process continuity under various failure conditions.

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use std::time::Duration;
use tokio::time::sleep;
use tracing::info;
use uuid::Uuid;

use super::chaos_framework::{ChaosExperiment, ChaosOrchestrator, FailureType, Signal};

/// ML-specific chaos test configurations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MLChaosConfig {
    pub ml_service_endpoint: String,
    pub checkpoint_base_path: PathBuf,
    pub model_types: Vec<ModelType>,
    pub training_timeout_secs: u64,
    pub max_recovery_time_ms: u64, // HFT requirement: sub-100ms
    pub gpu_memory_threshold_mb: u64,
}

impl Default for MLChaosConfig {
    fn default() -> Self {
        Self {
            ml_service_endpoint: "http://localhost:8080".to_string(),
            checkpoint_base_path: PathBuf::from("/tmp/test_checkpoints"),
            model_types: vec![ModelType::TLOB],
            training_timeout_secs: 300,
            max_recovery_time_ms: 100,
            gpu_memory_threshold_mb: 8192,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ModelType {
    TLOB,
    MAMBA2,
    DQN,
    PPO,
    Liquid,
    TFT,
}

impl ModelType {
    pub fn as_str(&self) -> &'static str {
        match self {
            ModelType::TLOB => "tlob",
            ModelType::MAMBA2 => "mamba2",
            ModelType::DQN => "dqn",
            ModelType::PPO => "ppo",
            ModelType::Liquid => "liquid",
            ModelType::TFT => "tft",
        }
    }

    pub fn typical_checkpoint_interval_secs(&self) -> u64 {
        match self {
            ModelType::TLOB => 30,   // Fast checkpointing for TLOB
            ModelType::MAMBA2 => 60, // MAMBA-2 SSM checkpointing
            ModelType::DQN => 120,   // DQN experience replay checkpoints
            ModelType::PPO => 90,    // PPO policy checkpoints
            ModelType::Liquid => 45, // Liquid network state checkpoints
            ModelType::TFT => 180,   // TFT transformer checkpoints
        }
    }

    pub fn expected_recovery_time_ms(&self) -> u64 {
        match self {
            ModelType::TLOB => 25,   // Ultra-fast TLOB recovery
            ModelType::MAMBA2 => 40, // MAMBA-2 state recovery
            ModelType::DQN => 80,    // DQN replay buffer recovery
            ModelType::PPO => 60,    // PPO policy recovery
            ModelType::Liquid => 35, // Liquid network recovery
            ModelType::TFT => 95,    // TFT transformer recovery
        }
    }
}

/// ML Training Chaos Test Results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MLChaosResult {
    pub experiment_id: Uuid,
    pub model_type: ModelType,
    pub training_job_id: Option<String>,
    pub checkpoint_before_failure: Option<CheckpointMetadata>,
    pub checkpoint_after_recovery: Option<CheckpointMetadata>,
    pub model_accuracy_before: Option<f64>,
    pub model_accuracy_after: Option<f64>,
    pub training_loss_continuity: bool,
    pub gpu_memory_recovery: Option<GpuMemoryMetrics>,
    pub performance_regression: Option<MLPerformanceRegression>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointMetadata {
    pub path: PathBuf,
    pub file_size_bytes: u64,
    pub created_at: std::time::SystemTime,
    pub model_epoch: u32,
    pub training_step: u64,
    pub loss_value: Option<f64>,
    pub checksum: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuMemoryMetrics {
    pub allocated_mb: u64,
    pub reserved_mb: u64,
    pub free_mb: u64,
    pub utilization_percent: f32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MLPerformanceRegression {
    pub inference_latency_before_ns: u64,
    pub inference_latency_after_ns: u64,
    pub training_throughput_before_samples_sec: f64,
    pub training_throughput_after_samples_sec: f64,
    pub memory_usage_increase_mb: i64,
}

/// ML Training Chaos Test Suite
pub struct MLTrainingChaosTests {
    config: MLChaosConfig,
    orchestrator: ChaosOrchestrator,
}

impl MLTrainingChaosTests {
    pub fn new(config: MLChaosConfig) -> Self {
        Self {
            config,
            orchestrator: ChaosOrchestrator::new(2), // Limit concurrent ML chaos tests
        }
    }

    /// Initialize all ML chaos experiments
    pub async fn initialize_experiments(&self) -> Result<Vec<Uuid>> {
        let mut experiment_ids = Vec::new();

        for model_type in &self.config.model_types {
            // 1. Process Kill/Restart Test
            let kill_experiment_id = self.create_process_kill_experiment(model_type).await?;
            experiment_ids.push(kill_experiment_id);

            // 2. Memory Pressure Test
            let memory_experiment_id = self.create_memory_pressure_experiment(model_type).await?;
            experiment_ids.push(memory_experiment_id);

            // 3. GPU Resource Exhaustion Test
            let gpu_experiment_id = self.create_gpu_exhaustion_experiment(model_type).await?;
            experiment_ids.push(gpu_experiment_id);

            // 4. Network Partition Test
            let network_experiment_id =
                self.create_network_partition_experiment(model_type).await?;
            experiment_ids.push(network_experiment_id);

            // 5. Disk I/O Failure Test
            let disk_experiment_id = self.create_disk_failure_experiment(model_type).await?;
            experiment_ids.push(disk_experiment_id);
        }

        info!("Initialized {} ML chaos experiments", experiment_ids.len());
        Ok(experiment_ids)
    }

    /// Create process kill/restart experiment for specific model type
    async fn create_process_kill_experiment(&self, model_type: &ModelType) -> Result<Uuid> {
        let experiment_id = Uuid::new_v4();

        let experiment = ChaosExperiment {
            id: experiment_id,
            name: format!(
                "MLTrainingService {} Process Kill Test",
                model_type.as_str()
            ),
            description: format!(
                "Test {} model training recovery from process termination",
                model_type.as_str()
            ),
            target_service: "ml_training_service".to_string(),
            failure_type: FailureType::ProcessKill {
                signal: Signal::SIGTERM,       // Graceful termination first
                delay_before_restart_ms: 2000, // 2 second delay
            },
            duration: Duration::from_secs(5), // Quick failure
            recovery_timeout: Duration::from_secs(30),
            max_recovery_time_ms: model_type.expected_recovery_time_ms(),
            enabled: true,
        };

        self.orchestrator.register_experiment(experiment).await?;
        Ok(experiment_id)
    }

    /// Create memory pressure experiment
    async fn create_memory_pressure_experiment(&self, model_type: &ModelType) -> Result<Uuid> {
        let experiment_id = Uuid::new_v4();

        let experiment = ChaosExperiment {
            id: experiment_id,
            name: format!(
                "MLTrainingService {} Memory Pressure Test",
                model_type.as_str()
            ),
            description: format!(
                "Test {} model training under memory pressure conditions",
                model_type.as_str()
            ),
            target_service: "ml_training_service".to_string(),
            failure_type: FailureType::MemoryPressure {
                target_mb: 4096,    // 4GB memory pressure
                duration_ms: 30000, // 30 seconds
            },
            duration: Duration::from_secs(35),
            recovery_timeout: Duration::from_secs(45),
            max_recovery_time_ms: 100, // HFT requirement
            enabled: true,
        };

        self.orchestrator.register_experiment(experiment).await?;
        Ok(experiment_id)
    }

    /// Create GPU resource exhaustion experiment
    async fn create_gpu_exhaustion_experiment(&self, model_type: &ModelType) -> Result<Uuid> {
        let experiment_id = Uuid::new_v4();

        let experiment = ChaosExperiment {
            id: experiment_id,
            name: format!(
                "MLTrainingService {} GPU Exhaustion Test",
                model_type.as_str()
            ),
            description: format!(
                "Test {} model training recovery from GPU resource exhaustion",
                model_type.as_str()
            ),
            target_service: "ml_training_service".to_string(),
            failure_type: FailureType::GpuResourceExhaustion {
                memory_fill_percent: 95, // Fill 95% of GPU memory
                duration_ms: 20000,      // 20 seconds
            },
            duration: Duration::from_secs(25),
            recovery_timeout: Duration::from_secs(60),
            max_recovery_time_ms: 150, // GPU recovery can be slower
            enabled: true,
        };

        self.orchestrator.register_experiment(experiment).await?;
        Ok(experiment_id)
    }

    /// Create network partition experiment
    async fn create_network_partition_experiment(&self, model_type: &ModelType) -> Result<Uuid> {
        let experiment_id = Uuid::new_v4();

        let experiment = ChaosExperiment {
            id: experiment_id,
            name: format!(
                "MLTrainingService {} Network Partition Test",
                model_type.as_str()
            ),
            description: format!(
                "Test {} model training resilience to network partitions",
                model_type.as_str()
            ),
            target_service: "ml_training_service".to_string(),
            failure_type: FailureType::NetworkPartition {
                target_ports: vec![8080, 5432, 6379], // gRPC, PostgreSQL, Redis
                duration_ms: 15000,                   // 15 seconds
            },
            duration: Duration::from_secs(20),
            recovery_timeout: Duration::from_secs(30),
            max_recovery_time_ms: 80,
            enabled: true,
        };

        self.orchestrator.register_experiment(experiment).await?;
        Ok(experiment_id)
    }

    /// Create disk I/O failure experiment
    async fn create_disk_failure_experiment(&self, model_type: &ModelType) -> Result<Uuid> {
        let experiment_id = Uuid::new_v4();

        let experiment = ChaosExperiment {
            id: experiment_id,
            name: format!(
                "MLTrainingService {} Disk I/O Failure Test",
                model_type.as_str()
            ),
            description: format!(
                "Test {} model training resilience to disk I/O failures",
                model_type.as_str()
            ),
            target_service: "ml_training_service".to_string(),
            failure_type: FailureType::DiskIoFailure {
                target_paths: vec![
                    self.config
                        .checkpoint_base_path
                        .to_string_lossy()
                        .to_string(),
                    "/tmp".to_string(),
                    "/var/log".to_string(),
                ],
                failure_rate_percent: 30, // 30% I/O failure rate
            },
            duration: Duration::from_secs(25),
            recovery_timeout: Duration::from_secs(40),
            max_recovery_time_ms: 120,
            enabled: true,
        };

        self.orchestrator.register_experiment(experiment).await?;
        Ok(experiment_id)
    }

    /// Execute comprehensive ML training chaos test suite
    pub async fn run_ml_chaos_suite(&self) -> Result<Vec<MLChaosResult>> {
        info!("Starting ML Training Chaos Test Suite");

        let experiment_ids = self.initialize_experiments().await?;
        let mut ml_results = Vec::new();

        for experiment_id in experiment_ids {
            info!("Executing ML chaos experiment: {}", experiment_id);

            // Start a training job for the experiment
            let training_job_id = self
                .start_training_job_for_experiment(experiment_id)
                .await?;

            // Capture pre-failure state
            let pre_failure_state = self.capture_ml_state(&training_job_id).await?;

            // Execute the chaos experiment
            let chaos_result = self.orchestrator.execute_experiment(experiment_id).await?;

            // Capture post-recovery state
            let post_recovery_state = self.capture_ml_state(&training_job_id).await?;

            // Validate checkpoint integrity and model continuity
            let checkpoint_valid = self
                .validate_model_checkpoint_integrity(&pre_failure_state, &post_recovery_state)
                .await?;

            // Create ML-specific result
            let ml_result = self
                .create_ml_result(
                    experiment_id,
                    training_job_id,
                    pre_failure_state,
                    post_recovery_state,
                    checkpoint_valid,
                )
                .await?;

            ml_results.push(ml_result);
        }

        info!(
            "ML Chaos Test Suite completed with {} results",
            ml_results.len()
        );
        Ok(ml_results)
    }

    /// Start a training job for chaos experiment
    async fn start_training_job_for_experiment(&self, experiment_id: Uuid) -> Result<String> {
        // TODO: Implement gRPC call to MLTrainingService to start training
        // This would call the StartTraining endpoint with appropriate model config

        let training_job_id = format!("chaos_training_{}", experiment_id);
        info!("Started training job: {}", training_job_id);

        // Wait for training to begin
        sleep(Duration::from_secs(5)).await;

        Ok(training_job_id)
    }

    /// Capture ML service state before/after chaos
    async fn capture_ml_state(&self, training_job_id: &str) -> Result<MLServiceState> {
        // TODO: Implement state capture via gRPC calls:
        // - GetTrainingJobDetails
        // - Get current checkpoint info
        // - Capture GPU metrics
        // - Capture performance metrics

        Ok(MLServiceState {
            training_job_id: training_job_id.to_string(),
            current_epoch: 42,
            training_step: 1000,
            current_loss: Some(0.125),
            checkpoint_metadata: None,
            gpu_metrics: None,
            inference_latency_ns: 25000, // 25μs
        })
    }

    /// Validate checkpoint integrity after recovery
    async fn validate_model_checkpoint_integrity(
        &self,
        pre_state: &MLServiceState,
        post_state: &MLServiceState,
    ) -> Result<bool> {
        // Validate that:
        // 1. Training can resume from checkpoint
        // 2. Model accuracy hasn't degraded significantly
        // 3. Training loss continuity is maintained
        // 4. No corruption in model weights

        let training_continuity = post_state.training_step >= pre_state.training_step;
        let loss_reasonable = match (pre_state.current_loss, post_state.current_loss) {
            (Some(pre), Some(post)) => (post - pre).abs() < 0.1, // Loss shouldn't jump
            _ => true,                                           // No loss data to compare
        };

        Ok(training_continuity && loss_reasonable)
    }

    /// Create ML-specific chaos result
    async fn create_ml_result(
        &self,
        experiment_id: Uuid,
        training_job_id: String,
        pre_state: MLServiceState,
        post_state: MLServiceState,
        checkpoint_valid: bool,
    ) -> Result<MLChaosResult> {
        let performance_regression =
            if post_state.inference_latency_ns > pre_state.inference_latency_ns {
                Some(MLPerformanceRegression {
                    inference_latency_before_ns: pre_state.inference_latency_ns,
                    inference_latency_after_ns: post_state.inference_latency_ns,
                    training_throughput_before_samples_sec: 1000.0, // Placeholder
                    training_throughput_after_samples_sec: 950.0,   // Placeholder
                    memory_usage_increase_mb: 50,                   // Placeholder
                })
            } else {
                None
            };

        Ok(MLChaosResult {
            experiment_id,
            model_type: ModelType::TLOB, // TODO: Extract from experiment
            training_job_id: Some(training_job_id),
            checkpoint_before_failure: pre_state.checkpoint_metadata,
            checkpoint_after_recovery: post_state.checkpoint_metadata,
            model_accuracy_before: None, // TODO: Implement accuracy capture
            model_accuracy_after: None,  // TODO: Implement accuracy capture
            training_loss_continuity: checkpoint_valid,
            gpu_memory_recovery: post_state.gpu_metrics,
            performance_regression,
        })
    }

    /// Generate chaos test report
    pub async fn generate_chaos_report(&self, results: &[MLChaosResult]) -> Result<String> {
        let mut report = String::new();

        report.push_str("# ML Training Chaos Engineering Report\n\n");
        report.push_str(&format!("**Generated:** {}\n", chrono::Utc::now()));
        report.push_str(&format!("**Total Tests:** {}\n\n", results.len()));

        let successful = results
            .iter()
            .filter(|r| r.training_loss_continuity)
            .count();
        let failed = results.len() - successful;

        report.push_str("## Summary\n");
        report.push_str(&format!("- ✅ **Successful:** {}\n", successful));
        report.push_str(&format!("- ❌ **Failed:** {}\n", failed));
        report.push_str(&format!(
            "- 📊 **Success Rate:** {:.1}%\n\n",
            (successful as f64 / results.len() as f64) * 100.0
        ));

        // Model type breakdown
        let mut model_stats: HashMap<String, (usize, usize)> = HashMap::new();
        for result in results {
            let model = result.model_type.as_str();
            let (success, total) = model_stats.entry(model.to_string()).or_insert((0, 0));
            *total += 1;
            if result.training_loss_continuity {
                *success += 1;
            }
        }

        report.push_str("## Results by Model Type\n");
        for (model, (success, total)) in model_stats {
            let rate = (success as f64 / total as f64) * 100.0;
            report.push_str(&format!(
                "- **{}:** {}/{} ({:.1}%)\n",
                model, success, total, rate
            ));
        }

        // Performance regression analysis
        report.push_str("\n## Performance Analysis\n");
        let regressions: Vec<_> = results
            .iter()
            .filter_map(|r| r.performance_regression.as_ref())
            .collect();

        if !regressions.is_empty() {
            report.push_str(&format!(
                "- **Performance Regressions:** {}\n",
                regressions.len()
            ));
            let avg_latency_increase = regressions
                .iter()
                .map(|r| r.inference_latency_after_ns - r.inference_latency_before_ns)
                .sum::<u64>() as f64
                / regressions.len() as f64;
            report.push_str(&format!(
                "- **Avg Latency Increase:** {:.1}ns\n",
                avg_latency_increase
            ));
        } else {
            report.push_str("- ✅ **No Performance Regressions Detected**\n");
        }

        Ok(report)
    }
}

#[derive(Debug, Clone)]
struct MLServiceState {
    training_job_id: String,
    current_epoch: u32,
    training_step: u64,
    current_loss: Option<f64>,
    checkpoint_metadata: Option<CheckpointMetadata>,
    gpu_metrics: Option<GpuMemoryMetrics>,
    inference_latency_ns: u64,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_ml_chaos_config() {
        let config = MLChaosConfig {
            ml_service_endpoint: "http://localhost:8080".to_string(),
            checkpoint_base_path: PathBuf::from("/tmp/checkpoints"),
            model_types: vec![ModelType::TLOB, ModelType::DQN],
            training_timeout_secs: 300,
            max_recovery_time_ms: 100,
            gpu_memory_threshold_mb: 8192,
        };

        let chaos_tests = MLTrainingChaosTests::new(config);
        assert_eq!(chaos_tests.config.model_types.len(), 2);
    }

    #[test]
    fn test_model_type_properties() {
        assert_eq!(ModelType::TLOB.as_str(), "tlob");
        assert_eq!(ModelType::TLOB.expected_recovery_time_ms(), 25);
        assert_eq!(ModelType::DQN.typical_checkpoint_interval_secs(), 120);
    }
}