foxhunt/tests/chaos/chaos_cli.rs

//! Chaos Engineering CLI Tool
//!
//! Command-line interface for running chaos engineering tests on the Foxhunt HFT system.

use anyhow::{Context, Result};
use clap::{Parser, Subcommand, ValueEnum};
use std::path::PathBuf;
use tracing::{error, info, warn};
use uuid::Uuid;

use super::chaos_framework::ChaosOrchestrator;
use super::ml_training_chaos::{MLChaosConfig, MLTrainingChaosTests, ModelType};
use super::nightly_chaos_runner::{NightlyChaosConfig, NightlyChaosRunner, ChaosJobEvent, AlertSeverity};

/// Chaos Engineering CLI for Foxhunt HFT System
#[derive(Parser)]
#[command(name = "foxhunt-chaos")]
#[command(about = "Chaos engineering tool for Foxhunt HFT trading system")]
#[command(version = "1.0.0")]
pub struct ChaosCliArgs {
    #[command(subcommand)]
    pub command: ChaosCommand,

    /// Enable verbose logging
    #[arg(short, long)]
    pub verbose: bool,

    /// Configuration file path
    #[arg(short, long, value_name = "FILE")]
    pub config: Option<PathBuf>,
}

#[derive(Subcommand)]
pub enum ChaosCommand {
    /// Run a single chaos experiment
    Run {
        /// Type of chaos experiment to run
        #[arg(short, long, value_enum)]
        experiment_type: ExperimentType,

        /// Target service name
        #[arg(short, long, default_value = "ml_training_service")]
        service: String,

        /// Model type for ML experiments
        #[arg(short, long, value_enum)]
        model: Option<ModelTypeArg>,

        /// Maximum recovery time in milliseconds
        #[arg(long, default_value = "100")]
        max_recovery_time_ms: u64,

        /// Experiment duration in seconds
        #[arg(short, long, default_value = "30")]
        duration: u64,

        /// Recovery timeout in seconds
        #[arg(long, default_value = "60")]
        recovery_timeout: u64,
    },

    /// Run ML training chaos test suite
    MlSuite {
        /// ML service endpoint
        #[arg(short, long, default_value = "http://localhost:8080")]
        endpoint: String,

        /// Checkpoint base path
        #[arg(short, long, default_value = "/tmp/ml_checkpoints")]
        checkpoint_path: PathBuf,

        /// Models to test (if not specified, tests all)
        #[arg(short, long, value_enum)]
        models: Vec<ModelTypeArg>,

        /// GPU memory threshold in MB
        #[arg(long, default_value = "8192")]
        gpu_memory_mb: u64,

        /// Generate report after completion
        #[arg(long)]
        generate_report: bool,

        /// Output report path
        #[arg(long)]
        output_report: Option<PathBuf>,
    },

    /// Start nightly chaos job scheduler
    Schedule {
        /// Schedule time (HH:MM format)
        #[arg(short, long, default_value = "02:00")]
        time: String,

        /// Timezone
        #[arg(short, long, default_value = "UTC")]
        timezone: String,

        /// Exclude weekends
        #[arg(long)]
        exclude_weekends: bool,

        /// Maximum duration in hours
        #[arg(long, default_value = "3")]
        max_duration: u8,

        /// Notification webhook URL
        #[arg(long)]
        webhook: Option<String>,

        /// Report storage path
        #[arg(long, default_value = "./chaos_reports")]
        report_path: PathBuf,
    },

    /// Validate system readiness for chaos testing
    Validate {
        /// Check ML service connectivity
        #[arg(long)]
        check_ml_service: bool,

        /// Check database connectivity
        #[arg(long)]
        check_database: bool,

        /// Check monitoring systems
        #[arg(long)]
        check_monitoring: bool,
    },

    /// List previous chaos experiment results
    History {
        /// Number of recent results to show
        #[arg(short, long, default_value = "10")]
        limit: usize,

        /// Filter by status
        #[arg(short, long)]
        status: Option<String>,

        /// Export to file
        #[arg(short, long)]
        export: Option<PathBuf>,
    },
}

#[derive(ValueEnum, Clone, Debug)]
pub enum ExperimentType {
    ProcessKill,
    MemoryPressure,
    NetworkPartition,
    DiskIoFailure,
    CpuThrottle,
    GpuExhaustion,
    DatabaseFailure,
}

#[derive(ValueEnum, Clone, Debug)]
pub enum ModelTypeArg {
    Tlob,
    Mamba2,
    Dqn,
    Ppo,
    Liquid,
    Tft,
}

impl From<ModelTypeArg> for ModelType {
    fn from(arg: ModelTypeArg) -> Self {
        match arg {
            ModelTypeArg::Tlob => ModelType::TLOB,
            ModelTypeArg::Mamba2 => ModelType::MAMBA2,
            ModelTypeArg::Dqn => ModelType::DQN,
            ModelTypeArg::Ppo => ModelType::PPO,
            ModelTypeArg::Liquid => ModelType::Liquid,
            ModelTypeArg::Tft => ModelType::TFT,
        }
    }
}

/// Main chaos CLI implementation
pub struct ChaosCli {
    orchestrator: ChaosOrchestrator,
    config: Option<NightlyChaosConfig>,
}

impl ChaosCli {
    pub fn new() -> Self {
        Self {
            orchestrator: ChaosOrchestrator::new(3), // Max 3 concurrent experiments
            config: None,
        }
    }

    /// Load configuration from file
    pub async fn load_config(&mut self, config_path: Option<PathBuf>) -> Result<()> {
        self.config = if let Some(path) = config_path {
            let content = tokio::fs::read_to_string(&path)
                .await
                .context(format!("Failed to read config file: {:?}", path))?;
            let config: NightlyChaosConfig =
                toml::from_str(&content).context("Failed to parse configuration file")?;
            Some(config)
        } else {
            Some(NightlyChaosConfig::default())
        };

        Ok(())
    }

    /// Execute the CLI command
    pub async fn execute(&self, args: ChaosCliArgs) -> Result<()> {
        // Initialize logging
        if args.verbose {
            tracing_subscriber::fmt()
                .with_max_level(tracing::Level::DEBUG)
                .init();
        } else {
            tracing_subscriber::fmt()
                .with_max_level(tracing::Level::INFO)
                .init();
        }

        match args.command {
            ChaosCommand::Run {
                experiment_type,
                service,
                model,
                max_recovery_time_ms,
                duration,
                recovery_timeout,
            } => {
                self.run_single_experiment(
                    experiment_type,
                    service,
                    model,
                    max_recovery_time_ms,
                    duration,
                    recovery_timeout,
                )
                .await
            }

            ChaosCommand::MlSuite {
                endpoint,
                checkpoint_path,
                models,
                gpu_memory_mb,
                generate_report,
                output_report,
            } => {
                self.run_ml_suite(
                    endpoint,
                    checkpoint_path,
                    models,
                    gpu_memory_mb,
                    generate_report,
                    output_report,
                )
                .await
            }

            ChaosCommand::Schedule {
                time,
                timezone,
                exclude_weekends,
                max_duration,
                webhook,
                report_path,
            } => {
                self.start_scheduler(
                    time,
                    timezone,
                    exclude_weekends,
                    max_duration,
                    webhook,
                    report_path,
                )
                .await
            }

            ChaosCommand::Validate {
                check_ml_service,
                check_database,
                check_monitoring,
            } => {
                self.validate_system(check_ml_service, check_database, check_monitoring)
                    .await
            }

            ChaosCommand::History {
                limit,
                status,
                export,
            } => self.show_history(limit, status, export).await,
        }
    }

    /// Run a single chaos experiment
    async fn run_single_experiment(
        &self,
        experiment_type: ExperimentType,
        service: String,
        model: Option<ModelTypeArg>,
        max_recovery_time_ms: u64,
        duration: u64,
        recovery_timeout: u64,
    ) -> Result<()> {
        info!(
            "Running single chaos experiment: {:?} on {}",
            experiment_type, service
        );

        let experiment_id = Uuid::new_v4();
        let failure_type = self.create_failure_type(&experiment_type)?;

        let experiment = super::chaos_framework::ChaosExperiment {
            id: experiment_id,
            name: format!("{:?} Test on {}", experiment_type, service),
            description: format!("Single chaos experiment: {:?}", experiment_type),
            target_service: service,
            failure_type,
            duration: std::time::Duration::from_secs(duration),
            recovery_timeout: std::time::Duration::from_secs(recovery_timeout),
            max_recovery_time_ms,
            enabled: true,
        };

        self.orchestrator.register_experiment(experiment).await?;
        let result = self.orchestrator.execute_experiment(experiment_id).await?;

        info!("Experiment completed with status: {:?}", result.status);
        if let Some(recovery_time) = result.recovery_time_ms {
            info!("Recovery time: {}ms", recovery_time);
            if recovery_time <= max_recovery_time_ms {
                info!("✅ Recovery time within HFT requirements");
            } else {
                warn!("⚠️ Recovery time exceeds HFT requirements");
            }
        }

        Ok(())
    }

    /// Run ML training chaos test suite
    async fn run_ml_suite(
        &self,
        endpoint: String,
        checkpoint_path: PathBuf,
        models: Vec<ModelTypeArg>,
        gpu_memory_mb: u64,
        generate_report: bool,
        output_report: Option<PathBuf>,
    ) -> Result<()> {
        info!("Running ML training chaos test suite");

        let model_types = if models.is_empty() {
            // Test all models if none specified
            vec![
                ModelType::TLOB,
                ModelType::MAMBA2,
                ModelType::DQN,
                ModelType::PPO,
                ModelType::Liquid,
                ModelType::TFT,
            ]
        } else {
            models.into_iter().map(Into::into).collect()
        };

        let ml_config = MLChaosConfig {
            ml_service_endpoint: endpoint,
            checkpoint_base_path: checkpoint_path,
            model_types,
            training_timeout_secs: 300,
            max_recovery_time_ms: 100, // HFT requirement
            gpu_memory_threshold_mb: gpu_memory_mb,
        };

        let ml_chaos = MLTrainingChaosTests::new(ml_config);
        let results = ml_chaos.run_ml_chaos_suite().await?;

        info!("ML chaos suite completed with {} results", results.len());

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

        info!("Results: {} successful, {} failed", successful, failed);

        if generate_report {
            let report = ml_chaos.generate_chaos_report(&results).await?;

            if let Some(output_path) = output_report {
                tokio::fs::write(&output_path, &report)
                    .await
                    .context("Failed to write report")?;
                info!("Report saved to: {:?}", output_path);
            } else {
                println!("{}", report);
            }
        }

        Ok(())
    }

    /// Start nightly chaos scheduler
    async fn start_scheduler(
        &self,
        time: String,
        timezone: String,
        exclude_weekends: bool,
        max_duration: u8,
        webhook: Option<String>,
        report_path: PathBuf,
    ) -> Result<()> {
        info!("Starting nightly chaos scheduler at {} {}", time, timezone);

        // Parse time
        let schedule_time = chrono::NaiveTime::parse_from_str(&time, "%H:%M")
            .context("Invalid time format (use HH:MM)")?;

        let config = NightlyChaosConfig {
            enabled: true,
            schedule_time,
            timezone,
            max_duration_hours: max_duration,
            notification_webhook: webhook,
            report_storage_path: report_path,
            ml_chaos_config: self
                .config
                .as_ref()
                .map(|c| c.ml_chaos_config.clone())
                .unwrap_or_default(),
            exclude_weekends,
            retry_on_failure: true,
            max_retries: 2,
        };

        let runner = NightlyChaosRunner::new(config);

        // Subscribe to events for logging
        let mut event_receiver = runner.subscribe_events();
        let event_handler = tokio::spawn(async move {
            while let Ok(event) = event_receiver.recv().await {
                match event {
                    ChaosJobEvent::JobScheduled {
                        id,
                        scheduled_time,
                    } => {
                        info!("📅 Chaos job {} scheduled for {}", id, scheduled_time);
                    }
                    ChaosJobEvent::JobStarted { id } => {
                        info!("🚀 Chaos job {} started", id);
                    }
                    ChaosJobEvent::JobCompleted { id, summary } => {
                        info!(
                            "✅ Chaos job {} completed. Avg recovery: {:.1}ms",
                            id, summary.average_recovery_time_ms
                        );
                    }
                    ChaosJobEvent::JobFailed { id, error } => {
                        error!("❌ Chaos job {} failed: {}", id, error);
                    }
                    ChaosJobEvent::AlertTriggered {
                        message,
                        severity,
                    } => match severity {
                        AlertSeverity::Critical => {
                            error!("🚨 {}", message)
                        }
                        AlertSeverity::Warning => {
                            warn!("⚠️ {}", message)
                        }
                        AlertSeverity::Info => info!("ℹ️ {}", message),
                    },
                    _ => {}
                }
            }
        });

        // Start the scheduler
        info!("Chaos scheduler starting... Press Ctrl+C to stop");
        tokio::select! {
            result = runner.start() => {
                if let Err(e) = result {
                    error!("Scheduler failed: {}", e);
                }
            }
            _ = tokio::signal::ctrl_c() => {
                info!("Received shutdown signal");
            }
        }

        runner.stop().await;
        event_handler.abort();
        info!("Chaos scheduler stopped");

        Ok(())
    }

    /// Validate system readiness
    async fn validate_system(
        &self,
        check_ml_service: bool,
        check_database: bool,
        check_monitoring: bool,
    ) -> Result<()> {
        info!("Validating system readiness for chaos testing");

        let mut all_checks_passed = true;

        if check_ml_service {
            info!("Checking ML service connectivity...");
            match self.validate_ml_service().await {
                Ok(()) => info!("✅ ML service connectivity: OK"),
                Err(e) => {
                    error!("❌ ML service connectivity: FAILED - {}", e);
                    all_checks_passed = false;
                }
            }
        }

        if check_database {
            info!("Checking database connectivity...");
            match self.validate_database().await {
                Ok(()) => info!("✅ Database connectivity: OK"),
                Err(e) => {
                    error!("❌ Database connectivity: FAILED - {}", e);
                    all_checks_passed = false;
                }
            }
        }

        if check_monitoring {
            info!("Checking monitoring systems...");
            match self.validate_monitoring().await {
                Ok(()) => info!("✅ Monitoring systems: OK"),
                Err(e) => {
                    error!("❌ Monitoring systems: FAILED - {}", e);
                    all_checks_passed = false;
                }
            }
        }

        if all_checks_passed {
            info!("🎉 All system checks passed. Ready for chaos testing!");
        } else {
            error!("💥 Some system checks failed. Fix issues before running chaos tests.");
            std::process::exit(1);
        }

        Ok(())
    }

    /// Show experiment history
    async fn show_history(
        &self,
        limit: usize,
        status: Option<String>,
        export: Option<PathBuf>,
    ) -> Result<()> {
        info!("Showing chaos experiment history (limit: {})", limit);

        // Get results from orchestrator
        let results = self.orchestrator.get_results().await;

        // Filter by status if specified
        let filtered_results: Vec<_> = if let Some(ref status_filter) = status {
            results
                .into_iter()
                .filter(|r| {
                    format!("{:?}", r.status)
                        .to_lowercase()
                        .contains(&status_filter.to_lowercase())
                })
                .take(limit)
                .collect()
        } else {
            results.into_iter().take(limit).collect()
        };

        if filtered_results.is_empty() {
            info!("No experiment results found");
            return Ok(());
        }

        // Display results
        println!("\n📊 Chaos Experiment History\n");
        println!(
            "{:<36} {:<20} {:<15} {:<10} {:<15}",
            "Experiment ID", "Started", "Status", "Recovery", "Checkpoint"
        );
        println!("{}", "─".repeat(100));

        for result in &filtered_results {
            let started = result
                .started_at
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs();
            let started_str = chrono::DateTime::from_timestamp(started as i64, 0)
                .map(|dt| dt.format("%Y-%m-%d %H:%M").to_string())
                .unwrap_or_else(|| "Unknown".to_string());

            let recovery_str = result
                .recovery_time_ms
                .map(|t| format!("{}ms", t))
                .unwrap_or_else(|| "N/A".to_string());

            let checkpoint_str = if result.checkpoint_integrity {
                "✅"
            } else {
                "❌"
            };

            println!(
                "{:<36} {:<20} {:<15} {:<10} {:<15}",
                result.experiment_id,
                started_str,
                format!("{:?}", result.status),
                recovery_str,
                checkpoint_str
            );
        }

        // Export if requested
        if let Some(export_path) = export {
            let export_data = serde_json::to_string_pretty(&filtered_results)
                .context("Failed to serialize results")?;
            tokio::fs::write(&export_path, export_data)
                .await
                .context("Failed to write export file")?;
            info!("Results exported to: {:?}", export_path);
        }

        Ok(())
    }

    /// Create failure type from experiment type
    fn create_failure_type(
        &self,
        experiment_type: &ExperimentType,
    ) -> Result<super::chaos_framework::FailureType> {
        use super::chaos_framework::{FailureType, Signal};

        let failure_type = match experiment_type {
            ExperimentType::ProcessKill => FailureType::ProcessKill {
                signal: Signal::SIGTERM,
                delay_before_restart_ms: 2000,
            },
            ExperimentType::MemoryPressure => FailureType::MemoryPressure {
                target_mb: 4096,
                duration_ms: 30000,
            },
            ExperimentType::NetworkPartition => FailureType::NetworkPartition {
                target_ports: vec![8080, 5432, 6379],
                duration_ms: 15000,
            },
            ExperimentType::DiskIoFailure => FailureType::DiskIoFailure {
                target_paths: vec!["/tmp".to_string(), "/var/log".to_string()],
                failure_rate_percent: 30,
            },
            ExperimentType::CpuThrottle => FailureType::CpuThrottle {
                cpu_limit_percent: 50,
                duration_ms: 20000,
            },
            ExperimentType::GpuExhaustion => FailureType::GpuResourceExhaustion {
                memory_fill_percent: 95,
                duration_ms: 25000,
            },
            ExperimentType::DatabaseFailure => FailureType::DatabaseConnectionFailure {
                connection_string: "postgresql://localhost:5432/foxhunt".to_string(),
                duration_ms: 10000,
            },
        };

        Ok(failure_type)
    }

    /// Validate ML service connectivity
    async fn validate_ml_service(&self) -> Result<()> {
        // TODO: Implement actual ML service health check
        // This would make a gRPC call to the ML service health endpoint
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        Ok(())
    }

    /// Validate database connectivity
    async fn validate_database(&self) -> Result<()> {
        // TODO: Implement actual database connectivity check
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        Ok(())
    }

    /// Validate monitoring systems
    async fn validate_monitoring(&self) -> Result<()> {
        // TODO: Implement actual monitoring system checks (Prometheus, etc.)
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        Ok(())
    }
}

impl Default for ChaosCli {
    fn default() -> Self {
        Self::new()
    }
}

/// Main CLI entry point
pub async fn main() -> Result<()> {
    let args = ChaosCliArgs::parse();
    let mut cli = ChaosCli::new();

    // Load configuration if specified
    cli.load_config(args.config.clone()).await?;

    // Execute the command
    cli.execute(args).await?;

    Ok(())
}

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

    #[test]
    fn test_model_type_conversion() {
        assert_eq!(ModelType::from(ModelTypeArg::Tlob), ModelType::TLOB);
        assert_eq!(ModelType::from(ModelTypeArg::Dqn), ModelType::DQN);
    }

    #[tokio::test]
    async fn test_cli_creation() {
        let cli = ChaosCli::new();
        assert!(true); // Basic creation test
    }
}