//! gRPC Service Implementation
//!
//! This module implements the MLTrainingService gRPC interface,
//! providing the external API for training job management.

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;

use anyhow::Result;
use tokio::time::{timeout, Duration};
use tokio_stream::Stream;
use tonic::{Request, Response, Status};
use tracing::{debug, error, info, warn, instrument};
use uuid::Uuid;

// Import the generated protobuf types
pub mod proto {
    pub mod ml_training {
        tonic::include_proto!("ml_training");
    }
    // Re-export for convenience
    pub use ml_training::*;
}

use proto::{
    ml_training_service_server::MlTrainingService, DqnParams,
    FinancialMetrics as ProtoFinancialMetrics, GetTrainingJobDetailsRequest,
    GetTrainingJobDetailsResponse, HealthCheckRequest, HealthCheckResponse, Hyperparameters,
    LiquidParams, ListAvailableModelsRequest, ListAvailableModelsResponse, ListTrainingJobsRequest,
    ListTrainingJobsResponse, MambaParams, ModelDefinition, PpoParams,
    ResourceUsage as ProtoResourceUsage, StartTrainingRequest, StartTrainingResponse,
    StopTrainingRequest, StopTrainingResponse, SubscribeToTrainingStatusRequest, TftParams,
    TlobParams, TrainingJobDetails, TrainingJobSummary, TrainingStatus as ProtoTrainingStatus,
    TrainingStatusUpdate as ProtoStatusUpdate,
    // Job completion & model promotion (on-demand training pipeline)
    JobCompletionReport, JobCompletionAck,
    ListPendingPromotionsRequest, ListPendingPromotionsResponse, PendingPromotion,
    ApprovePromotionRequest, ApprovePromotionResponse,
    RejectPromotionRequest, RejectPromotionResponse,
    // Dashboard-facing model approval / rejection
    ApproveModelRequest, ApproveModelResponse,
    RejectModelRequest, RejectModelResponse,
};

use crate::batch_tuning_manager::{BatchJobStatus, BatchTuningManager};
use crate::grpc_tuning_handlers::TuningHandlers;
use crate::orchestrator::{JobStatus, TrainingOrchestrator, TrainingStatusUpdate};
use crate::tuning_manager::{TuningJobStatus, TuningManager};
use config::MLConfig;
use ml::training_pipeline::{
    ModelArchitectureConfig, ProductionTrainingConfig, TrainingHyperparameters,
};

/// gRPC service implementation
pub struct MLTrainingServiceImpl {
    orchestrator: Arc<TrainingOrchestrator>,
    tuning_handlers: Arc<TuningHandlers>,
    batch_tuning_manager: Arc<BatchTuningManager>,
    /// Tracks pending model promotions and active model registry.
    pub promotion_manager: Arc<crate::promotion_manager::PromotionManager>,
    /// K8s job dispatcher -- None when running outside a K8s cluster.
    pub k8s_dispatcher: Option<Arc<crate::k8s_dispatcher::K8sDispatcher>>,
    /// Job spawner for persisting training jobs to PostgreSQL before K8s dispatch.
    /// None when no database pool is available.
    pub job_spawner: Option<Arc<crate::job_spawner::JobSpawner>>,
}

impl MLTrainingServiceImpl {
    /// Create a new service instance
    pub fn new(
        orchestrator: Arc<TrainingOrchestrator>,
        tuning_manager: Arc<TuningManager>,
        _config: MLConfig,
        promotion_manager: Arc<crate::promotion_manager::PromotionManager>,
        k8s_dispatcher: Option<Arc<crate::k8s_dispatcher::K8sDispatcher>>,
        job_spawner: Option<Arc<crate::job_spawner::JobSpawner>>,
    ) -> Self {
        let tuning_handlers = Arc::new(TuningHandlers::new(Arc::clone(&tuning_manager)));
        let working_dir = std::env::var("TUNING_WORKING_DIR").unwrap_or_else(|_| ".".to_string());
        let batch_tuning_manager = Arc::new(BatchTuningManager::new(tuning_manager, working_dir));
        Self {
            orchestrator,
            tuning_handlers,
            batch_tuning_manager,
            promotion_manager,
            k8s_dispatcher,
            job_spawner,
        }
    }

    /// Convert internal job status to protobuf status
    fn convert_job_status(status: &JobStatus) -> ProtoTrainingStatus {
        match status {
            JobStatus::Pending => ProtoTrainingStatus::Pending,
            JobStatus::Running => ProtoTrainingStatus::Running,
            JobStatus::Completed => ProtoTrainingStatus::Completed,
            JobStatus::Failed => ProtoTrainingStatus::Failed,
            JobStatus::Stopped => ProtoTrainingStatus::Stopped,
            JobStatus::Paused => ProtoTrainingStatus::Paused,
        }
    }

    /// Convert internal batch job status to protobuf batch status
    fn convert_batch_status(status: &BatchJobStatus) -> proto::BatchTuningStatus {
        match status {
            BatchJobStatus::Pending => proto::BatchTuningStatus::BatchPending,
            BatchJobStatus::Running => proto::BatchTuningStatus::BatchRunning,
            BatchJobStatus::Completed => proto::BatchTuningStatus::BatchCompleted,
            BatchJobStatus::PartiallyCompleted => {
                proto::BatchTuningStatus::BatchPartiallyCompleted
            },
            BatchJobStatus::Failed => proto::BatchTuningStatus::BatchFailed,
            BatchJobStatus::Stopped => proto::BatchTuningStatus::BatchStopped,
        }
    }

    /// Convert internal tuning job status to protobuf tuning status
    fn convert_tuning_job_status(status: &TuningJobStatus) -> proto::TuningJobStatus {
        match status {
            TuningJobStatus::Pending => proto::TuningJobStatus::TuningPending,
            TuningJobStatus::Running => proto::TuningJobStatus::TuningRunning,
            TuningJobStatus::Completed => proto::TuningJobStatus::TuningCompleted,
            TuningJobStatus::Failed => proto::TuningJobStatus::TuningFailed,
            TuningJobStatus::Stopped => proto::TuningJobStatus::TuningStopped,
        }
    }

    /// Convert protobuf hyperparameters to internal config
    fn convert_hyperparameters(
        &self,
        params: Option<Hyperparameters>,
    ) -> Result<ProductionTrainingConfig> {
        let mut config = ProductionTrainingConfig::default();

        if let Some(hyperparams) = params {
            match hyperparams.model_params {
                Some(proto::hyperparameters::ModelParams::TlobParams(tlob)) => {
                    config.model_config = ModelArchitectureConfig {
                        input_dim: 20, // Default for TLOB
                        hidden_dims: vec![
                            usize::try_from(tlob.hidden_dim).unwrap_or(256),
                            usize::try_from(tlob.hidden_dim).unwrap_or(256) / 2,
                        ],
                        output_dim: 1,
                        dropout_rate: tlob.dropout_rate as f64,
                        activation: "relu".to_string(),
                        batch_norm: true,
                        residual_connections: false,
                    };

                    config.training_params = TrainingHyperparameters {
                        learning_rate: tlob.learning_rate as f64,
                        batch_size: usize::try_from(tlob.batch_size).unwrap_or(64),
                        max_epochs: usize::try_from(tlob.epochs).unwrap_or(100),
                        patience: 50,
                        validation_split: 0.2,
                        l2_regularization: 1e-4,
                        lr_decay_factor: 0.5,
                        lr_decay_patience: 25,
                    };
                },

                Some(proto::hyperparameters::ModelParams::MambaParams(mamba)) => {
                    config.model_config = ModelArchitectureConfig {
                        input_dim: usize::try_from(mamba.state_dim).unwrap_or(128),
                        hidden_dims: vec![usize::try_from(mamba.hidden_dim).unwrap_or(512)],
                        output_dim: 1,
                        dropout_rate: 0.1,
                        activation: "relu".to_string(),
                        batch_norm: true,
                        residual_connections: false,
                    };

                    config.training_params = TrainingHyperparameters {
                        learning_rate: mamba.learning_rate as f64,
                        batch_size: usize::try_from(mamba.batch_size).unwrap_or(32),
                        max_epochs: usize::try_from(mamba.epochs).unwrap_or(150),
                        patience: 50,
                        validation_split: 0.2,
                        l2_regularization: 1e-4,
                        lr_decay_factor: 0.5,
                        lr_decay_patience: 25,
                    };
                },

                Some(proto::hyperparameters::ModelParams::DqnParams(dqn)) => {
                    config.training_params = TrainingHyperparameters {
                        learning_rate: dqn.learning_rate as f64,
                        batch_size: usize::try_from(dqn.batch_size).unwrap_or(128),
                        max_epochs: usize::try_from(dqn.epochs).unwrap_or(200),
                        patience: 100,
                        validation_split: 0.1,
                        l2_regularization: 1e-5,
                        lr_decay_factor: 0.8,
                        lr_decay_patience: 50,
                    };
                },

                // Add other model parameter conversions...
                _ => {
                    // Use default configuration
                },
            }
        }

        // Apply GPU preference
        if config.performance_config.device_preference == "cpu" {
            config.performance_config.device_preference = "cuda".to_string();
        }

        Ok(config)
    }

    /// Convert internal status update to protobuf
    fn convert_status_update(update: TrainingStatusUpdate) -> ProtoStatusUpdate {
        let financial_metrics = update.financial_metrics.map(|fm| ProtoFinancialMetrics {
            simulated_return: fm.simulated_return as f32,
            sharpe_ratio: fm.sharpe_ratio as f32,
            max_drawdown: fm.max_drawdown as f32,
            hit_rate: fm.hit_rate as f32,
            avg_prediction_error_bps: fm.avg_prediction_error_bps as f32,
            risk_adjusted_return: fm.risk_adjusted_return as f32,
            var_5pct: 0.0, // Would be populated from actual metrics
            expected_shortfall: 0.0,
        });

        let resource_usage = ProtoResourceUsage {
            cpu_usage_percent: update.resource_usage.cpu_usage_percent,
            memory_usage_gb: update.resource_usage.memory_usage_gb,
            gpu_usage_percent: update.resource_usage.gpu_usage_percent.unwrap_or(0.0),
            gpu_memory_usage_gb: update.resource_usage.gpu_memory_usage_gb.unwrap_or(0.0),
            active_workers: update.resource_usage.active_workers,
        };

        ProtoStatusUpdate {
            job_id: update.job_id.to_string(),
            status: (Self::convert_job_status(&update.status) as i32),
            progress_percentage: update.progress_percentage,
            current_epoch: update.current_epoch,
            total_epochs: update.total_epochs,
            metrics: update
                .metrics
                .into_iter()
                .map(|(k, v)| (k, v as f32))
                .collect(),
            message: update.message,
            timestamp: update.timestamp.timestamp(),
            financial_metrics,
            resource_usage: Some(resource_usage),
        }
    }
}

#[tonic::async_trait]
impl MlTrainingService for MLTrainingServiceImpl {
    /// Start a new training job
    #[instrument(skip_all)]
    async fn start_training(
        &self,
        request: Request<StartTrainingRequest>,
    ) -> Result<Response<StartTrainingResponse>, Status> {
        let req = request.into_inner();

        info!("Starting training job for model type: {}", req.model_type);

        // Detect hyperopt mode from tags
        let is_hyperopt = req.tags.get("mode").map(|m| m == "hyperopt").unwrap_or(false);

        // Determine the correct binary
        let binary = if is_hyperopt {
            crate::k8s_dispatcher::hyperopt_binary_for_model(&req.model_type)
        } else {
            crate::k8s_dispatcher::training_binary_for_model(&req.model_type)
        };

        if is_hyperopt {
            info!("Hyperopt mode: using binary {}", binary);
        }

        // Generate a job ID
        let job_id = Uuid::new_v4();

        // Extract symbol from data_source file_path (take last path component), or use default
        let symbol = req
            .data_source
            .as_ref()
            .and_then(|ds| match &ds.source {
                Some(proto::data_source::Source::FilePath(p)) => {
                    // file_path may be "data/cache/futures-baseline/ES.FUT" — extract just the symbol
                    std::path::Path::new(p)
                        .file_name()
                        .and_then(|f| f.to_str())
                        .map(String::from)
                }
                _ => None,
            })
            .unwrap_or_else(|| "ES.FUT".to_string());

        // Priority 1: Batch flow — persist job to DB, queue consumer dispatches to K8s
        if let Some(ref spawner) = self.job_spawner {
            let model = common::model_types::ModelType::from_str(&req.model_type).ok_or_else(
                || Status::invalid_argument(format!("Unknown model type: {}", req.model_type)),
            )?;

            let asset = crate::job_spawner::Asset {
                symbol: symbol.clone(),
                data_file: std::path::PathBuf::from(format!(
                    "/data/futures-baseline/{}",
                    symbol
                )),
            };

            match spawner.spawn_batch(vec![asset], vec![model]).await {
                Ok(batch) => {
                    info!(batch_id = %batch.batch_id, "Batch job created in database");
                    return Ok(Response::new(StartTrainingResponse {
                        job_id: batch.batch_id.to_string(),
                        status: ProtoTrainingStatus::Pending as i32,
                        message: format!("Batch training job queued: {}", batch.batch_id),
                    }));
                }
                Err(e) => {
                    warn!("JobSpawner failed, falling back to K8s/orchestrator: {}", e);
                    // Fall through to K8s dispatcher or in-process orchestrator
                }
            }
        }

        // Priority 2: Direct K8s dispatch
        if let Some(ref dispatcher) = self.k8s_dispatcher {
            let epochs = req
                .hyperparameters
                .as_ref()
                .and_then(|hp| match &hp.model_params {
                    Some(proto::hyperparameters::ModelParams::TlobParams(p)) => {
                        Some(p.epochs)
                    }
                    Some(proto::hyperparameters::ModelParams::DqnParams(p)) => {
                        Some(p.epochs)
                    }
                    Some(proto::hyperparameters::ModelParams::PpoParams(p)) => {
                        Some(p.epochs)
                    }
                    Some(proto::hyperparameters::ModelParams::MambaParams(p)) => {
                        Some(p.epochs)
                    }
                    Some(proto::hyperparameters::ModelParams::LiquidParams(p)) => {
                        Some(p.epochs)
                    }
                    Some(proto::hyperparameters::ModelParams::TftParams(p)) => {
                        Some(p.epochs)
                    }
                    None => None,
                })
                .unwrap_or(100);

            let extra_args = if is_hyperopt {
                let trials = req.tags.get("trials")
                    .and_then(|t| t.parse::<u32>().ok())
                    .unwrap_or(20);
                let parallel = req.tags.get("parallel")
                    .and_then(|p| p.parse::<u32>().ok())
                    .unwrap_or(0);
                vec![
                    "--output=/output/hyperopt_results.json".to_string(),
                    format!("--model={}", req.model_type.to_lowercase()),
                    format!("--trials={}", trials),
                    format!("--epochs={}", epochs),
                    format!("--parallel={}", parallel),
                ]
            } else {
                vec![]
            };

            let params = crate::k8s_dispatcher::TrainingJobParams {
                job_id,
                model_type: req.model_type.to_lowercase(),
                symbol,
                epochs,
                binary: binary.to_string(),
                extra_args,
            };

            match dispatcher.dispatch(&params).await {
                Ok(k8s_job_name) => {
                    info!("K8s Job dispatched: {} (job_id={})", k8s_job_name, job_id);
                    return Ok(Response::new(StartTrainingResponse {
                        job_id: job_id.to_string(),
                        status: (ProtoTrainingStatus::Pending as i32),
                        message: format!("K8s training job dispatched: {}", k8s_job_name),
                    }));
                }
                Err(e) => {
                    warn!("K8s dispatch failed, falling back to in-process: {}", e);
                    // Fall through to orchestrator
                }
            }
        }

        // Fallback: in-process orchestrator (for development or if K8s dispatch fails)
        let training_config = self
            .convert_hyperparameters(req.hyperparameters)
            .map_err(|e| Status::invalid_argument(format!("Invalid hyperparameters: {}", e)))?;

        let tags: HashMap<String, String> = req.tags.into_iter().collect();

        let job_id = self
            .orchestrator
            .submit_job(req.model_type, training_config, req.description, tags)
            .await
            .map_err(|e| Status::internal(format!("Failed to submit job: {}", e)))?;

        info!("Training job submitted to orchestrator: {}", job_id);

        Ok(Response::new(StartTrainingResponse {
            job_id: job_id.to_string(),
            status: (ProtoTrainingStatus::Pending as i32),
            message: "Training job submitted successfully".to_string(),
        }))
    }

    /// Subscribe to training status updates
    type SubscribeToTrainingStatusStream =
        Pin<Box<dyn Stream<Item = Result<ProtoStatusUpdate, Status>> + Send>>;

    #[instrument(skip_all)]
    async fn subscribe_to_training_status(
        &self,
        request: Request<SubscribeToTrainingStatusRequest>,
    ) -> Result<Response<Self::SubscribeToTrainingStatusStream>, Status> {
        let req = request.into_inner();

        let job_id = Uuid::parse_str(&req.job_id)
            .map_err(|_| Status::invalid_argument("Invalid job ID format"))?;

        debug!("Subscribing to status updates for job: {}", job_id);

        // Get status update receiver from orchestrator with timeout to prevent hanging
        let receiver = timeout(
            Duration::from_secs(5),
            self.orchestrator.subscribe_to_job_status(job_id),
        )
        .await
        .map_err(|_| Status::deadline_exceeded("Timeout subscribing to job status"))?
        .map_err(|e| Status::not_found(format!("Job not found: {}", e)))?;

        // Convert the broadcast receiver to a stream with proper error handling
        let stream = async_stream::stream! {
            let mut receiver = receiver;
            loop {
                match receiver.recv().await {
                    Ok(update) => yield Ok(Self::convert_status_update(update)),
                    Err(_) => break,
                }
            }
        };

        Ok(Response::new(Box::pin(stream)))
    }

    /// Stop a running training job
    #[instrument(skip_all)]
    async fn stop_training(
        &self,
        request: Request<StopTrainingRequest>,
    ) -> Result<Response<StopTrainingResponse>, Status> {
        let req = request.into_inner();

        let job_id = Uuid::parse_str(&req.job_id)
            .map_err(|_| Status::invalid_argument("Invalid job ID format"))?;

        info!("Stopping training job: {} (reason: {})", job_id, req.reason);

        let success = self
            .orchestrator
            .stop_job(job_id, req.reason)
            .await
            .map_err(|e| Status::internal(format!("Failed to stop job: {}", e)))?;

        let message = if success {
            "Training job stopped successfully".to_string()
        } else {
            "Job was not running or already completed".to_string()
        };

        let response = StopTrainingResponse { success, message };

        Ok(Response::new(response))
    }

    /// List available models for training
    #[instrument(skip_all)]
    async fn list_available_models(
        &self,
        _request: Request<ListAvailableModelsRequest>,
    ) -> Result<Response<ListAvailableModelsResponse>, Status> {
        debug!("Listing available models");

        // Define available models with their configurations
        let models = vec![
            ModelDefinition {
                model_type: "TLOB".to_string(),
                description:
                    "Time-Limit Order Book Transformer for ultra-low latency order book prediction"
                        .to_string(),
                default_hyperparameters: Some(Hyperparameters {
                    model_params: Some(proto::hyperparameters::ModelParams::TlobParams(
                        TlobParams {
                            epochs: 100,
                            learning_rate: 0.001,
                            batch_size: 64,
                            sequence_length: 50,
                            hidden_dim: 256,
                            num_heads: 8,
                            num_layers: 6,
                            dropout_rate: 0.1,
                            use_positional_encoding: true,
                        },
                    )),
                }),
                required_features: vec![
                    "order_book_levels".to_string(),
                    "trade_flow".to_string(),
                    "volume_imbalance".to_string(),
                ],
                estimated_training_time_minutes: 45,
                requires_gpu: true,
            },
            ModelDefinition {
                model_type: "MAMBA_2".to_string(),
                description: "MAMBA-2 State Space Model for long-sequence financial time series"
                    .to_string(),
                default_hyperparameters: Some(Hyperparameters {
                    model_params: Some(proto::hyperparameters::ModelParams::MambaParams(
                        MambaParams {
                            epochs: 150,
                            learning_rate: 0.0005,
                            batch_size: 32,
                            state_dim: 128,
                            hidden_dim: 512,
                            num_layers: 8,
                            dt_min: 0.001,
                            dt_max: 0.1,
                            use_cuda_kernels: true,
                        },
                    )),
                }),
                required_features: vec![
                    "price_series".to_string(),
                    "volume_series".to_string(),
                    "technical_indicators".to_string(),
                ],
                estimated_training_time_minutes: 90,
                requires_gpu: true,
            },
            ModelDefinition {
                model_type: "DQN".to_string(),
                description: "Deep Q-Network for reinforcement learning-based trading strategies"
                    .to_string(),
                default_hyperparameters: Some(Hyperparameters {
                    model_params: Some(proto::hyperparameters::ModelParams::DqnParams(DqnParams {
                        epochs: 200,
                        learning_rate: 0.0001,
                        batch_size: 128,
                        replay_buffer_size: 100000,
                        epsilon_start: 1.0,
                        epsilon_end: 0.01,
                        epsilon_decay_steps: 50000,
                        gamma: 0.99,
                        target_update_frequency: 1000,
                        use_double_dqn: true,
                        use_dueling: true,
                        use_prioritized_replay: true,
                    })),
                }),
                required_features: vec![
                    "market_state".to_string(),
                    "portfolio_state".to_string(),
                    "risk_metrics".to_string(),
                ],
                estimated_training_time_minutes: 120,
                requires_gpu: true,
            },
            ModelDefinition {
                model_type: "PPO".to_string(),
                description: "Proximal Policy Optimization for continuous action space trading"
                    .to_string(),
                default_hyperparameters: Some(Hyperparameters {
                    model_params: Some(proto::hyperparameters::ModelParams::PpoParams(PpoParams {
                        epochs: 100,
                        learning_rate: 0.0003,
                        batch_size: 64,
                        clip_ratio: 0.2,
                        value_loss_coef: 0.5,
                        entropy_coef: 0.01,
                        rollout_steps: 2048,
                        minibatch_size: 64,
                        gae_lambda: 0.95,
                    })),
                }),
                required_features: vec![
                    "market_state".to_string(),
                    "position_state".to_string(),
                    "risk_state".to_string(),
                ],
                estimated_training_time_minutes: 75,
                requires_gpu: true,
            },
            ModelDefinition {
                model_type: "LIQUID".to_string(),
                description: "Liquid Neural Network for adaptive market regime detection"
                    .to_string(),
                default_hyperparameters: Some(Hyperparameters {
                    model_params: Some(proto::hyperparameters::ModelParams::LiquidParams(
                        LiquidParams {
                            epochs: 80,
                            learning_rate: 0.002,
                            batch_size: 48,
                            num_neurons: 128,
                            tau: 0.1,
                            sigma: 0.5,
                            use_adaptive_tau: true,
                        },
                    )),
                }),
                required_features: vec![
                    "volatility_regime".to_string(),
                    "liquidity_metrics".to_string(),
                    "market_microstructure".to_string(),
                ],
                estimated_training_time_minutes: 60,
                requires_gpu: false,
            },
            ModelDefinition {
                model_type: "TFT".to_string(),
                description: "Temporal Fusion Transformer for multi-horizon forecasting"
                    .to_string(),
                default_hyperparameters: Some(Hyperparameters {
                    model_params: Some(proto::hyperparameters::ModelParams::TftParams(TftParams {
                        epochs: 120,
                        learning_rate: 0.001,
                        batch_size: 32,
                        hidden_dim: 240,
                        num_heads: 4,
                        num_layers: 3,
                        lookback_window: 168,
                        forecast_horizon: 24,
                        dropout_rate: 0.3,
                    })),
                }),
                required_features: vec![
                    "historical_prices".to_string(),
                    "external_regressors".to_string(),
                    "calendar_features".to_string(),
                ],
                estimated_training_time_minutes: 100,
                requires_gpu: true,
            },
        ];

        let response = ListAvailableModelsResponse { models };
        Ok(Response::new(response))
    }

    /// List training jobs with filtering
    #[instrument(skip_all)]
    async fn list_training_jobs(
        &self,
        request: Request<ListTrainingJobsRequest>,
    ) -> Result<Response<ListTrainingJobsResponse>, Status> {
        let req = request.into_inner();

        debug!("Listing training jobs with filters");

        // Convert status filter
        let status_filter = if req.status_filter != ProtoTrainingStatus::Unknown as i32 {
            Some(match req.status_filter {
                x if x == ProtoTrainingStatus::Pending as i32 => JobStatus::Pending,
                x if x == ProtoTrainingStatus::Running as i32 => JobStatus::Running,
                x if x == ProtoTrainingStatus::Completed as i32 => JobStatus::Completed,
                x if x == ProtoTrainingStatus::Failed as i32 => JobStatus::Failed,
                x if x == ProtoTrainingStatus::Stopped as i32 => JobStatus::Stopped,
                x if x == ProtoTrainingStatus::Paused as i32 => JobStatus::Paused,
                _ => return Err(Status::invalid_argument("Invalid status filter")),
            })
        } else {
            None
        };

        let model_type_filter = if req.model_type_filter.is_empty() {
            None
        } else {
            Some(req.model_type_filter)
        };

        let limit = if req.page_size == 0 {
            None
        } else {
            Some(usize::try_from(req.page_size).unwrap_or(50))
        };
        let offset = if req.page == 0 {
            None
        } else {
            Some(usize::try_from((req.page - 1) * req.page_size).unwrap_or(0))
        };

        // Get jobs from orchestrator
        let jobs = self
            .orchestrator
            .list_jobs(status_filter, model_type_filter, limit, offset)
            .await
            .map_err(|e| Status::internal(format!("Failed to list jobs: {}", e)))?;

        // Convert to protobuf format
        let job_summaries: Vec<TrainingJobSummary> = jobs
            .into_iter()
            .map(|job| TrainingJobSummary {
                job_id: job.id.to_string(),
                model_type: job.model_type,
                status: Self::convert_job_status(&job.status) as i32,
                created_at: job.created_at.timestamp(),
                started_at: job.started_at.map(|dt| dt.timestamp()).unwrap_or(0),
                completed_at: job.completed_at.map(|dt| dt.timestamp()).unwrap_or(0),
                description: job.description,
                final_loss: job.metrics.get("final_train_loss").copied().unwrap_or(0.0) as f32,
                best_validation_score: job.metrics.get("final_val_loss").copied().unwrap_or(0.0)
                    as f32,
                tags: job.tags,
            })
            .collect();

        let response = ListTrainingJobsResponse {
            jobs: job_summaries,
            total_count: 0, // Would implement proper counting
            page: req.page,
            page_size: req.page_size,
        };

        Ok(Response::new(response))
    }

    /// Get detailed information about a training job
    #[instrument(skip_all)]
    async fn get_training_job_details(
        &self,
        request: Request<GetTrainingJobDetailsRequest>,
    ) -> Result<Response<GetTrainingJobDetailsResponse>, Status> {
        let req = request.into_inner();

        let job_id = Uuid::parse_str(&req.job_id)
            .map_err(|_| Status::invalid_argument("Invalid job ID format"))?;

        debug!("Getting details for training job: {}", job_id);

        // Get job from orchestrator
        let job = self
            .orchestrator
            .get_job(job_id)
            .await
            .map_err(|e| Status::not_found(format!("Job not found: {}", e)))?;

        // Convert to detailed response
        let job_details = TrainingJobDetails {
            job_id: job.id.to_string(),
            model_type: job.model_type,
            status: Self::convert_job_status(&job.status) as i32,
            created_at: job.created_at.timestamp(),
            started_at: job.started_at.map(|dt| dt.timestamp()).unwrap_or(0),
            completed_at: job.completed_at.map(|dt| dt.timestamp()).unwrap_or(0),
            description: job.description,
            hyperparameters: None,         // Would serialize from job.config
            data_source: None,             // Would populate from job configuration
            status_history: Vec::new(),    // Would populate from database
            final_financial_metrics: None, // Would populate from final results
            model_artifact_path: job.model_artifact_path.unwrap_or_default(),
            tags: job.tags,
            error_message: job.error_message.unwrap_or_default(),
        };

        let response = GetTrainingJobDetailsResponse {
            job_details: Some(job_details),
        };

        Ok(Response::new(response))
    }

    /// Health check
    #[instrument(skip_all)]
    async fn health_check(
        &self,
        _request: Request<HealthCheckRequest>,
    ) -> Result<Response<HealthCheckResponse>, Status> {
        debug!("Health check requested");

        let mut details = HashMap::new();
        details.insert("service".to_string(), "ml_training_service".to_string());
        details.insert("version".to_string(), "0.1.0".to_string());
        details.insert("uptime".to_string(), "active".to_string());

        let response = HealthCheckResponse {
            healthy: true,
            message: "ML Training Service is healthy".to_string(),
            details,
        };

        Ok(Response::new(response))
    }

    /// Start a new hyperparameter tuning job
    #[instrument(skip_all)]
    async fn start_tuning_job(
        &self,
        request: Request<proto::StartTuningJobRequest>,
    ) -> Result<Response<proto::StartTuningJobResponse>, Status> {
        self.tuning_handlers.start_tuning_job(request).await
    }

    /// Get status of a tuning job
    #[instrument(skip_all)]
    async fn get_tuning_job_status(
        &self,
        request: Request<proto::GetTuningJobStatusRequest>,
    ) -> Result<Response<proto::GetTuningJobStatusResponse>, Status> {
        self.tuning_handlers.get_tuning_job_status(request).await
    }

    /// Stop a running tuning job
    #[instrument(skip_all)]
    async fn stop_tuning_job(
        &self,
        request: Request<proto::StopTuningJobRequest>,
    ) -> Result<Response<proto::StopTuningJobResponse>, Status> {
        self.tuning_handlers.stop_tuning_job(request).await
    }

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

    #[instrument(skip_all)]
    async fn stream_tuning_progress(
        &self,
        request: Request<proto::StreamProgressRequest>,
    ) -> Result<Response<Self::StreamTuningProgressStream>, Status> {
        self.tuning_handlers.stream_tuning_progress(request).await
    }

    /// INTERNAL: Train a single model with specific hyperparameters (called by Optuna subprocess)
    #[instrument(skip_all)]
    async fn train_model(
        &self,
        request: Request<proto::TrainModelRequest>,
    ) -> Result<Response<proto::TrainModelResponse>, Status> {
        let req = request.into_inner();

        info!(
            "INTERNAL TrainModel called by Optuna - trial_id: {}, model_type: {}",
            req.trial_id, req.model_type
        );

        let start_time = std::time::Instant::now();

        // 1. Parse hyperparameters from request
        let hyperparams = req.hyperparameters;

        // 2. Convert hyperparameters to training config
        let training_config =
            match self.convert_hyperparameters_from_map(&req.model_type, &hyperparams) {
                Ok(config) => config,
                Err(e) => {
                    return Ok(Response::new(proto::TrainModelResponse {
                        success: false,
                        sharpe_ratio: 0.0,
                        training_loss: f32::INFINITY,
                        validation_metrics: HashMap::new(),
                        error_message: format!("Failed to parse hyperparameters: {}", e),
                        training_duration_seconds: 0,
                    }));
                },
            };

        // 3. Load training data (reuse existing pipeline)
        let (training_data, validation_data) =
            match crate::orchestrator::TrainingOrchestrator::load_training_data().await {
                Ok(data) => data,
                Err(e) => {
                    return Ok(Response::new(proto::TrainModelResponse {
                        success: false,
                        sharpe_ratio: 0.0,
                        training_loss: f32::INFINITY,
                        validation_metrics: HashMap::new(),
                        error_message: format!("Failed to load training data: {}", e),
                        training_duration_seconds: 0,
                    }));
                },
            };

        // 4. Create training system and execute training
        let training_system =
            match ml::training_pipeline::ProductionMLTrainingSystem::new(training_config).await {
                Ok(system) => system,
                Err(e) => {
                    return Ok(Response::new(proto::TrainModelResponse {
                        success: false,
                        sharpe_ratio: 0.0,
                        training_loss: f32::INFINITY,
                        validation_metrics: HashMap::new(),
                        error_message: format!("Failed to create training system: {:?}", e),
                        training_duration_seconds: 0,
                    }));
                },
            };

        // 5. Train model
        let result = match training_system
            .train_model(training_data, Some(validation_data.clone()))
            .await
        {
            Ok(r) => r,
            Err(e) => {
                return Ok(Response::new(proto::TrainModelResponse {
                    success: false,
                    sharpe_ratio: 0.0,
                    training_loss: f32::INFINITY,
                    validation_metrics: HashMap::new(),
                    error_message: format!("Training failed: {:?}", e),
                    training_duration_seconds: 0,
                }));
            },
        };

        // 6. Calculate Sharpe ratio on validation set using backtesting
        let sharpe_ratio = self
            .calculate_sharpe_ratio_from_validation(&validation_data, &result)
            .await;

        // 7. Collect validation metrics
        let mut validation_metrics = HashMap::new();
        validation_metrics.insert("validation_loss".to_string(), result.final_val_loss as f32);
        validation_metrics.insert("train_loss".to_string(), result.final_train_loss as f32);

        // Add financial metrics if available
        if let Some(last_metrics) = result.metrics_history.last() {
            validation_metrics.insert(
                "hit_rate".to_string(),
                last_metrics.financial_metrics.hit_rate as f32,
            );
            validation_metrics.insert(
                "avg_prediction_error_bps".to_string(),
                last_metrics.financial_metrics.avg_prediction_error_bps as f32,
            );
        }

        let duration = start_time.elapsed().as_secs() as i64;

        info!(
            "INTERNAL TrainModel completed - trial_id: {}, sharpe_ratio: {:.4}, training_loss: {:.6}, duration: {}s",
            req.trial_id, sharpe_ratio, result.final_train_loss, duration
        );

        // 8. Return response with Sharpe ratio as primary objective
        Ok(Response::new(proto::TrainModelResponse {
            success: true,
            sharpe_ratio,
            training_loss: result.final_train_loss as f32,
            validation_metrics,
            error_message: String::new(),
            training_duration_seconds: duration,
        }))
    }

    /// Start batch tuning job for multiple models
    #[instrument(skip_all)]
    async fn batch_start_tuning_jobs(
        &self,
        request: Request<proto::BatchStartTuningJobsRequest>,
    ) -> Result<Response<proto::BatchStartTuningJobsResponse>, Status> {
        let req = request.into_inner();

        info!(
            "Starting batch tuning job for models: {:?} with {} trials each",
            req.model_types, req.trials_per_model
        );

        // Validate inputs
        if req.model_types.is_empty() {
            return Err(Status::invalid_argument(
                "At least one model type must be specified",
            ));
        }

        if req.trials_per_model == 0 {
            return Err(Status::invalid_argument(
                "Number of trials per model must be > 0",
            ));
        }

        if req.config_path.is_empty() {
            return Err(Status::invalid_argument(
                "Configuration path cannot be empty",
            ));
        }

        // Determine data source path from the proto DataSource
        let data_source_path = req.data_source.and_then(|ds| {
            ds.source.map(|s| match s {
                proto::data_source::Source::FilePath(p) => p,
                proto::data_source::Source::HistoricalDbQuery(q) => q,
                proto::data_source::Source::RealTimeStreamTopic(t) => t,
            })
        });

        // Determine YAML export path (use custom or default)
        let yaml_export_path = if req.yaml_export_path.is_empty() {
            None
        } else {
            Some(req.yaml_export_path)
        };

        // Start batch tuning via manager
        let batch_id = self
            .batch_tuning_manager
            .start_batch_tuning(
                req.model_types,
                req.trials_per_model,
                req.config_path,
                data_source_path,
                req.auto_export_yaml,
                yaml_export_path,
            )
            .await
            .map_err(|e| Status::internal(format!("Failed to start batch tuning: {}", e)))?;

        // Get the job to return the execution order
        let job = self
            .batch_tuning_manager
            .get_batch_status(batch_id)
            .await
            .map_err(|e| {
                Status::internal(format!("Batch started but failed to retrieve status: {}", e))
            })?;

        info!(
            "Batch tuning job started: {} with execution order: {:?}",
            batch_id, job.execution_order
        );

        let response = proto::BatchStartTuningJobsResponse {
            batch_id: batch_id.to_string(),
            execution_order: job.execution_order,
            message: format!("Batch tuning started for {} models", job.models.len()),
            status: Self::convert_batch_status(&job.status) as i32,
        };

        Ok(Response::new(response))
    }

    /// Get batch tuning job status
    #[instrument(skip_all)]
    async fn get_batch_tuning_status(
        &self,
        request: Request<proto::GetBatchTuningStatusRequest>,
    ) -> Result<Response<proto::GetBatchTuningStatusResponse>, Status> {
        let req = request.into_inner();

        let batch_id = Uuid::parse_str(&req.batch_id)
            .map_err(|_| Status::invalid_argument("Invalid batch ID format"))?;

        debug!("Getting batch tuning status for: {}", batch_id);

        let job = self
            .batch_tuning_manager
            .get_batch_status(batch_id)
            .await
            .map_err(|e| Status::not_found(format!("Batch job not found: {}", e)))?;

        // Convert model results to proto format
        let results: Vec<proto::ModelTuningResult> = job
            .results
            .iter()
            .map(|r| proto::ModelTuningResult {
                model_type: r.model_type.clone(),
                job_id: r.job_id.to_string(),
                status: Self::convert_tuning_job_status(&r.status) as i32,
                best_params: r.best_params.clone(),
                best_metrics: r.best_metrics.clone(),
                trials_completed: r.trials_completed,
                started_at: r.started_at.timestamp(),
                completed_at: r.completed_at.map(|dt| dt.timestamp()).unwrap_or(0),
                error_message: r.error_message.clone().unwrap_or_default(),
            })
            .collect();

        // Determine the current model name
        let current_model = job
            .execution_order
            .get(job.current_model_index)
            .cloned()
            .unwrap_or_default();

        let response = proto::GetBatchTuningStatusResponse {
            batch_id: job.batch_id.to_string(),
            status: Self::convert_batch_status(&job.status) as i32,
            current_model_index: job.current_model_index as u32,
            total_models: job.execution_order.len() as u32,
            results,
            current_model,
            started_at: job.started_at.timestamp(),
            updated_at: job.updated_at.timestamp(),
            estimated_completion_time: 0, // Not yet estimated
            yaml_export_path: job.yaml_export_path.clone(),
        };

        Ok(Response::new(response))
    }

    /// Stop a running batch tuning job
    #[instrument(skip_all)]
    async fn stop_batch_tuning_job(
        &self,
        request: Request<proto::StopBatchTuningJobRequest>,
    ) -> Result<Response<proto::StopBatchTuningJobResponse>, Status> {
        let req = request.into_inner();

        let batch_id = Uuid::parse_str(&req.batch_id)
            .map_err(|_| Status::invalid_argument("Invalid batch ID format"))?;

        info!(
            "Stopping batch tuning job: {} (reason: {})",
            batch_id, req.reason
        );

        // Stop the batch job
        self.batch_tuning_manager
            .stop_batch_job(batch_id, req.reason.clone())
            .await
            .map_err(|e| Status::internal(format!("Failed to stop batch job: {}", e)))?;

        // Get final status
        let job = self
            .batch_tuning_manager
            .get_batch_status(batch_id)
            .await
            .map_err(|e| Status::not_found(format!("Batch job not found after stop: {}", e)))?;

        // Convert completed results
        let completed_results: Vec<proto::ModelTuningResult> = job
            .results
            .iter()
            .map(|r| proto::ModelTuningResult {
                model_type: r.model_type.clone(),
                job_id: r.job_id.to_string(),
                status: Self::convert_tuning_job_status(&r.status) as i32,
                best_params: r.best_params.clone(),
                best_metrics: r.best_metrics.clone(),
                trials_completed: r.trials_completed,
                started_at: r.started_at.timestamp(),
                completed_at: r.completed_at.map(|dt| dt.timestamp()).unwrap_or(0),
                error_message: r.error_message.clone().unwrap_or_default(),
            })
            .collect();

        let completed_count = job
            .results
            .iter()
            .filter(|r| r.status == TuningJobStatus::Completed)
            .count();

        let response = proto::StopBatchTuningJobResponse {
            success: true,
            message: format!(
                "Batch job stopped. {} of {} models completed before stop.",
                completed_count,
                job.execution_order.len()
            ),
            final_status: Self::convert_batch_status(&job.status) as i32,
            completed_results,
        };

        Ok(Response::new(response))
    }

    // -- Job completion & model promotion (on-demand training pipeline) --------

    #[instrument(skip_all)]
    async fn report_job_completion(
        &self,
        request: Request<JobCompletionReport>,
    ) -> Result<Response<JobCompletionAck>, Status> {
        let report = request.into_inner();
        info!(
            job_id = %report.job_id,
            success = report.success,
            s3_path = %report.s3_path,
            "received job completion report"
        );

        // Parse the job UUID from the report
        let job_uuid = uuid::Uuid::parse_str(&report.job_id).map_err(|e| {
            Status::invalid_argument(format!("invalid job_id UUID: {}", e))
        })?;

        // ── Failed job path ──────────────────────────────────────────────
        if !report.success {
            // Best-effort DB status update; log warning if it fails
            if let Some(ref spawner) = self.job_spawner {
                if let Err(e) = spawner.update_job_status(job_uuid, "Failed").await {
                    error!(job_id = %report.job_id, error = %e, "failed to update job status to Failed");
                }
            }
            return Ok(Response::new(JobCompletionAck {
                accepted: true,
                promotion_status: "failed".to_string(),
            }));
        }

        // ── Successful job path ──────────────────────────────────────────
        // 1. Update DB status to Completed
        if let Some(ref spawner) = self.job_spawner {
            if let Err(e) = spawner.update_job_status(job_uuid, "Completed").await {
                error!(job_id = %report.job_id, error = %e, "failed to update job status to Completed");
            }
        }

        // 2. Look up model_type and symbol from the child_jobs table
        let (model_type, symbol) = match self.job_spawner {
            Some(ref spawner) => {
                match spawner.get_job_by_id(job_uuid).await {
                    Ok(Some(child_job)) => {
                        // Symbol is stored in config_json.asset
                        let sym = child_job
                            .config_json
                            .get("asset")
                            .and_then(|v| v.as_str())
                            .unwrap_or("UNKNOWN")
                            .to_string();
                        (child_job.model_type, sym)
                    }
                    Ok(None) => {
                        warn!(job_id = %report.job_id, "child job not found in DB, using defaults");
                        ("UNKNOWN".to_string(), "UNKNOWN".to_string())
                    }
                    Err(e) => {
                        warn!(job_id = %report.job_id, error = %e, "failed to look up child job");
                        ("UNKNOWN".to_string(), "UNKNOWN".to_string())
                    }
                }
            }
            None => {
                // No DB — cannot resolve model_type/symbol
                debug!(job_id = %report.job_id, "no JobSpawner, cannot resolve model metadata");
                ("UNKNOWN".to_string(), "UNKNOWN".to_string())
            }
        };

        // 3. Register with PromotionManager
        let metrics: std::collections::HashMap<String, f64> = report.metrics.into_iter().collect();
        let status = self
            .promotion_manager
            .register_completion(&report.job_id, &model_type, &symbol, &report.s3_path, metrics)
            .await;

        let promotion_status = match status {
            crate::promotion_manager::PromotionStatus::PendingPromotion => "pending_promotion",
            crate::promotion_manager::PromotionStatus::NoImprovement => "no_improvement",
            crate::promotion_manager::PromotionStatus::Registered => "registered",
            crate::promotion_manager::PromotionStatus::Error(ref msg) => {
                warn!(job_id = %report.job_id, error = %msg, "promotion registration error");
                "error"
            }
        };

        info!(
            job_id = %report.job_id,
            model_type = %model_type,
            symbol = %symbol,
            promotion_status = %promotion_status,
            "job completion processed"
        );

        Ok(Response::new(JobCompletionAck {
            accepted: true,
            promotion_status: promotion_status.to_string(),
        }))
    }

    #[instrument(skip_all)]
    async fn list_pending_promotions(
        &self,
        _request: Request<ListPendingPromotionsRequest>,
    ) -> Result<Response<ListPendingPromotionsResponse>, Status> {
        let pending = self.promotion_manager.list_pending().await;
        info!(count = pending.len(), "listing pending promotions");

        let promotions = pending
            .into_iter()
            .map(|m| PendingPromotion {
                model_id: m.model_id,
                model_type: m.model_type,
                symbol: m.symbol,
                s3_path: m.s3_path,
                new_metrics: m.new_metrics.into_iter().collect(),
                current_metrics: m.current_metrics.into_iter().collect(),
                trained_at: m.trained_at.timestamp(),
                job_id: m.job_id,
            })
            .collect();

        Ok(Response::new(ListPendingPromotionsResponse { promotions }))
    }

    #[instrument(skip_all)]
    async fn approve_promotion(
        &self,
        request: Request<ApprovePromotionRequest>,
    ) -> Result<Response<ApprovePromotionResponse>, Status> {
        let req = request.into_inner();
        if req.model_id.is_empty() {
            return Err(Status::invalid_argument("model_id is required"));
        }
        info!(model_id = %req.model_id, "approve_promotion request received");

        match self.promotion_manager.approve(&req.model_id).await {
            Ok(()) => {
                info!(model_id = %req.model_id, "model promotion approved and activated");
                Ok(Response::new(ApprovePromotionResponse {
                    success: true,
                    message: format!("Model {} promoted to active", req.model_id),
                }))
            }
            Err(e) => {
                warn!(model_id = %req.model_id, error = %e, "approve_promotion failed");
                Ok(Response::new(ApprovePromotionResponse {
                    success: false,
                    message: format!("Failed to approve promotion: {}", e),
                }))
            }
        }
    }

    #[instrument(skip_all)]
    async fn reject_promotion(
        &self,
        request: Request<RejectPromotionRequest>,
    ) -> Result<Response<RejectPromotionResponse>, Status> {
        let req = request.into_inner();
        if req.model_id.is_empty() {
            return Err(Status::invalid_argument("model_id is required"));
        }
        let reason = if req.reason.is_empty() {
            "no reason provided".to_string()
        } else {
            req.reason
        };
        info!(model_id = %req.model_id, reason = %reason, "reject_promotion request received");

        match self.promotion_manager.reject(&req.model_id, &reason).await {
            Ok(()) => {
                info!(model_id = %req.model_id, reason = %reason, "model promotion rejected");
                Ok(Response::new(RejectPromotionResponse {
                    success: true,
                    message: format!("Promotion of model {} rejected: {}", req.model_id, reason),
                }))
            }
            Err(e) => {
                warn!(model_id = %req.model_id, error = %e, "reject_promotion failed");
                Ok(Response::new(RejectPromotionResponse {
                    success: false,
                    message: format!("Failed to reject promotion: {}", e),
                }))
            }
        }
    }

    // -- Dashboard-facing model approval / rejection ----------------------------

    #[instrument(skip_all)]
    async fn approve_model(
        &self,
        request: Request<ApproveModelRequest>,
    ) -> Result<Response<ApproveModelResponse>, Status> {
        let req = request.into_inner();
        if req.model_id.is_empty() {
            return Err(Status::invalid_argument("model_id is required"));
        }
        info!(model_id = %req.model_id, promoted_to = %req.promoted_to, "approve_model request received");

        match self.promotion_manager.approve(&req.model_id).await {
            Ok(()) => {
                let target = if req.promoted_to.is_empty() {
                    "active".to_string()
                } else {
                    req.promoted_to
                };
                info!(model_id = %req.model_id, promoted_to = %target, "model approved");
                Ok(Response::new(ApproveModelResponse {
                    success: true,
                    message: format!("Model '{}' approved for '{}'", req.model_id, target),
                }))
            }
            Err(e) => {
                warn!(model_id = %req.model_id, error = %e, "approve_model failed");
                Ok(Response::new(ApproveModelResponse {
                    success: false,
                    message: format!("Failed to approve model: {}", e),
                }))
            }
        }
    }

    #[instrument(skip_all)]
    async fn reject_model(
        &self,
        request: Request<RejectModelRequest>,
    ) -> Result<Response<RejectModelResponse>, Status> {
        let req = request.into_inner();
        if req.model_id.is_empty() {
            return Err(Status::invalid_argument("model_id is required"));
        }
        let reason = if req.reason.is_empty() {
            "no reason provided".to_string()
        } else {
            req.reason
        };
        info!(model_id = %req.model_id, reason = %reason, "reject_model request received");

        match self.promotion_manager.reject(&req.model_id, &reason).await {
            Ok(()) => {
                info!(model_id = %req.model_id, reason = %reason, "model rejected");
                Ok(Response::new(RejectModelResponse {
                    success: true,
                    message: format!("Model '{}' rejected: {}", req.model_id, reason),
                }))
            }
            Err(e) => {
                warn!(model_id = %req.model_id, error = %e, "reject_model failed");
                Ok(Response::new(RejectModelResponse {
                    success: false,
                    message: format!("Failed to reject model: {}", e),
                }))
            }
        }
    }
}

impl MLTrainingServiceImpl {
    /// Convert hyperparameters from flat map to ProductionTrainingConfig
    fn convert_hyperparameters_from_map(
        &self,
        model_type: &str,
        params: &HashMap<String, f32>,
    ) -> Result<ProductionTrainingConfig> {
        let mut config = ProductionTrainingConfig::default();

        // Extract common parameters
        let learning_rate = params.get("learning_rate").copied().unwrap_or(0.001) as f64;
        let batch_size = params.get("batch_size").copied().unwrap_or(64.0) as usize;
        let epochs = params.get("epochs").copied().unwrap_or(100.0) as usize;

        // Update training parameters
        config.training_params.learning_rate = learning_rate;
        config.training_params.batch_size = batch_size;
        config.training_params.max_epochs = epochs;

        // Model-specific parameter extraction
        match model_type {
            "TLOB" => {
                let hidden_dim = params.get("hidden_dim").copied().unwrap_or(256.0) as usize;
                let dropout_rate = params.get("dropout_rate").copied().unwrap_or(0.1) as f64;

                config.model_config.input_dim = 20;
                config.model_config.hidden_dims = vec![hidden_dim, hidden_dim / 2];
                config.model_config.dropout_rate = dropout_rate;
                config.model_config.output_dim = 1;
            },
            "MAMBA_2" => {
                let state_dim = params.get("state_dim").copied().unwrap_or(128.0) as usize;
                let hidden_dim = params.get("hidden_dim").copied().unwrap_or(512.0) as usize;

                config.model_config.input_dim = state_dim;
                config.model_config.hidden_dims = vec![hidden_dim];
                config.model_config.output_dim = 1;
            },
            "DQN" => {
                let gamma = params.get("gamma").copied().unwrap_or(0.99);
                // DQN uses default config with gamma stored in tags/metadata
                config.training_params.l2_regularization = 1e-5;
                // Store gamma for later use
                let _ = gamma; // Placeholder for future DQN-specific config
            },
            "PPO" => {
                let clip_ratio = params.get("clip_ratio").copied().unwrap_or(0.2);
                // PPO uses default config with clip_ratio stored in metadata
                let _ = clip_ratio; // Placeholder for future PPO-specific config
            },
            "TFT" => {
                let hidden_dim = params.get("hidden_dim").copied().unwrap_or(240.0) as usize;
                let dropout_rate = params.get("dropout_rate").copied().unwrap_or(0.3) as f64;

                config.model_config.hidden_dims = vec![hidden_dim];
                config.model_config.dropout_rate = dropout_rate;
                config.model_config.output_dim = 1;
            },
            "LIQUID" => {
                let num_neurons = params.get("num_neurons").copied().unwrap_or(128.0) as usize;
                config.model_config.hidden_dims = vec![num_neurons];
                config.model_config.output_dim = 1;
            },
            _ => {
                return Err(anyhow::anyhow!("Unknown model type: {}", model_type));
            },
        }

        Ok(config)
    }

    /// Calculate Sharpe ratio from validation data
    ///
    /// This performs a mini-backtest using the trained model to generate predictions
    /// on the validation set and calculates financial performance metrics.
    async fn calculate_sharpe_ratio_from_validation(
        &self,
        _validation_data: &[(ml::training_pipeline::FinancialFeatures, Vec<f64>)],
        result: &ml::training_pipeline::TrainingResult,
    ) -> f32 {
        // Extract returns from validation predictions
        // In a real implementation, this would:
        // 1. Use the trained model to predict on validation data
        // 2. Simulate trades based on predictions
        // 3. Calculate returns from simulated trades
        // 4. Compute Sharpe ratio from returns

        // For now, use financial metrics from training if available
        if let Some(last_metrics) = result.metrics_history.last() {
            // Use the financial performance metrics from training
            let sharpe = last_metrics.financial_metrics.sharpe_ratio;

            // Validate and clamp to reasonable range
            if sharpe.is_finite() {
                return sharpe.max(-5.0).min(10.0) as f32;
            }
        }

        // Fallback: Estimate Sharpe ratio from validation loss
        // Lower validation loss suggests better predictions -> higher Sharpe ratio
        // This is a rough approximation: Sharpe ≈ -log(val_loss)
        let val_loss = result.final_val_loss;
        if val_loss > 0.0 && val_loss.is_finite() {
            let estimated_sharpe = -(val_loss.ln()) / 2.0; // Normalized to ~0-3 range
            estimated_sharpe.max(-2.0).min(5.0) as f32
        } else {
            0.0
        }
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use crate::orchestrator::TrainingJob;

    #[test]
    fn test_status_conversion() {
        assert_eq!(
            MLTrainingServiceImpl::convert_job_status(&JobStatus::Pending),
            ProtoTrainingStatus::Pending
        );

        assert_eq!(
            MLTrainingServiceImpl::convert_job_status(&JobStatus::Running),
            ProtoTrainingStatus::Running
        );

        assert_eq!(
            MLTrainingServiceImpl::convert_job_status(&JobStatus::Completed),
            ProtoTrainingStatus::Completed
        );

        assert_eq!(
            MLTrainingServiceImpl::convert_job_status(&JobStatus::Failed),
            ProtoTrainingStatus::Failed
        );

        assert_eq!(
            MLTrainingServiceImpl::convert_job_status(&JobStatus::Stopped),
            ProtoTrainingStatus::Stopped
        );

        assert_eq!(
            MLTrainingServiceImpl::convert_job_status(&JobStatus::Paused),
            ProtoTrainingStatus::Paused
        );
    }

    #[test]
    fn test_training_job_creation() {
        let config = ProductionTrainingConfig::default();
        let mut tags = HashMap::new();
        tags.insert("env".to_string(), "test".to_string());

        let job = TrainingJob::new(
            "TLOB".to_string(),
            config,
            "Test job".to_string(),
            tags.clone(),
        );

        assert_eq!(job.model_type, "TLOB");
        assert_eq!(job.status, JobStatus::Pending);
        assert_eq!(job.description, "Test job");
        assert_eq!(job.tags.get("env"), Some(&"test".to_string()));
        assert_eq!(job.progress_percentage, 0.0);
        assert_eq!(job.current_epoch, 0);
        assert!(job.started_at.is_none());
        assert!(job.completed_at.is_none());
    }

    #[test]
    fn test_job_id_uniqueness() {
        let config = ProductionTrainingConfig::default();
        let mut job_ids = vec![];

        for i in 0..100 {
            let job = TrainingJob::new(
                "TLOB".to_string(),
                config.clone(),
                format!("Job {}", i),
                HashMap::new(),
            );
            job_ids.push(job.id);
        }

        // Verify all IDs are unique
        let unique_ids: std::collections::HashSet<_> = job_ids.iter().collect();
        assert_eq!(unique_ids.len(), 100);
    }

    #[test]
    fn test_hyperparameter_protobuf_structure() {
        let tlob_params = Hyperparameters {
            model_params: Some(proto::hyperparameters::ModelParams::TlobParams(
                TlobParams {
                    epochs: 100,
                    learning_rate: 0.001,
                    batch_size: 64,
                    sequence_length: 50,
                    hidden_dim: 256,
                    num_heads: 8,
                    num_layers: 6,
                    dropout_rate: 0.1,
                    use_positional_encoding: true,
                },
            )),
        };

        assert!(matches!(
            tlob_params.model_params,
            Some(proto::hyperparameters::ModelParams::TlobParams(_))
        ));

        if let Some(proto::hyperparameters::ModelParams::TlobParams(params)) =
            tlob_params.model_params
        {
            assert_eq!(params.epochs, 100);
            assert_eq!(params.batch_size, 64);
            assert_eq!(params.num_heads, 8);
        }
    }

    #[test]
    fn test_mamba_hyperparameters() {
        let mamba_params = Hyperparameters {
            model_params: Some(proto::hyperparameters::ModelParams::MambaParams(
                MambaParams {
                    epochs: 150,
                    learning_rate: 0.0005,
                    batch_size: 32,
                    state_dim: 128,
                    hidden_dim: 512,
                    num_layers: 8,
                    dt_min: 0.001,
                    dt_max: 0.1,
                    use_cuda_kernels: true,
                },
            )),
        };

        if let Some(proto::hyperparameters::ModelParams::MambaParams(params)) =
            mamba_params.model_params
        {
            assert_eq!(params.state_dim, 128);
            assert_eq!(params.hidden_dim, 512);
            assert!(params.use_cuda_kernels);
        }
    }

    #[test]
    fn test_dqn_hyperparameters() {
        let dqn_params = Hyperparameters {
            model_params: Some(proto::hyperparameters::ModelParams::DqnParams(DqnParams {
                epochs: 200,
                learning_rate: 0.0001,
                batch_size: 128,
                replay_buffer_size: 100000,
                epsilon_start: 1.0,
                epsilon_end: 0.01,
                epsilon_decay_steps: 50000,
                gamma: 0.99,
                target_update_frequency: 1000,
                use_double_dqn: true,
                use_dueling: true,
                use_prioritized_replay: true,
            })),
        };

        if let Some(proto::hyperparameters::ModelParams::DqnParams(params)) =
            dqn_params.model_params
        {
            assert_eq!(params.replay_buffer_size, 100000);
            assert!(params.use_prioritized_replay);
        }
    }

    #[test]
    fn test_ppo_hyperparameters() {
        let ppo_params = Hyperparameters {
            model_params: Some(proto::hyperparameters::ModelParams::PpoParams(PpoParams {
                epochs: 100,
                learning_rate: 0.0003,
                batch_size: 64,
                clip_ratio: 0.2,
                value_loss_coef: 0.5,
                entropy_coef: 0.01,
                rollout_steps: 2048,
                minibatch_size: 64,
                gae_lambda: 0.95,
            })),
        };

        if let Some(proto::hyperparameters::ModelParams::PpoParams(params)) =
            ppo_params.model_params
        {
            assert_eq!(params.clip_ratio, 0.2);
            assert_eq!(params.rollout_steps, 2048);
        }
    }

    #[test]
    fn test_liquid_hyperparameters() {
        let liquid_params = Hyperparameters {
            model_params: Some(proto::hyperparameters::ModelParams::LiquidParams(
                LiquidParams {
                    epochs: 80,
                    learning_rate: 0.002,
                    batch_size: 48,
                    num_neurons: 128,
                    tau: 0.1,
                    sigma: 0.5,
                    use_adaptive_tau: true,
                },
            )),
        };

        if let Some(proto::hyperparameters::ModelParams::LiquidParams(params)) =
            liquid_params.model_params
        {
            assert_eq!(params.num_neurons, 128);
            assert!(params.use_adaptive_tau);
        }
    }

    #[test]
    fn test_tft_hyperparameters() {
        let tft_params = Hyperparameters {
            model_params: Some(proto::hyperparameters::ModelParams::TftParams(TftParams {
                epochs: 120,
                learning_rate: 0.001,
                batch_size: 32,
                hidden_dim: 240,
                num_heads: 4,
                num_layers: 3,
                lookback_window: 168,
                forecast_horizon: 24,
                dropout_rate: 0.3,
            })),
        };

        if let Some(proto::hyperparameters::ModelParams::TftParams(params)) =
            tft_params.model_params
        {
            assert_eq!(params.lookback_window, 168);
            assert_eq!(params.forecast_horizon, 24);
        }
    }

    #[test]
    fn test_model_types() {
        let models = vec!["TLOB", "MAMBA_2", "DQN", "PPO", "LIQUID", "TFT"];

        // Verify unique count
        let unique_count = models
            .iter()
            .collect::<std::collections::HashSet<_>>()
            .len();
        assert_eq!(unique_count, 6);

        // Verify naming conventions
        for model in models {
            assert!(!model.is_empty());
            assert!(model
                .chars()
                .all(|c| c.is_uppercase() || c.is_numeric() || c == '_'));
        }
    }

    #[test]
    fn test_job_progress_updates() {
        let config = ProductionTrainingConfig::default();
        let mut job = TrainingJob::new(
            "DQN".to_string(),
            config,
            "Progress test".to_string(),
            HashMap::new(),
        );

        // Initial state
        assert_eq!(job.progress_percentage, 0.0);
        assert_eq!(job.current_epoch, 0);

        // Simulate progress
        job.progress_percentage = 50.0;
        job.current_epoch = 50;
        job.total_epochs = 100;

        assert_eq!(job.progress_percentage, 50.0);
        assert_eq!(job.current_epoch, 50);
        assert_eq!(job.total_epochs, 100);
    }

    #[test]
    fn test_job_metrics_tracking() {
        let config = ProductionTrainingConfig::default();
        let mut job = TrainingJob::new(
            "TLOB".to_string(),
            config,
            "Metrics test".to_string(),
            HashMap::new(),
        );

        // Add metrics
        job.metrics.insert("train_loss".to_string(), 0.5);
        job.metrics.insert("val_loss".to_string(), 0.6);
        job.metrics.insert("accuracy".to_string(), 0.85);

        assert_eq!(job.metrics.get("train_loss"), Some(&0.5));
        assert_eq!(job.metrics.get("val_loss"), Some(&0.6));
        assert_eq!(job.metrics.get("accuracy"), Some(&0.85));
        assert_eq!(job.metrics.len(), 3);
    }

    #[test]
    fn test_training_config_defaults() {
        let config = ProductionTrainingConfig::default();

        // Verify default training parameters exist
        assert!(config.training_params.learning_rate > 0.0);
        assert!(config.training_params.batch_size > 0);
        assert!(config.training_params.max_epochs > 0);
        assert!(config.training_params.validation_split >= 0.0);
        assert!(config.training_params.validation_split <= 1.0);
    }

    #[test]
    fn test_service_impl_has_job_spawner() {
        let _: fn(&MLTrainingServiceImpl) -> bool = |s| s.job_spawner.is_some();
    }

    #[test]
    fn test_start_training_model_binary_mapping() {
        use crate::k8s_dispatcher::training_binary_for_model;
        assert_eq!(training_binary_for_model("tft"), "train_baseline_supervised");
        assert_eq!(training_binary_for_model("dqn"), "train_baseline_rl");
        assert_eq!(training_binary_for_model("ppo"), "train_baseline_rl");
        assert_eq!(
            training_binary_for_model("mamba2"),
            "train_baseline_supervised"
        );
    }

    #[test]
    fn test_promotion_status_string_mapping() {
        // Verify the string representations match what gRPC clients expect
        use crate::promotion_manager::PromotionStatus;

        let cases: Vec<(PromotionStatus, &str)> = vec![
            (PromotionStatus::PendingPromotion, "pending_promotion"),
            (PromotionStatus::NoImprovement, "no_improvement"),
            (PromotionStatus::Registered, "registered"),
            (PromotionStatus::Error("test".to_string()), "error"),
        ];

        for (status, expected) in cases {
            let result = match status {
                PromotionStatus::PendingPromotion => "pending_promotion",
                PromotionStatus::NoImprovement => "no_improvement",
                PromotionStatus::Registered => "registered",
                PromotionStatus::Error(_) => "error",
            };
            assert_eq!(result, expected);
        }
    }

    #[tokio::test]
    async fn test_report_job_completion_registers_with_promotion_manager() {
        // Integration test: verify that a successful completion report
        // flows through to the PromotionManager and returns the correct status
        use crate::promotion_manager::PromotionManager;

        let pm = PromotionManager::new();
        let metrics = HashMap::from([
            ("best_val_loss".to_string(), 0.05),
            ("sharpe_ratio".to_string(), 1.5),
        ]);

        // First registration for a type+symbol returns Registered
        let status = pm
            .register_completion("job-1", "DQN", "ES.FUT", "s3://b/m.bin", metrics)
            .await;
        assert_eq!(
            status,
            crate::promotion_manager::PromotionStatus::Registered
        );
    }

    #[test]
    fn test_list_pending_promotions_maps_to_proto() {
        // Verify that PendingModel fields map correctly to proto PendingPromotion
        use crate::promotion_manager::PendingModel;
        use chrono::Utc;

        let now = Utc::now();
        let model = PendingModel {
            model_id: "m-123".to_string(),
            model_type: "DQN".to_string(),
            symbol: "ES.FUT".to_string(),
            s3_path: "s3://bucket/better.bin".to_string(),
            new_metrics: HashMap::from([
                ("best_val_loss".to_string(), 0.03),
                ("sharpe_ratio".to_string(), 2.5),
            ]),
            current_metrics: HashMap::from([
                ("best_val_loss".to_string(), 0.10),
                ("sharpe_ratio".to_string(), 1.0),
            ]),
            trained_at: now,
            job_id: "job-99".to_string(),
        };

        // Map to proto the same way the handler does
        let proto = super::proto::PendingPromotion {
            model_id: model.model_id.clone(),
            model_type: model.model_type.clone(),
            symbol: model.symbol.clone(),
            s3_path: model.s3_path.clone(),
            new_metrics: model.new_metrics.clone().into_iter().collect(),
            current_metrics: model.current_metrics.clone().into_iter().collect(),
            trained_at: model.trained_at.timestamp(),
            job_id: model.job_id.clone(),
        };

        assert_eq!(proto.model_id, "m-123");
        assert_eq!(proto.model_type, "DQN");
        assert_eq!(proto.symbol, "ES.FUT");
        assert_eq!(proto.s3_path, "s3://bucket/better.bin");
        assert_eq!(proto.job_id, "job-99");
        assert_eq!(proto.trained_at, now.timestamp());
        assert_eq!(proto.new_metrics.len(), 2);
        assert_eq!(proto.current_metrics.len(), 2);
    }

    #[tokio::test]
    async fn test_approve_promotion_activates_pending_model() {
        use crate::promotion_manager::{ActiveModel, PromotionManager};
        use chrono::Utc;

        let pm = Arc::new(PromotionManager::new());

        // Seed an active model so register_completion returns PendingPromotion
        {
            let active = pm.active_models_for_test();
            let mut w = active.write().await;
            w.insert(
                ("DQN".to_string(), "ES.FUT".to_string()),
                ActiveModel {
                    model_id: "old-model".to_string(),
                    s3_path: "s3://bucket/old.bin".to_string(),
                    metrics: HashMap::from([
                        ("best_val_loss".to_string(), 0.10),
                        ("sharpe_ratio".to_string(), 1.0),
                    ]),
                    promoted_at: Utc::now(),
                },
            );
        }

        // Register a better model
        let new_metrics = HashMap::from([
            ("best_val_loss".to_string(), 0.03),
            ("sharpe_ratio".to_string(), 2.5),
        ]);
        let status = pm
            .register_completion("job-a", "DQN", "ES.FUT", "s3://bucket/better.bin", new_metrics)
            .await;
        assert_eq!(
            status,
            crate::promotion_manager::PromotionStatus::PendingPromotion
        );

        let pending = pm.list_pending().await;
        assert_eq!(pending.len(), 1);
        let model_id = pending[0].model_id.clone();

        // Approve it
        pm.approve(&model_id).await.unwrap();

        // Pending list should be empty
        assert!(pm.list_pending().await.is_empty());

        // Active model should be updated
        let active = pm.active_models_for_test();
        let r = active.read().await;
        let am = r
            .get(&("DQN".to_string(), "ES.FUT".to_string()))
            .unwrap();
        assert_eq!(am.model_id, model_id);
        assert_eq!(am.s3_path, "s3://bucket/better.bin");
    }

    #[tokio::test]
    async fn test_reject_promotion_removes_pending_model() {
        use crate::promotion_manager::{ActiveModel, PromotionManager};
        use chrono::Utc;

        let pm = Arc::new(PromotionManager::new());

        // Seed an active model
        {
            let active = pm.active_models_for_test();
            let mut w = active.write().await;
            w.insert(
                ("PPO".to_string(), "NQ.FUT".to_string()),
                ActiveModel {
                    model_id: "old-ppo".to_string(),
                    s3_path: "s3://bucket/old-ppo.bin".to_string(),
                    metrics: HashMap::from([
                        ("best_val_loss".to_string(), 0.10),
                        ("sharpe_ratio".to_string(), 1.0),
                    ]),
                    promoted_at: Utc::now(),
                },
            );
        }

        // Register a better model
        let new_metrics = HashMap::from([
            ("best_val_loss".to_string(), 0.04),
            ("sharpe_ratio".to_string(), 2.0),
        ]);
        let status = pm
            .register_completion("job-r", "PPO", "NQ.FUT", "s3://bucket/cand.bin", new_metrics)
            .await;
        assert_eq!(
            status,
            crate::promotion_manager::PromotionStatus::PendingPromotion
        );

        let pending = pm.list_pending().await;
        assert_eq!(pending.len(), 1);
        let model_id = pending[0].model_id.clone();

        // Reject it
        pm.reject(&model_id, "needs more walk-forward windows")
            .await
            .unwrap();

        // Pending list should be empty
        assert!(pm.list_pending().await.is_empty());

        // Active model should NOT have changed
        let active = pm.active_models_for_test();
        let r = active.read().await;
        let am = r
            .get(&("PPO".to_string(), "NQ.FUT".to_string()))
            .unwrap();
        assert_eq!(am.model_id, "old-ppo");
    }

    #[tokio::test]
    async fn test_approve_nonexistent_model_returns_error() {
        use crate::promotion_manager::PromotionManager;

        let pm = PromotionManager::new();
        let result = pm.approve("nonexistent-id").await;
        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .to_string()
                .contains("pending model not found")
        );
    }

    #[tokio::test]
    async fn test_reject_nonexistent_model_returns_error() {
        use crate::promotion_manager::PromotionManager;

        let pm = PromotionManager::new();
        let result = pm.reject("nonexistent-id", "bad model").await;
        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .to_string()
                .contains("pending model not found")
        );
    }
}
