foxhunt/services/integration_tests/tests/ml_training_service_e2e.rs

//! End-to-End Integration Tests: API Gateway → ML Training Service
//!
//! This test suite validates the complete ML training service flow through the API Gateway:
//! - Training job management (start, stop, list)
//! - Real-time training progress monitoring
//! - Resource utilization tracking
//! - Training configuration validation
//! - Model deployment workflow
//!
//! Test Count: 12 tests
//! Coverage: Full ML training service integration

use anyhow::Result;
use chrono::{Duration, Utc};
use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Duration as StdDuration;
use tokio::time::timeout;
use tonic::{metadata::MetadataValue, transport::Channel, Request};
use uuid::Uuid;

// Generated proto code
pub mod ml {
    tonic::include_proto!("foxhunt.ml");
}

use ml::{
    ml_training_service_client::MlTrainingServiceClient,
    ListTrainingJobsRequest,
    ResourceRequest,
    ResourceRequirements,
    StartTrainingRequest,
    StopTrainingRequest,
    TrainingConfigRequest,
    TrainingHyperparameters,
    TrainingStatus,
    TrainingTemplatesRequest,
    WatchTrainingRequest,
};

const API_GATEWAY_ADDR: &str = "http://localhost:50050";
const JWT_SECRET: &str = "dev_secret_key_change_in_production";

#[derive(Debug, Serialize, Deserialize)]
struct Claims {
    sub: String,
    exp: usize,
    iat: usize,
    roles: Vec<String>,
    permissions: Vec<String>,
    session_id: Option<String>,
    jti: String,
}

/// Generate a test JWT token for authentication
fn generate_test_token(user_id: &str, roles: Vec<String>, permissions: Vec<String>) -> Result<String> {
    let now = Utc::now();
    let claims = Claims {
        sub: user_id.to_string(),
        exp: (now + Duration::hours(1)).timestamp() as usize,
        iat: now.timestamp() as usize,
        roles,
        permissions,
        session_id: Some(Uuid::new_v4().to_string()),
        jti: Uuid::new_v4().to_string(),
    };

    let token = encode(
        &Header::new(Algorithm::HS256),
        &claims,
        &EncodingKey::from_secret(JWT_SECRET.as_ref()),
    )?;

    Ok(token)
}

/// Create an authenticated ML training service client
async fn create_authenticated_client() -> Result<MlTrainingServiceClient<tonic::service::interceptor::InterceptedService<Channel, impl Fn(Request<()>) -> Result<Request<()>, tonic::Status> + Clone>>> {
    let user_id = "test_ml_engineer_001";
    let role = "ml_engineer";

    let token = generate_test_token(
        user_id,
        vec![role.to_string()],
        vec![
            "api.access".to_string(),
            "ml.train".to_string(),
            "ml.view".to_string(),
            "ml.manage".to_string(),
        ],
    )?;

    let channel = Channel::from_static(API_GATEWAY_ADDR)
        .connect()
        .await?;

    // Create interceptor that injects JWT token AND user context into request metadata
    let user_id_owned = user_id.to_string();
    let role_owned = role.to_string();

    let interceptor = move |mut req: Request<()>| -> Result<Request<()>, tonic::Status> {
        // JWT token in authorization header
        let token_value = format!("Bearer {}", token);
        let metadata_value = MetadataValue::try_from(token_value)
            .map_err(|_| tonic::Status::internal("Failed to create metadata value"))?;
        req.metadata_mut().insert("authorization", metadata_value);

        // User context in metadata headers
        let user_id_value = MetadataValue::try_from(user_id_owned.clone())
            .map_err(|_| tonic::Status::internal("Failed to create user_id metadata"))?;
        req.metadata_mut().insert("x-user-id", user_id_value);

        let role_value = MetadataValue::try_from(role_owned.clone())
            .map_err(|_| tonic::Status::internal("Failed to create role metadata"))?;
        req.metadata_mut().insert("x-user-role", role_value);

        Ok(req)
    };

    let client = MlTrainingServiceClient::with_interceptor(channel, interceptor);

    Ok(client)
}

// ============================================================================
// SECTION 1: TRAINING JOB LIFECYCLE E2E TESTS (5 tests)
// ============================================================================

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_training_job_start() -> Result<()> {
    println!("\n=== E2E Test: Start Training Job via API Gateway ===");

    let mut client = create_authenticated_client().await?;

    let hyperparameters = TrainingHyperparameters {
        learning_rate: 0.001,
        batch_size: 32,
        epochs: 10,
        dropout_rate: Some(0.2),
        hidden_layers: Some(3),
        hidden_units: Some(128),
        custom_params: HashMap::new(),
    };

    let resource_requirements = ResourceRequirements {
        gpu_count: 1,
        cpu_cores: 4,
        memory_gb: 16,
        gpu_type: Some("RTX 3050 Ti".to_string()),
        disk_gb: 50,
    };

    let request = Request::new(StartTrainingRequest {
        model_name: "dqn_trading_model".to_string(),
        dataset_id: "btc_eth_historical_2024".to_string(),
        hyperparameters: Some(hyperparameters),
        resource_requirements: Some(resource_requirements),
        tags: vec!["e2e_test".to_string(), "dqn".to_string()],
        description: "E2E test training job - DQN model".to_string(),
        auto_deploy: false,
    });

    let response = client.start_training(request).await?;
    let training_job = response.into_inner();

    assert!(!training_job.job_id.is_empty(), "Should return job ID");
    assert_eq!(training_job.model_name, "dqn_trading_model");

    println!("✓ Training job started successfully");
    println!("  Job ID: {}", training_job.job_id);
    println!("  Model: {}", training_job.model_name);
    println!("  Status: {:?}", training_job.status);

    Ok(())
}

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_training_job_stop() -> Result<()> {
    println!("\n=== E2E Test: Stop Training Job via API Gateway ===");

    let mut client = create_authenticated_client().await?;

    // Start a training job first
    let start_request = Request::new(StartTrainingRequest {
        model_name: "ppo_model".to_string(),
        dataset_id: "crypto_dataset_2024".to_string(),
        hyperparameters: Some(TrainingHyperparameters {
            learning_rate: 0.001,
            batch_size: 64,
            epochs: 100, // Long training
            dropout_rate: None,
            hidden_layers: None,
            hidden_units: None,
            custom_params: HashMap::new(),
        }),
        resource_requirements: Some(ResourceRequirements {
            gpu_count: 1,
            cpu_cores: 2,
            memory_gb: 8,
            gpu_type: None,
            disk_gb: 20,
        }),
        tags: vec!["e2e_stop_test".to_string()],
        description: "E2E stop test - PPO model".to_string(),
        auto_deploy: false,
    });

    let start_response = client.start_training(start_request).await?;
    let job_id = start_response.into_inner().job_id;

    println!("✓ Training job started: {}", job_id);

    // Wait for training to actually start
    tokio::time::sleep(StdDuration::from_secs(1)).await;

    // Stop the training job
    let stop_request = Request::new(StopTrainingRequest {
        job_id: job_id.clone(),
        force: false,
    });

    let stop_response = client.stop_training(stop_request).await?;
    let stopped_job = stop_response.into_inner();

    assert_eq!(stopped_job.job_id, job_id);

    println!("✓ Training job stopped successfully");
    println!("  Final Status: {:?}", stopped_job.status);

    Ok(())
}

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_list_training_jobs() -> Result<()> {
    println!("\n=== E2E Test: List Training Jobs via API Gateway ===");

    let mut client = create_authenticated_client().await?;

    let request = Request::new(ListTrainingJobsRequest {
        model_name: None,
        status: None,
        start_time_after: None,
        start_time_before: None,
        tags: vec![],
        limit: 10,
        cursor: String::new(),
    });

    let response = client.list_training_jobs(request).await?;
    let list_result = response.into_inner();

    println!("✓ Training jobs list retrieved");
    println!("  Total Count: {}", list_result.total_count);
    println!("  Returned: {}", list_result.jobs.len());

    for (i, job) in list_result.jobs.into_iter().enumerate().take(5) {
        println!("  {}. {} - {} (Status: {:?})",
            i + 1,
            job.job_id,
            job.model_name,
            job.status
        );
    }

    Ok(())
}

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_filter_training_jobs_by_model() -> Result<()> {
    println!("\n=== E2E Test: Filter Training Jobs by Model Name ===");

    let mut client = create_authenticated_client().await?;

    let request = Request::new(ListTrainingJobsRequest {
        model_name: Some("dqn_trading_model".to_string()),
        status: None,
        start_time_after: None,
        start_time_before: None,
        tags: vec![],
        limit: 10,
        cursor: String::new(),
    });

    let response = client.list_training_jobs(request).await?;
    let list_result = response.into_inner();

    println!("✓ Filtered training jobs retrieved");
    println!("  Filter: model_name = dqn_trading_model");
    println!("  Count: {}", list_result.jobs.len());

    for job in &list_result.jobs {
        assert_eq!(job.model_name, "dqn_trading_model");
    }

    Ok(())
}

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_filter_training_jobs_by_status() -> Result<()> {
    println!("\n=== E2E Test: Filter Training Jobs by Status ===");

    let mut client = create_authenticated_client().await?;

    let request = Request::new(ListTrainingJobsRequest {
        model_name: None,
        status: Some(TrainingStatus::Completed as i32),
        start_time_after: None,
        start_time_before: None,
        tags: vec![],
        limit: 10,
        cursor: String::new(),
    });

    let response = client.list_training_jobs(request).await?;
    let list_result = response.into_inner();

    println!("✓ Filtered training jobs retrieved");
    println!("  Filter: status = COMPLETED");
    println!("  Count: {}", list_result.jobs.len());

    for job in &list_result.jobs {
        assert_eq!(job.status, TrainingStatus::Completed as i32);
    }

    Ok(())
}

// ============================================================================
// SECTION 2: REAL-TIME MONITORING E2E TESTS (3 tests)
// ============================================================================

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_watch_training_progress() -> Result<()> {
    println!("\n=== E2E Test: Watch Training Progress via API Gateway ===");

    let mut client = create_authenticated_client().await?;

    // Start a training job
    let start_request = Request::new(StartTrainingRequest {
        model_name: "mamba_model".to_string(),
        dataset_id: "high_freq_trading_data".to_string(),
        hyperparameters: Some(TrainingHyperparameters {
            learning_rate: 0.0001,
            batch_size: 128,
            epochs: 50,
            dropout_rate: Some(0.3),
            hidden_layers: Some(4),
            hidden_units: Some(256),
            custom_params: HashMap::new(),
        }),
        resource_requirements: Some(ResourceRequirements {
            gpu_count: 1,
            cpu_cores: 8,
            memory_gb: 32,
            gpu_type: Some("RTX 3050 Ti".to_string()),
            disk_gb: 100,
        }),
        tags: vec!["e2e_progress_test".to_string(), "mamba".to_string()],
        description: "E2E progress monitoring test".to_string(),
        auto_deploy: false,
    });

    let start_response = client.start_training(start_request).await?;
    let job_id = start_response.into_inner().job_id;

    println!("✓ Training job started: {}", job_id);

    // Watch training progress
    let watch_request = Request::new(WatchTrainingRequest {
        job_id: job_id.clone(),
        include_logs: true,
        include_metrics: true,
    });

    let mut stream = client.watch_training_progress(watch_request).await?.into_inner();

    println!("✓ Progress stream established");

    // Receive progress updates (with timeout)
    let mut updates_received = 0;
    while let Ok(update_result) = timeout(
        StdDuration::from_secs(10),
        stream.message()
    ).await {
        if let Ok(Some(update)) = update_result {
            updates_received += 1;

            println!("  Update #{}: Epoch {}/{} - Progress: {:.1}%",
                updates_received,
                update.current_epoch,
                update.total_epochs,
                update.progress_percentage
            );

            if let Some(metrics) = &update.metrics {
                println!("    Loss: {:.4}, Accuracy: {:.2}%",
                    metrics.loss,
                    metrics.accuracy * 100.0
                );
            }

            if let Some(log) = &update.log_message {
                println!("    Log: {}", log);
            }

            // Stop after 5 updates or if completed
            if updates_received >= 5 || update.status == TrainingStatus::Completed as i32 {
                break;
            }
        }
    }

    assert!(updates_received > 0, "Should receive at least one progress update");
    println!("✓ Received {} progress updates", updates_received);

    Ok(())
}

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_resource_utilization() -> Result<()> {
    println!("\n=== E2E Test: Get Resource Utilization via API Gateway ===");

    let mut client = create_authenticated_client().await?;

    let request = Request::new(ResourceRequest {});

    let response = client.get_resource_utilization(request).await?;
    let resources = response.into_inner();

    println!("✓ Resource utilization retrieved");

    if let Some(utilization) = resources.current_utilization {
        println!("  GPU Utilization: {:.1}%", utilization.gpu_utilization * 100.0);
        println!("  GPU Memory: {:.1}%", utilization.gpu_memory_used * 100.0);
        println!("  CPU Utilization: {:.1}%", utilization.cpu_utilization * 100.0);
        println!("  Memory Used: {:.1}%", utilization.memory_used * 100.0);
    }

    println!("  Available GPUs: {} / {}", resources.available_gpus, resources.total_gpus);
    println!("  Active Training Jobs: {}", resources.active_training_jobs.len());

    Ok(())
}

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_stream_resource_metrics() -> Result<()> {
    println!("\n=== E2E Test: Stream Resource Metrics via API Gateway ===");

    let mut client = create_authenticated_client().await?;

    let request = Request::new(ResourceRequest {});

    let mut stream = client.stream_resource_metrics(request).await?.into_inner();

    println!("✓ Resource metrics stream established");

    // Receive resource metric updates (with timeout)
    let mut updates_received = 0;
    while let Ok(update_result) = timeout(
        StdDuration::from_secs(5),
        stream.message()
    ).await {
        if let Ok(Some(update)) = update_result {
            updates_received += 1;

            if let Some(utilization) = &update.utilization {
                println!("  Metrics Update #{}: GPU: {:.1}%, CPU: {:.1}%",
                    updates_received,
                    utilization.gpu_utilization * 100.0,
                    utilization.cpu_utilization * 100.0
                );
            }

            println!("    Active Jobs: {}", update.active_jobs.len());

            if updates_received >= 3 {
                break;
            }
        }
    }

    assert!(updates_received > 0, "Should receive at least one resource metric update");
    println!("✓ Received {} resource metric updates", updates_received);

    Ok(())
}

// ============================================================================
// SECTION 3: CONFIGURATION AND VALIDATION E2E TESTS (4 tests)
// ============================================================================

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_validate_training_config() -> Result<()> {
    println!("\n=== E2E Test: Validate Training Configuration ===");

    let mut client = create_authenticated_client().await?;

    let hyperparameters = TrainingHyperparameters {
        learning_rate: 0.01,
        batch_size: 64,
        epochs: 20,
        dropout_rate: Some(0.5),
        hidden_layers: Some(2),
        hidden_units: Some(64),
        custom_params: HashMap::new(),
    };

    let resource_requirements = ResourceRequirements {
        gpu_count: 1,
        cpu_cores: 4,
        memory_gb: 16,
        gpu_type: None,
        disk_gb: 50,
    };

    let request = Request::new(TrainingConfigRequest {
        model_name: "test_model".to_string(),
        hyperparameters: Some(hyperparameters),
        resource_requirements: Some(resource_requirements),
    });

    let response = client.validate_training_config(request).await?;
    let validation = response.into_inner();

    println!("✓ Configuration validation completed");
    println!("  Valid: {}", validation.valid);
    println!("  Errors: {}", validation.validation_errors.len());
    println!("  Warnings: {}", validation.validation_warnings.len());
    println!("  Estimated Duration: {:.1}h", validation.estimated_duration_hours);

    for error in &validation.validation_errors {
        println!("    Error: {}", error);
    }

    for warning in &validation.validation_warnings {
        println!("    Warning: {}", warning);
    }

    Ok(())
}

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_get_training_templates() -> Result<()> {
    println!("\n=== E2E Test: Get Training Templates via API Gateway ===");

    let mut client = create_authenticated_client().await?;

    let request = Request::new(TrainingTemplatesRequest {
        model_type: None, // Get all templates
    });

    let response = client.get_training_templates(request).await?;
    let templates = response.into_inner();

    println!("✓ Training templates retrieved");
    println!("  Total Templates: {}", templates.templates.len());

    for (i, template) in templates.templates.into_iter().enumerate().take(5) {
        println!("  {}. {} - {}",
            i + 1,
            template.template_id,
            template.name
        );
        println!("     Description: {}", template.description);
        println!("     Model Type: {}", template.model_type);
    }

    Ok(())
}

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_filter_templates_by_model_type() -> Result<()> {
    println!("\n=== E2E Test: Filter Templates by Model Type ===");

    let mut client = create_authenticated_client().await?;

    let request = Request::new(TrainingTemplatesRequest {
        model_type: Some("DQN".to_string()),
    });

    let response = client.get_training_templates(request).await?;
    let templates = response.into_inner();

    println!("✓ Filtered templates retrieved");
    println!("  Filter: model_type = DQN");
    println!("  Count: {}", templates.templates.len());

    for template in &templates.templates {
        assert_eq!(template.model_type, "DQN");
    }

    Ok(())
}

#[tokio::test]
#[ignore] // Requires running services
async fn test_e2e_training_with_auto_deploy() -> Result<()> {
    println!("\n=== E2E Test: Training with Auto-Deploy ===");

    let mut client = create_authenticated_client().await?;

    let request = Request::new(StartTrainingRequest {
        model_name: "auto_deploy_model".to_string(),
        dataset_id: "quick_training_dataset".to_string(),
        hyperparameters: Some(TrainingHyperparameters {
            learning_rate: 0.001,
            batch_size: 32,
            epochs: 5, // Quick training for auto-deploy test
            dropout_rate: None,
            hidden_layers: None,
            hidden_units: None,
            custom_params: HashMap::new(),
        }),
        resource_requirements: Some(ResourceRequirements {
            gpu_count: 1,
            cpu_cores: 2,
            memory_gb: 8,
            gpu_type: None,
            disk_gb: 20,
        }),
        tags: vec!["e2e_auto_deploy".to_string()],
        description: "E2E test with auto-deploy enabled".to_string(),
        auto_deploy: true, // Enable auto-deploy
    });

    let response = client.start_training(request).await?;
    let training_job = response.into_inner();

    println!("✓ Training with auto-deploy started");
    println!("  Job ID: {}", training_job.job_id);
    println!("  Auto Deploy: Enabled");
    println!("  Description: {}", training_job.description);

    Ok(())
}