foxhunt/services/ml_training_service/src/main.rs

//! ML Training Service
//!
//! Production-ready ML training service for the Foxhunt HFT trading system.
//! Provides gRPC API for training job management with comprehensive orchestration,
//! resource management, and progress tracking.

use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result};
use clap::{Args, Parser, Subcommand};
use tonic::transport::Server;
use tonic_reflection::server::Builder as ReflectionBuilder;
use tracing::{debug, error, info, warn};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

// Import from library instead of duplicating module declarations
use ml_training_service::{database, encryption, gpu_config, orchestrator, service, storage};

mod tls_config;

use config::database::DatabaseConfig;
use config::manager::ConfigManager;
use config::MLConfig;
use database::DatabaseManager;
use encryption::EncryptionKeyManager;
use gpu_config::GpuConfigManager;
use orchestrator::TrainingOrchestrator;
use service::{proto::ml_training_service_server::MlTrainingServiceServer, MLTrainingServiceImpl};
use storage::ModelStorageManager;
use tls_config::MLTrainingServiceTlsConfig;

/// ML Training Service CLI
#[derive(Parser)]
#[command(name = "ml_training_service")]
#[command(about = "ML Training Service for Foxhunt HFT Trading System")]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Start the ML training service
    Serve(ServeArgs),
    /// Health check
    Health(HealthArgs),
    /// Database operations
    Database(DatabaseArgs),
    /// Configuration validation
    Config(ConfigArgs),
}

#[derive(Args)]
struct ServeArgs {
    /// Configuration file path
    #[arg(short, long)]
    config: Option<String>,

    /// Override server port
    #[arg(short, long)]
    port: Option<u16>,

    /// Enable development mode with debug logging
    #[arg(long)]
    dev: bool,
}

#[derive(Args)]
struct HealthArgs {
    /// Service endpoint to check
    #[arg(long, default_value = "http://localhost:50053")]
    endpoint: String,
}

#[derive(Args)]
struct DatabaseArgs {
    #[command(subcommand)]
    action: DatabaseAction,
}

#[derive(Subcommand)]
enum DatabaseAction {
    /// Run database migrations
    Migrate,
    /// Check database health
    Health,
    /// Clean up old training jobs
    Cleanup {
        /// Days to retain completed jobs
        #[arg(long, default_value = "30")]
        retain_days: u32,
    },
}

#[derive(Args)]
struct ConfigArgs {
    /// Configuration file to validate
    #[arg(short, long)]
    file: Option<String>,
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();

    match cli.command {
        Commands::Serve(args) => serve(args).await,
        Commands::Health(args) => health_check(args).await,
        Commands::Database(args) => database_operations(args).await,
        Commands::Config(args) => config_operations(args).await,
    }
}

/// Start the ML training service
async fn serve(args: ServeArgs) -> Result<()> {
    // Initialize logging
    init_logging(args.dev)?;

    info!("Starting ML Training Service");

    // Load configuration using central config manager
    let service_config = config::ServiceConfig {
        name: "ml_training_service".to_string(),
        environment: std::env::var("ENVIRONMENT").unwrap_or_else(|_| "development".to_string()),
        version: env!("CARGO_PKG_VERSION").to_string(),
        settings: serde_json::json!({}),
    };
    let config_manager = ConfigManager::new(service_config);

    // Get ML training configuration - use defaults for now
    let ml_config = MLConfig::default();

    info!("ML Training configuration loaded from central config manager");

    // Get database configuration from environment
    let database_url = std::env::var("DATABASE_URL")
        .unwrap_or_else(|_| "postgresql://localhost/foxhunt".to_string());

    // HFT-optimized database configuration for ML training workloads
    // Wave 67 Agent 2: Updated pool configuration based on Wave 66 Agent 8 analysis
    let database_config = DatabaseConfig {
        url: database_url.clone(),
        max_connections: 20,  // Increased from 10 to support parallel training
        min_connections: 5,   // Increased from 1 for sustained throughput
        connect_timeout: std::time::Duration::from_secs(30),
        query_timeout: std::time::Duration::from_secs(60),
        enable_query_logging: false,
        application_name: Some("ml_training_service".to_string()),
        pool: config::PoolConfig {
            min_connections: 5,              // Increased from 1 to maintain warm connections
            max_connections: 20,             // Increased from 10 for parallel training jobs
            acquire_timeout_secs: 5,         // Reduced from 30s to 5s for ML training responsiveness
            max_lifetime_secs: 7200,         // Increased to 2 hours for long-running training
            idle_timeout_secs: 900,          // Increased to 15 minutes for training workloads
            test_before_acquire: true,
            database_url: database_url.clone(),
            health_check_enabled: true,
            health_check_interval_secs: 60,
        },
        transaction: config::TransactionConfig::default(),
    };

    info!("Configuration loaded and validated");

    // Default server address
    let grpc_port = std::env::var("GRPC_PORT")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(50053);
    let server_address = format!("0.0.0.0:{}", grpc_port);
    info!("Server will bind to: {}", server_address);

    // Make config_manager Arc for sharing
    let config_manager = Arc::new(config_manager);

    info!("ConfigManager initialized successfully");

    // Initialize GPU configuration manager
    let mut gpu_config_manager = GpuConfigManager::new(
        ml_config.training_config.clone(),
        Arc::clone(&config_manager),
    );

    // Load and validate GPU configuration
    match gpu_config_manager.load_config().await {
        Ok(gpu_config) => {
            info!(
                "GPU configuration loaded: device={}, memory={}GB",
                gpu_config.device_id, gpu_config.max_memory_gb
            );

            // Validate GPU availability
            match gpu_config_manager.validate_gpu_availability().await {
                Ok(validation) => {
                    if validation.is_ready_for_training() {
                        info!("GPU validation successful - ready for training");
                    } else {
                        warn!(
                            "GPU validation issues detected: {:?}",
                            validation.get_issues()
                        );
                        info!("Proceeding with training despite GPU issues");
                    }
                },
                Err(e) => {
                    warn!("GPU validation failed: {}", e);
                },
            }
        },
        Err(e) => {
            error!("Failed to load GPU configuration: {}", e);
            return Err(anyhow::anyhow!(
                "GPU configuration is required for ML training"
            ));
        },
    }

    // Initialize encryption key manager - use default encryption config
    let default_encryption_config = config::EncryptionConfig::default();
    let encryption_manager = EncryptionKeyManager::new(default_encryption_config.clone(), None);

    if encryption_manager.is_encryption_enabled() {
        match encryption_manager.load_encryption_keys().await {
            Ok(_) => {
                info!("Encryption keys loaded successfully");

                // Check if key rotation is needed
                match encryption_manager.keys_need_rotation().await {
                    Ok(needs_rotation) => {
                        if needs_rotation {
                            warn!("Encryption keys need rotation - schedule key rotation soon");
                        } else {
                            debug!("Encryption keys are current");
                        }
                    },
                    Err(e) => {
                        warn!("Failed to check key rotation status: {}", e);
                    },
                }
            },
            Err(e) => {
                if default_encryption_config.enable_encryption {
                    error!("Failed to load encryption keys: {}", e);
                    return Err(anyhow::anyhow!(
                        "Encryption keys are required when encryption is enabled"
                    ));
                } else {
                    warn!(
                        "Failed to load encryption keys (encryption disabled): {}",
                        e
                    );
                }
            },
        }
    } else {
        info!("Model encryption is disabled");
    }

    // Initialize database
    let database = Arc::new(
        DatabaseManager::new(&database_config)
            .await
            .context("Failed to initialize database")?,
    );

    info!("Database connection established");

    // Initialize storage with S3 configuration from environment
    let storage_config = config::storage_config::StorageConfig::from_env().unwrap_or_else(|e| {
        warn!(
            "Failed to load S3 config from environment: {}, using defaults",
            e
        );
        config::storage_config::StorageConfig::default()
    });

    info!(
        "Storage configuration: type={}, compression={}",
        storage_config.storage_type, storage_config.enable_compression
    );

    let storage = Arc::new(
        ModelStorageManager::new_with_config_manager(
            storage_config.into(), // Convert to storage::StorageConfig
            Arc::clone(&config_manager),
        )
        .await
        .context("Failed to initialize storage")?,
    );

    info!("Storage backend initialized with ConfigManager integration");

    // Initialize orchestrator
    let mut orchestrator = TrainingOrchestrator::new(
        ml_config.clone(),
        Arc::clone(&database),
        Arc::clone(&storage),
    )
    .await
    .context("Failed to initialize orchestrator")?;

    // Start orchestrator workers
    orchestrator
        .start()
        .await
        .context("Failed to start orchestrator")?;
    let orchestrator = Arc::new(orchestrator);

    info!("Training orchestrator started");

    // Initialize TLS configuration for mTLS
    let tls_config = MLTrainingServiceTlsConfig::from_config(&config_manager).await
        .context("Failed to initialize TLS configuration")?;

    info!("TLS configuration initialized with mutual TLS");

    // Create gRPC service
    let training_service = MLTrainingServiceImpl::new(Arc::clone(&orchestrator), ml_config.clone());

    // Build server with reflection
    let service = MlTrainingServiceServer::new(training_service);

    // Wave 67 Agent 3: HTTP/2 streaming performance optimizations
    // Apply same optimizations as Trading Service for consistency
    let enable_http2_opts = std::env::var("ENABLE_HTTP2_OPTIMIZATIONS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(true);

    let mut server = if enable_http2_opts {
        info!("✅ HTTP/2 optimizations enabled:");
        info!("   - tcp_nodelay: true (-40ms Nagle delay)");
        info!("   - Stream window: 1MB");
        info!("   - Connection window: 10MB");
        info!("   - Adaptive window: true");
        info!("   - Max streams: 1000");

        Server::builder()
            .tcp_nodelay(true)  // Critical: eliminates 40ms Nagle delay
            .tls_config(tls_config.to_server_tls_config())?
            .http2_keepalive_interval(Some(Duration::from_secs(30)))
            .http2_keepalive_timeout(Some(Duration::from_secs(10)))
            .initial_stream_window_size(Some(1024 * 1024))  // 1MB
            .initial_connection_window_size(Some(10 * 1024 * 1024))  // 10MB
            .http2_adaptive_window(Some(true))
            .max_concurrent_streams(Some(1000))
            .add_service(service)
    } else {
        info!("⚠️  HTTP/2 optimizations disabled via feature flag");
        Server::builder().tls_config(tls_config.to_server_tls_config())?.add_service(service)
    };

    // Add reflection service for development
    if args.dev {
        let reflection_service = ReflectionBuilder::configure()
            .build_v1alpha()
            .context("Failed to build reflection service")?;

        server = server.add_service(reflection_service);
        info!("gRPC reflection enabled for development");
    }

    // Start the server
    let server = server.serve(server_address.parse::<SocketAddr>()?);

    // Skip metrics server for now - would need proper monitoring config
    let _metrics_handle: Option<tokio::task::JoinHandle<()>> = None;

    info!("ML Training Service ready");
    info!("gRPC server listening on {}", server_address);

    // Run server
    if let Err(e) = server.await {
        error!("Server error: {}", e);
        return Err(e.into());
    }

    info!("ML Training Service stopped");
    Ok(())
}

/// Start Prometheus metrics server
#[allow(dead_code)]
async fn start_metrics_server(_config: MLConfig) -> Result<tokio::task::JoinHandle<()>> {
    use std::time::Duration;

    let metrics_handle = tokio::spawn(async move {
        // Keep the metrics server running
        loop {
            tokio::time::sleep(Duration::from_secs(60)).await;
        }
    });

    info!("Prometheus metrics server started");
    Ok(metrics_handle)
}

/// Perform health check
async fn health_check(args: HealthArgs) -> Result<()> {
    use service::proto::{ml_training_service_client::MlTrainingServiceClient, HealthCheckRequest};

    println!("Checking service health at: {}", args.endpoint);

    let mut client = MlTrainingServiceClient::connect(args.endpoint)
        .await
        .context("Failed to connect to service")?;

    let request = tonic::Request::new(HealthCheckRequest {});

    let response = client
        .health_check(request)
        .await
        .context("Health check failed")?;

    let health = response.into_inner();

    if health.healthy {
        println!("✅ Service is healthy: {}", health.message);
        for (key, value) in health.details {
            println!("   {}: {}", key, value);
        }
        Ok(())
    } else {
        println!("❌ Service is unhealthy: {}", health.message);
        std::process::exit(1);
    }
}

/// Database operations
async fn database_operations(args: DatabaseArgs) -> Result<()> {
    init_logging(false)?;

    // Get database configuration from environment
    let database_url = std::env::var("DATABASE_URL")
        .unwrap_or_else(|_| "postgresql://localhost/foxhunt".to_string());
    // HFT-optimized database configuration for ML training database operations
    // Wave 67 Agent 2: Updated pool configuration based on Wave 66 Agent 8 analysis
    let database_config = DatabaseConfig {
        url: database_url.clone(),
        max_connections: 20,  // Increased from 10 to support parallel operations
        min_connections: 5,   // Increased from 1 for sustained throughput
        connect_timeout: std::time::Duration::from_secs(30),
        query_timeout: std::time::Duration::from_secs(60),
        enable_query_logging: false,
        application_name: Some("ml_training_service".to_string()),
        pool: config::PoolConfig {
            min_connections: 5,              // Increased from 1 to maintain warm connections
            max_connections: 20,             // Increased from 10 for parallel database operations
            acquire_timeout_secs: 5,         // Reduced from 30s to 5s for ML training responsiveness
            max_lifetime_secs: 7200,         // Increased to 2 hours for long-running operations
            idle_timeout_secs: 900,          // Increased to 15 minutes for training workloads
            test_before_acquire: true,
            database_url: database_url.clone(),
            health_check_enabled: true,
            health_check_interval_secs: 60,
        },
        transaction: config::TransactionConfig::default(),
    };
    let database = DatabaseManager::new(&database_config)
        .await
        .context("Failed to connect to database")?;

    match args.action {
        DatabaseAction::Migrate => {
            info!("Running database migrations");
            database.run_migrations().await?;
            println!("✅ Database migrations completed successfully");
        },

        DatabaseAction::Health => {
            info!("Checking database health");
            database.health_check().await?;
            println!("✅ Database is healthy");
        },

        DatabaseAction::Cleanup { retain_days } => {
            info!(
                "Cleaning up old training jobs (retain {} days)",
                retain_days
            );
            // Would implement cleanup logic
            println!("✅ Database cleanup completed");
        },
    }

    Ok(())
}

/// Configuration operations
async fn config_operations(_args: ConfigArgs) -> Result<()> {
    println!("Validating configuration...");

    // Get ML config defaults
    let _ml_config = MLConfig::default();

    // Get database configuration from environment
    let database_url = std::env::var("DATABASE_URL")
        .unwrap_or_else(|_| "postgresql://localhost/foxhunt".to_string());

    println!("✅ Configuration is valid");
    println!("\nConfiguration summary:");
    println!("  Server: 0.0.0.0:50053");
    println!("  Database URL: {}", database_url);
    println!("  ML Config: Using defaults");

    Ok(())
}

/// Initialize logging
fn init_logging(dev_mode: bool) -> Result<()> {
    let log_level = if dev_mode { "debug" } else { "info" };

    let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(log_level));

    tracing_subscriber::registry()
        .with(env_filter)
        .with(tracing_subscriber::fmt::layer())
        .init();

    Ok(())
}

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

    #[test]
    fn test_cli_parsing() {
        // Test basic serve command
        let cli = Cli::parse_from(&["ml_training_service", "serve"]);
        matches!(cli.command, Commands::Serve(_));

        // Test serve with options
        let cli = Cli::parse_from(&["ml_training_service", "serve", "--port", "8080", "--dev"]);
        if let Commands::Serve(args) = cli.command {
            assert_eq!(args.port, Some(8080));
            assert!(args.dev);
        }

        // Test health check
        let cli = Cli::parse_from(&["ml_training_service", "health"]);
        matches!(cli.command, Commands::Health(_));

        // Test database operations
        let cli = Cli::parse_from(&["ml_training_service", "database", "migrate"]);
        matches!(cli.command, Commands::Database(_));
    }

    #[test]
    fn test_config_validation() {
        let config = MLConfig::default();
        // Basic validation test - config should have sensible defaults
        assert!(!config.model_config.model_type.is_empty());
    }
}