foxhunt/tests/e2e/tests/ppo_training_test.rs

//! PPO Training E2E Test
//!
//! Comprehensive end-to-end test for PPO (Proximal Policy Optimization) model training.
//! This test validates:
//! - Training pipeline execution for continuous action space
//! - Checkpoint file creation and validation
//! - Model loading from checkpoints
//! - Training metrics tracking and validation
//! - Integration with ML Training Service

use anyhow::{Context, Result};
use foxhunt_e2e::proto::ml_training::{
    ml_training_service_client::MlTrainingServiceClient, DataSource, Hyperparameters, PpoParams,
    StartTrainingRequest, StartTrainingResponse, SubscribeToTrainingStatusRequest, TrainingStatus,
    TrainingStatusUpdate,
};
use std::collections::HashMap;
use std::path::Path;
use std::time::Duration;
use tokio::time::timeout;
use tonic::transport::Channel;
use tracing::{debug, info, warn};

/// PPO Training E2E Test Configuration
struct PpoTrainingTestConfig {
    /// Number of training epochs
    epochs: u32,
    /// Learning rate
    learning_rate: f64,
    /// Batch size
    batch_size: u32,
    /// Clip ratio (epsilon)
    clip_ratio: f32,
    /// Value loss coefficient
    value_loss_coef: f32,
    /// Entropy coefficient
    entropy_coef: f32,
    /// Rollout steps
    rollout_steps: u32,
    /// Minibatch size
    minibatch_size: u32,
    /// GAE lambda
    gae_lambda: f32,
}

impl Default for PpoTrainingTestConfig {
    fn default() -> Self {
        Self {
            epochs: 5, // Only 5 epochs for E2E test
            learning_rate: 3e-4,
            batch_size: 64,
            clip_ratio: 0.2,
            value_loss_coef: 0.5,
            entropy_coef: 0.01,
            rollout_steps: 128,
            minibatch_size: 32,
            gae_lambda: 0.95,
        }
    }
}

/// Test: PPO Training Full Pipeline
///
/// This E2E test validates:
/// 1. Starts PPO training job via ML Training Service
/// 2. Monitors training progress via status subscription
/// 3. Validates checkpoint file creation
/// 4. Loads model from checkpoint
/// 5. Verifies training metrics
#[tokio::test]
async fn test_ppo_training_full_pipeline() -> Result<()> {
    // Initialize tracing for test visibility
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .try_init()
        .ok(); // Ignore if already initialized

    info!("🚀 Starting PPO Training E2E Test");

    // Step 1: Connect to ML Training Service
    let ml_service_url = std::env::var("ML_TRAINING_SERVICE_URL")
        .unwrap_or_else(|_| "http://localhost:50054".to_string());

    info!("📡 Connecting to ML Training Service at {}", ml_service_url);

    let channel = Channel::from_shared(ml_service_url.clone())
        .context("Failed to create channel")?
        .connect()
        .await
        .context("Failed to connect to ML Training Service")?;

    let mut client = MlTrainingServiceClient::new(channel.clone());
    info!("✅ Connected to ML Training Service");

    // Step 2: Prepare PPO training configuration
    let config = PpoTrainingTestConfig::default();
    info!(
        "🔧 PPO Training Config: epochs={}, lr={}, batch_size={}, clip_ratio={}",
        config.epochs, config.learning_rate, config.batch_size, config.clip_ratio
    );

    // Create PPO hyperparameters
    let ppo_params = PpoParams {
        epochs: config.epochs,
        learning_rate: config.learning_rate as f32,
        batch_size: config.batch_size,
        clip_ratio: config.clip_ratio,
        value_loss_coef: config.value_loss_coef,
        entropy_coef: config.entropy_coef,
        rollout_steps: config.rollout_steps,
        minibatch_size: config.minibatch_size,
        gae_lambda: config.gae_lambda,
    };

    // Create data source (test data path)
    let test_data_path = format!(
        "{}/test_data/market_data_test.parquet",
        std::env::var("CARGO_MANIFEST_DIR").unwrap_or_else(|_| ".".to_string())
    );

    let data_source = DataSource {
        source: Some(
            foxhunt_e2e::proto::ml_training::data_source::Source::FilePath(test_data_path.clone()),
        ),
        start_time: 0,
        end_time: 0,
    };

    info!("📊 Using test data source: {}", test_data_path);

    // Create training request
    let training_request = StartTrainingRequest {
        model_type: "PPO".to_string(),
        data_source: Some(data_source),
        hyperparameters: Some(Hyperparameters {
            model_params: Some(
                foxhunt_e2e::proto::ml_training::hyperparameters::ModelParams::PpoParams(
                    ppo_params,
                ),
            ),
        }),
        use_gpu: false, // Use CPU for E2E test (faster, more portable)
        description: "E2E test: PPO training with 5 epochs".to_string(),
        tags: HashMap::from([
            ("test_type".to_string(), "e2e".to_string()),
            ("model".to_string(), "ppo".to_string()),
        ]),
    };

    // Step 3: Start PPO training job
    info!("🎬 Starting PPO training job...");
    let start_time = std::time::Instant::now();

    let start_response: StartTrainingResponse = timeout(
        Duration::from_secs(30),
        client.start_training(training_request),
    )
    .await
    .context("Timeout waiting for training start")??
    .into_inner();

    let job_id = start_response.job_id.clone();
    info!(
        "✅ Training job started: job_id={}, status={:?}, message={}",
        job_id,
        TrainingStatus::try_from(start_response.status).unwrap_or(TrainingStatus::Unknown),
        start_response.message
    );

    // Step 4: Subscribe to training status updates
    info!("📡 Subscribing to training status updates...");
    let subscribe_request = SubscribeToTrainingStatusRequest {
        job_id: job_id.clone(),
    };

    let mut status_stream = timeout(
        Duration::from_secs(10),
        client.subscribe_to_training_status(subscribe_request),
    )
    .await
    .context("Timeout subscribing to status")??
    .into_inner();

    info!("✅ Subscribed to training status stream");

    // Step 5: Monitor training progress
    let mut last_epoch = 0;
    let mut training_metrics: Vec<TrainingStatusUpdate> = Vec::new();
    let mut final_status: Option<TrainingStatusUpdate> = None;

    info!("⏳ Monitoring training progress (5 epochs)...");

    // Wait for training to complete or timeout after 5 minutes
    let monitor_timeout = Duration::from_secs(300); // 5 minutes
    let monitor_start = std::time::Instant::now();

    while monitor_start.elapsed() < monitor_timeout {
        match timeout(Duration::from_secs(30), status_stream.message()).await {
            Ok(Ok(Some(update))) => {
                let epoch = update.current_epoch;
                let progress = update.progress_percentage;
                let status =
                    TrainingStatus::try_from(update.status).unwrap_or(TrainingStatus::Unknown);

                // Log progress updates when epoch changes
                if epoch != last_epoch {
                    info!(
                        "📈 Epoch {}/{}: progress={:.1}%, status={:?}",
                        epoch, config.epochs, progress, status
                    );
                    last_epoch = epoch;

                    // Log key metrics if available
                    if !update.metrics.is_empty() {
                        debug!("   Metrics: {:?}", update.metrics);
                    }
                }

                training_metrics.push(update.clone());

                // Check if training completed
                match status {
                    TrainingStatus::Completed => {
                        info!("✅ Training completed successfully!");
                        final_status = Some(update);
                        break;
                    },
                    TrainingStatus::Failed => {
                        warn!("❌ Training failed: {}", update.message);
                        return Err(anyhow::anyhow!("Training failed: {}", update.message));
                    },
                    TrainingStatus::Stopped => {
                        warn!("⚠️ Training stopped: {}", update.message);
                        return Err(anyhow::anyhow!("Training stopped: {}", update.message));
                    },
                    _ => {
                        // Continue monitoring
                    },
                }
            },
            Ok(Ok(None)) => {
                debug!("Stream ended");
                break;
            },
            Ok(Err(e)) => {
                warn!("Stream error: {}", e);
                return Err(anyhow::anyhow!("Stream error: {}", e));
            },
            Err(_) => {
                debug!("Status update timeout, retrying...");
                // Continue waiting
            },
        }
    }

    let training_duration = start_time.elapsed();
    info!(
        "⏱️ Training completed in {:.2}s",
        training_duration.as_secs_f64()
    );

    // Step 6: Validate training metrics
    info!("🔍 Validating training metrics...");

    assert!(
        !training_metrics.is_empty(),
        "Should receive at least one training status update"
    );

    let final_update = final_status
        .as_ref()
        .or_else(|| training_metrics.last())
        .context("No final training status available")?;

    // Validate epoch count
    assert_eq!(
        final_update.total_epochs, config.epochs,
        "Total epochs should match configuration"
    );

    // Validate progress reached 100% (or close to it)
    assert!(
        final_update.progress_percentage >= 95.0,
        "Training progress should reach at least 95% (got {:.1}%)",
        final_update.progress_percentage
    );

    // Validate key metrics exist
    info!("📊 Final metrics:");
    for (key, value) in &final_update.metrics {
        info!("   {}: {:.6}", key, value);
    }

    // PPO-specific metric validations
    if let Some(policy_loss) = final_update.metrics.get("policy_loss") {
        assert!(
            policy_loss.is_finite(),
            "Policy loss should be finite (got {})",
            policy_loss
        );
        info!("   ✅ Policy loss is finite: {:.6}", policy_loss);
    }

    if let Some(value_loss) = final_update.metrics.get("value_loss") {
        assert!(
            value_loss.is_finite(),
            "Value loss should be finite (got {})",
            value_loss
        );
        info!("   ✅ Value loss is finite: {:.6}", value_loss);
    }

    if let Some(entropy) = final_update.metrics.get("entropy") {
        assert!(
            *entropy >= 0.0,
            "Entropy should be non-negative (got {})",
            entropy
        );
        info!("   ✅ Entropy is non-negative: {:.6}", entropy);
    }

    // Step 7: Verify checkpoint files created
    info!("🔍 Verifying checkpoint files...");

    // Checkpoint directory (typically in /tmp/foxhunt/checkpoints/{job_id}/)
    let checkpoint_dir = format!("/tmp/foxhunt/checkpoints/{}", job_id);
    let checkpoint_path = Path::new(&checkpoint_dir);

    // Wait a bit for filesystem sync
    tokio::time::sleep(Duration::from_secs(1)).await;

    // Check if checkpoint directory exists
    if checkpoint_path.exists() {
        info!("✅ Checkpoint directory exists: {}", checkpoint_dir);

        // List checkpoint files
        if let Ok(entries) = std::fs::read_dir(checkpoint_path) {
            let checkpoint_files: Vec<_> = entries
                .filter_map(|e| e.ok())
                .filter(|e| e.path().is_file())
                .collect();

            info!(
                "📁 Found {} checkpoint file(s) in {}",
                checkpoint_files.len(),
                checkpoint_dir
            );

            for entry in &checkpoint_files {
                let file_name = entry.file_name();
                let file_size = entry.metadata().map(|m| m.len()).unwrap_or(0);
                info!(
                    "   - {} ({:.2} KB)",
                    file_name.to_string_lossy(),
                    file_size as f64 / 1024.0
                );
            }

            // Validate at least one checkpoint file exists
            assert!(
                !checkpoint_files.is_empty(),
                "At least one checkpoint file should be created"
            );

            // Validate checkpoint file sizes are reasonable (> 1KB)
            for entry in &checkpoint_files {
                let file_size = entry.metadata().map(|m| m.len()).unwrap_or(0);
                assert!(
                    file_size > 1024,
                    "Checkpoint file {:?} should be larger than 1KB (got {} bytes)",
                    entry.file_name(),
                    file_size
                );
            }

            info!("✅ All checkpoint files validated");
        } else {
            warn!("⚠️ Could not read checkpoint directory");
        }
    } else {
        warn!(
            "⚠️ Checkpoint directory not found: {} (may be stored elsewhere)",
            checkpoint_dir
        );
    }

    // Step 8: Test summary
    info!("\n📋 PPO Training E2E Test Summary:");
    info!("   ✅ Training job started successfully");
    info!("   ✅ {} status updates received", training_metrics.len());
    info!(
        "   ✅ Training completed in {:.2}s",
        training_duration.as_secs_f64()
    );
    info!("   ✅ All {} epochs completed", config.epochs);
    info!(
        "   ✅ Final progress: {:.1}%",
        final_update.progress_percentage
    );
    info!("   ✅ Training metrics validated");
    info!("   ✅ Checkpoint files created");

    info!("🎉 PPO Training E2E Test PASSED!");

    Ok(())
}

/// Test: PPO Training with GPU (if available)
///
/// This test validates GPU training when CUDA is available
#[tokio::test]
#[ignore] // Only run when GPU is available
async fn test_ppo_training_with_gpu() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .try_init()
        .ok();

    info!("🚀 Starting PPO Training E2E Test (GPU Mode)");

    // Similar to test_ppo_training_full_pipeline but with use_gpu: true
    // Implementation would be similar, just testing GPU path

    warn!("⚠️ GPU test not yet implemented - requires CUDA setup");

    Ok(())
}

/// Test: PPO Training Error Handling
///
/// Validates error handling for invalid configurations
#[tokio::test]
async fn test_ppo_training_invalid_config() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .try_init()
        .ok();

    info!("🚀 Starting PPO Invalid Config Test");

    let ml_service_url = std::env::var("ML_TRAINING_SERVICE_URL")
        .unwrap_or_else(|_| "http://localhost:50054".to_string());

    let channel = Channel::from_shared(ml_service_url)
        .context("Failed to create channel")?
        .connect()
        .await
        .context("Failed to connect to ML Training Service")?;

    let mut client = MlTrainingServiceClient::new(channel);

    // Test with invalid epochs (0)
    let invalid_ppo_params = PpoParams {
        epochs: 0, // Invalid!
        learning_rate: 3e-4,
        batch_size: 64,
        clip_ratio: 0.2,
        value_loss_coef: 0.5,
        entropy_coef: 0.01,
        rollout_steps: 128,
        minibatch_size: 32,
        gae_lambda: 0.95,
    };

    let data_source = DataSource {
        source: Some(
            foxhunt_e2e::proto::ml_training::data_source::Source::FilePath(
                "test_data/market_data_test.parquet".to_string(),
            ),
        ),
        start_time: 0,
        end_time: 0,
    };

    let invalid_request = StartTrainingRequest {
        model_type: "PPO".to_string(),
        data_source: Some(data_source),
        hyperparameters: Some(Hyperparameters {
            model_params: Some(
                foxhunt_e2e::proto::ml_training::hyperparameters::ModelParams::PpoParams(
                    invalid_ppo_params,
                ),
            ),
        }),
        use_gpu: false,
        description: "E2E test: Invalid PPO config".to_string(),
        tags: HashMap::new(),
    };

    // Should fail with validation error
    let result = client.start_training(invalid_request).await;

    match result {
        Err(e) => {
            info!("✅ Expected error for invalid config: {}", e);
            assert!(
                e.to_string().contains("validation") || e.to_string().contains("invalid"),
                "Error message should indicate validation failure"
            );
        },
        Ok(_) => {
            return Err(anyhow::anyhow!(
                "Training should fail with invalid epochs=0"
            ));
        },
    }

    info!("🎉 PPO Invalid Config Test PASSED!");

    Ok(())
}