foxhunt/services/api/src/grpc/ml_training_proxy.rs

//! ML Training Service Proxy - Zero-copy gRPC forwarding for ML training operations
//!
//! This module implements a high-performance proxy for the ML Training Service with:
//! - Zero-copy message forwarding (routing overhead <10μs)
//! - Connection pooling via tonic::transport::Channel
//! - Circuit breaker integration for backend failures
//! - Efficient streaming support for training metrics
//! - Health checking integration

use futures::Stream;
use std::pin::Pin;
use tonic::{Request, Response, Status};
use tracing::{error, info, instrument, warn};

// Import the generated ML training service protobuf definitions from lib.rs
use crate::ml_training::ml_training_service_client::MlTrainingServiceClient;
use crate::ml_training::ml_training_service_server::{MlTrainingService, MlTrainingServiceServer};
use crate::ml_training::{
    // Batch tuning types
    BatchStartTuningJobsRequest,
    BatchStartTuningJobsResponse,
    GetBatchTuningStatusRequest,
    GetBatchTuningStatusResponse,
    GetTrainingJobDetailsRequest,
    GetTrainingJobDetailsResponse,
    GetTuningJobStatusRequest,
    GetTuningJobStatusResponse,
    HealthCheckRequest,
    HealthCheckResponse,
    ListAvailableModelsRequest,
    ListAvailableModelsResponse,
    ListTrainingJobsRequest,
    ListTrainingJobsResponse,
    ProgressUpdate,
    StartTrainingRequest,
    StartTrainingResponse,
    // Tuning job types
    StartTuningJobRequest,
    StartTuningJobResponse,
    StopBatchTuningJobRequest,
    StopBatchTuningJobResponse,
    StopTrainingRequest,
    StopTrainingResponse,
    StopTuningJobRequest,
    StopTuningJobResponse,
    StreamProgressRequest,
    SubscribeToTrainingStatusRequest,
    TrainModelRequest,
    TrainModelResponse,
    TrainingStatusUpdate,
    // Job completion and model promotion types
    JobCompletionReport,
    JobCompletionAck,
    ListPendingPromotionsRequest,
    ListPendingPromotionsResponse,
    ApprovePromotionRequest,
    ApprovePromotionResponse,
    RejectPromotionRequest,
    RejectPromotionResponse,
    ApproveModelRequest,
    ApproveModelResponse,
    RejectModelRequest,
    RejectModelResponse,
};

/// ML Training Service Proxy
///
/// Provides zero-copy forwarding of gRPC requests to the backend ML Training Service.
///
/// Uses connection pooling and circuit breakers for high availability and performance.
#[derive(Debug, Clone)]
pub struct MlTrainingProxy {
    /// Backend ML Training Service client with connection pooling
    client: MlTrainingServiceClient<tonic::transport::Channel>,
}

impl MlTrainingProxy {
    /// Create a new ML Training Service proxy
    ///
    /// # Arguments
    /// * `client` - Pre-configured ML Training Service client with circuit breaker
    ///
    /// # Performance
    /// - Uses Arc-based channel cloning for zero-copy client reuse
    ///
    /// - Connection pooling managed by tonic::transport::Channel
    pub fn new(client: MlTrainingServiceClient<tonic::transport::Channel>) -> Self {
        Self { client }
    }

    /// Convert proxy into a tonic server instance
    pub fn into_server(self) -> MlTrainingServiceServer<MlTrainingProxy> {
        MlTrainingServiceServer::new(self)
    }
}

#[tonic::async_trait]
impl MlTrainingService for MlTrainingProxy {
    /// Server streaming type for training status updates
    type SubscribeToTrainingStatusStream =
        Pin<Box<dyn Stream<Item = Result<TrainingStatusUpdate, Status>> + Send>>;

    /// Server streaming type for tuning progress updates
    type StreamTuningProgressStream =
        Pin<Box<dyn Stream<Item = Result<ProgressUpdate, Status>> + Send>>;

    /// Start a new model training job
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    ///
    /// - Routing overhead target: <10μs
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn start_training(
        &self,
        request: Request<StartTrainingRequest>,
    ) -> Result<Response<StartTrainingResponse>, Status> {
        info!("Proxying StartTraining request");

        // Clone client (cheap Arc increment) for concurrent request handling
        let mut client = self.client.clone();

        // Forward request with zero-copy
        let response = client.start_training(request).await.map_err(|e| {
            error!("Backend StartTraining failed: {}", e);
            e
        })?;

        info!("StartTraining request forwarded successfully");
        Ok(response)
    }

    /// Subscribe to real-time training status updates (server streaming)
    ///
    /// # Performance
    /// - Zero-copy stream forwarding
    ///
    /// - No intermediate buffering
    /// - Direct stream passthrough from backend
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn subscribe_to_training_status(
        &self,
        request: Request<SubscribeToTrainingStatusRequest>,
    ) -> Result<Response<Self::SubscribeToTrainingStatusStream>, Status> {
        info!("Proxying SubscribeToTrainingStatus streaming request");

        let mut client = self.client.clone();

        // Get backend stream response
        let stream_response = client
            .subscribe_to_training_status(request)
            .await
            .map_err(|e| {
                error!("Backend SubscribeToTrainingStatus failed: {}", e);
                e
            })?;

        // Extract inner stream and forward directly (zero-copy)
        let stream = stream_response.into_inner();
        let boxed_stream = Box::pin(stream) as Self::SubscribeToTrainingStatusStream;

        info!("SubscribeToTrainingStatus streaming request forwarded successfully");
        Ok(Response::new(boxed_stream))
    }

    /// Stop a running training job
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    ///
    /// - Routing overhead target: <10μs
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn stop_training(
        &self,
        request: Request<StopTrainingRequest>,
    ) -> Result<Response<StopTrainingResponse>, Status> {
        info!("Proxying StopTraining request");

        let mut client = self.client.clone();
        let response = client.stop_training(request).await.map_err(|e| {
            error!("Backend StopTraining failed: {}", e);
            e
        })?;

        info!("StopTraining request forwarded successfully");
        Ok(response)
    }

    /// List available ML models
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    ///
    /// - Routing overhead target: <10μs
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn list_available_models(
        &self,
        request: Request<ListAvailableModelsRequest>,
    ) -> Result<Response<ListAvailableModelsResponse>, Status> {
        info!("Proxying ListAvailableModels request");

        let mut client = self.client.clone();
        let response = client.list_available_models(request).await.map_err(|e| {
            error!("Backend ListAvailableModels failed: {}", e);
            e
        })?;

        info!("ListAvailableModels request forwarded successfully");
        Ok(response)
    }

    /// List training job history
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    ///
    /// - Routing overhead target: <10μs
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn list_training_jobs(
        &self,
        request: Request<ListTrainingJobsRequest>,
    ) -> Result<Response<ListTrainingJobsResponse>, Status> {
        info!("Proxying ListTrainingJobs request");

        let mut client = self.client.clone();
        let response = client.list_training_jobs(request).await.map_err(|e| {
            error!("Backend ListTrainingJobs failed: {}", e);
            e
        })?;

        info!("ListTrainingJobs request forwarded successfully");
        Ok(response)
    }

    /// Get detailed information about a training job
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    ///
    /// - Routing overhead target: <10μs
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn get_training_job_details(
        &self,
        request: Request<GetTrainingJobDetailsRequest>,
    ) -> Result<Response<GetTrainingJobDetailsResponse>, Status> {
        info!("Proxying GetTrainingJobDetails request");

        let mut client = self.client.clone();
        let response = client
            .get_training_job_details(request)
            .await
            .map_err(|e| {
                error!("Backend GetTrainingJobDetails failed: {}", e);
                e
            })?;

        info!("GetTrainingJobDetails request forwarded successfully");
        Ok(response)
    }

    /// Health check for ML Training Service backend
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    ///
    /// - Routing overhead target: <10μs
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn health_check(
        &self,
        request: Request<HealthCheckRequest>,
    ) -> Result<Response<HealthCheckResponse>, Status> {
        info!("Proxying HealthCheck request");

        let mut client = self.client.clone();
        let response = client.health_check(request).await.map_err(|e| {
            warn!("Backend HealthCheck failed: {}", e);
            e
        })?;

        info!("HealthCheck request forwarded successfully");
        Ok(response)
    }

    /// Start a new hyperparameter tuning job
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    ///
    /// - Routing overhead target: <10μs
    ///
    /// # Security
    /// - Requires "ml.tune" permission in JWT metadata
    /// - JWT validation handled by interceptor layer
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn start_tuning_job(
        &self,
        request: Request<StartTuningJobRequest>,
    ) -> Result<Response<StartTuningJobResponse>, Status> {
        info!("Proxying StartTuningJob request");

        // Clone client (cheap Arc increment) for concurrent request handling
        let mut client = self.client.clone();

        // Forward request with zero-copy
        // Note: JWT validation and "ml.tune" permission check handled by interceptor
        let response = client.start_tuning_job(request).await.map_err(|e| {
            error!("Backend StartTuningJob failed: {}", e);
            e
        })?;

        info!("StartTuningJob request forwarded successfully");
        Ok(response)
    }

    /// Get status and best parameters from a tuning job
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    ///
    /// - Routing overhead target: <10μs
    ///
    /// # Security
    /// - Job ownership validation handled by backend service
    /// - User can only query their own tuning jobs
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn get_tuning_job_status(
        &self,
        request: Request<GetTuningJobStatusRequest>,
    ) -> Result<Response<GetTuningJobStatusResponse>, Status> {
        info!("Proxying GetTuningJobStatus request");

        let mut client = self.client.clone();

        // Forward request - backend validates job ownership via user_id from JWT
        let response = client.get_tuning_job_status(request).await.map_err(|e| {
            error!("Backend GetTuningJobStatus failed: {}", e);
            e
        })?;

        info!("GetTuningJobStatus request forwarded successfully");
        Ok(response)
    }

    /// Stop a running hyperparameter tuning job
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    ///
    /// - Routing overhead target: <10μs
    ///
    /// # Security
    /// - Job ownership validation handled by backend service
    /// - User can only stop their own tuning jobs
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn stop_tuning_job(
        &self,
        request: Request<StopTuningJobRequest>,
    ) -> Result<Response<StopTuningJobResponse>, Status> {
        info!("Proxying StopTuningJob request");

        let mut client = self.client.clone();

        // Forward stop request - backend validates job ownership via user_id from JWT
        let response = client.stop_tuning_job(request).await.map_err(|e| {
            error!("Backend StopTuningJob failed: {}", e);
            e
        })?;

        info!("StopTuningJob request forwarded successfully");
        Ok(response)
    }

    /// INTERNAL: Train a single model instance with specific hyperparameters
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    ///
    /// - Routing overhead target: <10μs
    ///
    /// # Note
    /// - Called by Optuna subprocess during tuning trials
    /// - Not typically called directly by clients
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn train_model(
        &self,
        request: Request<TrainModelRequest>,
    ) -> Result<Response<TrainModelResponse>, Status> {
        info!("Proxying TrainModel request (INTERNAL)");

        let mut client = self.client.clone();

        // Forward internal training request
        let response = client.train_model(request).await.map_err(|e| {
            error!("Backend TrainModel failed: {}", e);
            e
        })?;

        info!("TrainModel request forwarded successfully");
        Ok(response)
    }

    /// Stream real-time tuning progress updates (server streaming)
    ///
    /// # Performance
    /// - Zero-copy stream forwarding
    ///
    /// - No intermediate buffering
    /// - Direct stream passthrough from backend
    ///
    /// # Security
    /// - Tuning job ownership validation handled by backend service
    /// - User can only subscribe to their own tuning jobs
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn stream_tuning_progress(
        &self,
        request: Request<StreamProgressRequest>,
    ) -> Result<Response<Self::StreamTuningProgressStream>, Status> {
        info!("Proxying StreamTuningProgress streaming request");

        let mut client = self.client.clone();

        // Get backend stream response
        let stream_response = client.stream_tuning_progress(request).await.map_err(|e| {
            error!("Backend StreamTuningProgress failed: {}", e);
            e
        })?;

        // Extract inner stream and forward directly (zero-copy)
        let stream = stream_response.into_inner();
        let boxed_stream = Box::pin(stream) as Self::StreamTuningProgressStream;

        info!("StreamTuningProgress streaming request forwarded successfully");
        Ok(Response::new(boxed_stream))
    }

    /// Start batch tuning for multiple models with automatic dependency resolution
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    /// - Routing overhead target: <10μs
    ///
    /// # Security
    /// - Requires "ml.tune" permission in JWT metadata
    /// - JWT validation handled by interceptor layer
    ///
    /// # Features
    /// - Sequential model tuning with dependency resolution
    /// - Automatic YAML export of best hyperparameters
    /// - Supports all model types (DQN, PPO, MAMBA_2, TFT, etc.)
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn batch_start_tuning_jobs(
        &self,
        request: Request<BatchStartTuningJobsRequest>,
    ) -> Result<Response<BatchStartTuningJobsResponse>, Status> {
        info!("Proxying BatchStartTuningJobs request");

        // Clone client (cheap Arc increment) for concurrent request handling
        let mut client = self.client.clone();

        // Forward batch tuning request with zero-copy
        // Note: JWT validation and "ml.tune" permission check handled by interceptor
        let response = client.batch_start_tuning_jobs(request).await.map_err(|e| {
            error!("Backend BatchStartTuningJobs failed: {}", e);
            e
        })?;

        info!("BatchStartTuningJobs request forwarded successfully");
        Ok(response)
    }

    /// Get batch tuning job status with per-model results
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    /// - Routing overhead target: <10μs
    ///
    /// # Security
    /// - Batch job ownership validation handled by backend service
    /// - User can only query their own batch jobs
    ///
    /// # Returns
    /// - Batch status (PENDING, RUNNING, COMPLETED, FAILED, STOPPED)
    /// - Per-model tuning results
    /// - Current execution progress
    /// - Estimated completion time
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn get_batch_tuning_status(
        &self,
        request: Request<GetBatchTuningStatusRequest>,
    ) -> Result<Response<GetBatchTuningStatusResponse>, Status> {
        info!("Proxying GetBatchTuningStatus request");

        let mut client = self.client.clone();

        // Forward request - backend validates batch job ownership via user_id from JWT
        let response = client.get_batch_tuning_status(request).await.map_err(|e| {
            error!("Backend GetBatchTuningStatus failed: {}", e);
            e
        })?;

        info!("GetBatchTuningStatus request forwarded successfully");
        Ok(response)
    }

    /// Stop a running batch tuning job
    ///
    /// # Performance
    /// - Zero-copy message forwarding
    /// - Routing overhead target: <10μs
    ///
    /// # Security
    /// - Batch job ownership validation handled by backend service
    /// - User can only stop their own batch jobs
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn stop_batch_tuning_job(
        &self,
        request: Request<StopBatchTuningJobRequest>,
    ) -> Result<Response<StopBatchTuningJobResponse>, Status> {
        info!("Proxying StopBatchTuningJob request");

        let mut client = self.client.clone();

        // Forward stop request - backend validates batch job ownership via user_id from JWT
        let response = client.stop_batch_tuning_job(request).await.map_err(|e| {
            error!("Backend StopBatchTuningJob failed: {}", e);
            e
        })?;

        info!("StopBatchTuningJob request forwarded successfully");
        Ok(response)
    }

    /// Report job completion (called by training-uploader sidecar)
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn report_job_completion(
        &self,
        request: Request<JobCompletionReport>,
    ) -> Result<Response<JobCompletionAck>, Status> {
        info!("Proxying ReportJobCompletion request");
        let mut client = self.client.clone();
        let response = client.report_job_completion(request).await.map_err(|e| {
            error!("Backend ReportJobCompletion failed: {}", e);
            e
        })?;
        Ok(response)
    }

    /// List models pending promotion approval
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn list_pending_promotions(
        &self,
        request: Request<ListPendingPromotionsRequest>,
    ) -> Result<Response<ListPendingPromotionsResponse>, Status> {
        info!("Proxying ListPendingPromotions request");
        let mut client = self.client.clone();
        let response = client.list_pending_promotions(request).await.map_err(|e| {
            error!("Backend ListPendingPromotions failed: {}", e);
            e
        })?;
        Ok(response)
    }

    /// Approve a pending model promotion
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn approve_promotion(
        &self,
        request: Request<ApprovePromotionRequest>,
    ) -> Result<Response<ApprovePromotionResponse>, Status> {
        info!("Proxying ApprovePromotion request");
        let mut client = self.client.clone();
        let response = client.approve_promotion(request).await.map_err(|e| {
            error!("Backend ApprovePromotion failed: {}", e);
            e
        })?;
        Ok(response)
    }

    /// Reject a pending model promotion
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn reject_promotion(
        &self,
        request: Request<RejectPromotionRequest>,
    ) -> Result<Response<RejectPromotionResponse>, Status> {
        info!("Proxying RejectPromotion request");
        let mut client = self.client.clone();
        let response = client.reject_promotion(request).await.map_err(|e| {
            error!("Backend RejectPromotion failed: {}", e);
            e
        })?;
        Ok(response)
    }

    /// Approve a model for promotion
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn approve_model(
        &self,
        request: Request<ApproveModelRequest>,
    ) -> Result<Response<ApproveModelResponse>, Status> {
        info!("Proxying ApproveModel request");
        let mut client = self.client.clone();
        let response = client.approve_model(request).await.map_err(|e| {
            error!("Backend ApproveModel failed: {}", e);
            e
        })?;
        Ok(response)
    }

    /// Reject a model promotion
    #[instrument(skip(self, request), fields(request_id = %uuid::Uuid::new_v4()), err)]
    async fn reject_model(
        &self,
        request: Request<RejectModelRequest>,
    ) -> Result<Response<RejectModelResponse>, Status> {
        info!("Proxying RejectModel request");
        let mut client = self.client.clone();
        let response = client.reject_model(request).await.map_err(|e| {
            error!("Backend RejectModel failed: {}", e);
            e
        })?;
        Ok(response)
    }
}

#[cfg(test)]
mod tests {

    #[test]
    fn test_proxy_creation() {
        // This test validates the proxy struct can be created
        // Full integration tests require running backend service
    }
}