//! Training Orchestrator
//!
//! This module coordinates training jobs using the existing ML training infrastructure
//! from the ml crate. It manages job queues, resource allocation, and progress tracking.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Result;
use chrono::{DateTime, Utc};
use tokio::sync::{broadcast, mpsc, Mutex, RwLock};
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info, warn};
use uuid::Uuid;

// Import from existing ML infrastructure
use ml::training_pipeline::{
    FinancialFeatures, ProductionMLTrainingSystem, ProductionTrainingConfig, TrainingResult,
};

use crate::data_config::TrainingDataSourceConfig;
use crate::database::DatabaseManager;
use crate::storage::ModelStorageManager;
use config::MLConfig;

/// Training job status
#[derive(Debug, Clone, PartialEq)]
pub enum JobStatus {
    Pending,
    Running,
    Completed,
    Failed,
    Stopped,
    Paused,
}

/// Training job metadata
#[derive(Debug, Clone)]
pub struct TrainingJob {
    pub id: Uuid,
    pub model_type: String,
    pub status: JobStatus,
    pub config: ProductionTrainingConfig,
    pub created_at: DateTime<Utc>,
    pub started_at: Option<DateTime<Utc>>,
    pub completed_at: Option<DateTime<Utc>>,
    pub description: String,
    pub tags: HashMap<String, String>,
    pub progress_percentage: f32,
    pub current_epoch: u32,
    pub total_epochs: u32,
    pub metrics: HashMap<String, f64>,
    pub error_message: Option<String>,
    pub model_artifact_path: Option<String>,
}

impl TrainingJob {
    pub fn new(
        model_type: String,
        config: ProductionTrainingConfig,
        description: String,
        tags: HashMap<String, String>,
    ) -> Self {
        Self {
            id: Uuid::new_v4(),
            model_type,
            status: JobStatus::Pending,
            config,
            created_at: Utc::now(),
            started_at: None,
            completed_at: None,
            description,
            tags,
            progress_percentage: 0.0,
            current_epoch: 0,
            total_epochs: 0,
            metrics: HashMap::new(),
            error_message: None,
            model_artifact_path: None,
        }
    }
}

/// Resource allocation for training jobs
#[derive(Debug, Clone)]
pub struct ResourceAllocation {
    pub gpu_id: Option<u32>,
    pub cpu_cores: u32,
    pub memory_gb: f64,
    pub worker_id: u32,
}

/// Training orchestrator that manages job lifecycle and resources
pub struct TrainingOrchestrator {
    config: MLConfig,

    // Job management
    jobs: Arc<RwLock<HashMap<Uuid, TrainingJob>>>,
    job_queue: Arc<Mutex<mpsc::Sender<Uuid>>>,
    job_receiver: Arc<Mutex<mpsc::Receiver<Uuid>>>,

    // Resource management
    available_resources: Arc<Mutex<Vec<ResourceAllocation>>>,
    resource_assignments: Arc<RwLock<HashMap<Uuid, ResourceAllocation>>>,

    // Progress tracking
    status_broadcasters: Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,

    // External dependencies
    database: Arc<crate::database::DatabaseManager>,
    storage: Arc<ModelStorageManager>,

    // Worker pool
    worker_handles: Vec<tokio::task::JoinHandle<()>>,

    // Graceful shutdown
    cancellation_token: CancellationToken,
}

/// Status update message for broadcasting
#[derive(Debug, Clone)]
pub struct TrainingStatusUpdate {
    pub job_id: Uuid,
    pub status: JobStatus,
    pub progress_percentage: f32,
    pub current_epoch: u32,
    pub total_epochs: u32,
    pub metrics: HashMap<String, f64>,
    pub message: String,
    pub timestamp: DateTime<Utc>,
    pub financial_metrics: Option<ml::training_pipeline::FinancialPerformanceMetrics>,
    pub resource_usage: ResourceUsage,
}

#[derive(Debug, Clone)]
pub struct ResourceUsage {
    pub cpu_usage_percent: f32,
    pub memory_usage_gb: f32,
    pub gpu_usage_percent: Option<f32>,
    pub gpu_memory_usage_gb: Option<f32>,
    pub active_workers: u32,
}

impl TrainingOrchestrator {
    /// Create a new training orchestrator
    pub async fn new(
        config: MLConfig,
        database: Arc<crate::database::DatabaseManager>,
        storage: Arc<ModelStorageManager>,
    ) -> Result<Self> {
        let (job_sender, job_receiver) = mpsc::channel(1000); // Default queue capacity

        // Initialize available resources based on system capabilities
        let available_resources = Self::detect_system_resources(&config).await?;

        let orchestrator = Self {
            config: config.clone(),
            jobs: Arc::new(RwLock::new(HashMap::new())),
            job_queue: Arc::new(Mutex::new(job_sender)),
            job_receiver: Arc::new(Mutex::new(job_receiver)),
            available_resources: Arc::new(Mutex::new(available_resources)),
            resource_assignments: Arc::new(RwLock::new(HashMap::new())),
            status_broadcasters: Arc::new(RwLock::new(HashMap::new())),
            database,
            storage,
            worker_handles: Vec::new(),
            cancellation_token: CancellationToken::new(),
        };

        info!("Training orchestrator initialized with default configuration");

        Ok(orchestrator)
    }

    /// Start the orchestrator worker pool
    pub async fn start(&mut self) -> Result<()> {
        info!("Starting training orchestrator with default worker configuration");

        // Start worker tasks with default thread count
        let worker_count = 4; // Default worker threads
        for worker_id in 0..worker_count {
            let worker_handle = self.spawn_worker(worker_id).await?;
            self.worker_handles.push(worker_handle);
        }

        // Start resource monitor
        let resource_monitor_handle = self.spawn_resource_monitor().await?;
        self.worker_handles.push(resource_monitor_handle);

        // Start snapshot broadcaster for fallback status updates
        let snapshot_broadcaster_handle = self.spawn_snapshot_broadcaster().await?;
        self.worker_handles.push(snapshot_broadcaster_handle);

        info!("Training orchestrator started successfully");
        Ok(())
    }

    /// Gracefully shutdown the orchestrator
    pub async fn shutdown(&mut self) -> Result<()> {
        info!("Initiating graceful shutdown of training orchestrator");

        // Signal all workers to stop
        self.cancellation_token.cancel();

        // Wait for all workers to finish with timeout
        let shutdown_timeout = Duration::from_secs(30);
        let mut remaining_handles = Vec::new();

        for handle in self.worker_handles.drain(..) {
            remaining_handles.push(handle);
        }

        // Wait for workers with timeout
        if let Err(_) = tokio::time::timeout(shutdown_timeout, async {
            for handle in remaining_handles {
                let _ = handle.await;
            }
        })
        .await
        {
            warn!("Some workers did not shutdown gracefully within timeout");
        }

        // Clean up any remaining broadcasters for proper resource cleanup
        {
            let mut broadcasters = self.status_broadcasters.write().await;
            broadcasters.clear();
        }

        info!("Training orchestrator shutdown completed");
        Ok(())
    }

    /// Clean up disconnected broadcasters to prevent memory leaks
    pub async fn cleanup_disconnected_broadcasters(&self) {
        let mut broadcasters = self.status_broadcasters.write().await;
        let mut to_remove = Vec::new();

        for (job_id, broadcaster) in broadcasters.iter() {
            if broadcaster.receiver_count() == 0 {
                to_remove.push(*job_id);
            }
        }

        for job_id in to_remove {
            broadcasters.remove(&job_id);
            debug!("Cleaned up broadcaster for job {}", job_id);
        }
    }

    /// Submit a new training job
    pub async fn submit_job(
        &self,
        model_type: String,
        config: ProductionTrainingConfig,
        description: String,
        tags: HashMap<String, String>,
    ) -> Result<Uuid> {
        let job = TrainingJob::new(model_type, config, description, tags);
        let job_id = job.id;

        // Store job in database
        let job_record = crate::database::TrainingJobRecord::from_training_job(&job);
        if let Err(e) = self.database.insert_training_job(&job_record).await {
            error!("Failed to store job {} in database: {}", job_id, e);
            // Continue - in-memory storage still works
        } else {
            info!("Job {} successfully stored in database", job_id);
        }

        // Add to in-memory jobs
        {
            let mut jobs = self.jobs.write().await;
            jobs.insert(job_id, job);
        }

        // Create status broadcaster for this job
        let (tx, _) = broadcast::channel(100); // Default broadcast capacity
        {
            let mut broadcasters = self.status_broadcasters.write().await;
            broadcasters.insert(job_id, tx);
        }

        // Queue the job for execution
        {
            let queue = self.job_queue.lock().await;
            match queue.try_send(job_id) {
                Ok(()) => {},
                Err(mpsc::error::TrySendError::Full(_)) => {
                    return Err(anyhow::anyhow!(
                        "Job queue is full. System is under high load. Please try again later."
                    ));
                },
                Err(mpsc::error::TrySendError::Closed(_)) => {
                    return Err(anyhow::anyhow!(
                        "Job queue is closed. System is shutting down."
                    ));
                },
            }
        }

        info!(
            "Submitted training job {} for model type {}",
            job_id,
            self.get_job_model_type(job_id)
                .await
                .unwrap_or("unknown".to_string())
        );

        Ok(job_id)
    }

    /// Stop a running training job
    pub async fn stop_job(&self, job_id: Uuid, reason: String) -> Result<bool> {
        let mut job_updated = false;

        // Update job status
        {
            let mut jobs = self.jobs.write().await;
            if let Some(job) = jobs.get_mut(&job_id) {
                if matches!(job.status, JobStatus::Running | JobStatus::Pending) {
                    job.status = JobStatus::Stopped;
                    job.completed_at = Some(Utc::now());
                    job.error_message = Some(format!("Stopped: {}", reason));
                    job_updated = true;
                }
            }
        }

        if job_updated {
            // Update database
            let jobs_read = self.jobs.read().await;
            if let Some(job) = jobs_read.get(&job_id) {
                let job_record = crate::database::TrainingJobRecord::from_training_job(job);
                drop(jobs_read); // Release lock before async call

                if let Err(e) = self.database.update_training_job(&job_record).await {
                    error!("Failed to update job {} in database: {}", job_id, e);
                } else {
                    info!("Job {} successfully updated in database", job_id);
                }
            }

            // Broadcast status update
            self.broadcast_status_update(job_id, "Job stopped by user request".to_string())
                .await;

            info!("Stopped training job {}: {}", job_id, reason);
        }

        Ok(job_updated)
    }

    /// Get job status
    pub async fn get_job(&self, job_id: Uuid) -> Result<TrainingJob> {
        let jobs = self.jobs.read().await;
        jobs.get(&job_id)
            .cloned()
            .ok_or_else(|| anyhow::anyhow!("Job {} not found", job_id))
    }

    /// List all jobs with optional filtering
    pub async fn list_jobs(
        &self,
        status_filter: Option<JobStatus>,
        model_type_filter: Option<String>,
        limit: Option<usize>,
        offset: Option<usize>,
    ) -> Result<Vec<TrainingJob>> {
        let jobs = self.jobs.read().await;
        let mut filtered_jobs: Vec<_> = jobs
            .values()
            .filter(|job| {
                if let Some(ref status) = status_filter {
                    if job.status != *status {
                        return false;
                    }
                }
                if let Some(ref model_type) = model_type_filter {
                    if job.model_type != *model_type {
                        return false;
                    }
                }
                true
            })
            .cloned()
            .collect();

        // Sort by creation time (newest first)
        filtered_jobs.sort_by(|a, b| b.created_at.cmp(&a.created_at));

        // Apply pagination
        if let Some(offset) = offset {
            if offset < filtered_jobs.len() {
                filtered_jobs = filtered_jobs.into_iter().skip(offset).collect();
            } else {
                filtered_jobs.clear();
            }
        }

        if let Some(limit) = limit {
            filtered_jobs.truncate(limit);
        }

        Ok(filtered_jobs)
    }

    /// Subscribe to status updates for a job
    pub async fn subscribe_to_job_status(
        &self,
        job_id: Uuid,
    ) -> Result<broadcast::Receiver<TrainingStatusUpdate>> {
        let broadcasters = self.status_broadcasters.read().await;
        if let Some(broadcaster) = broadcasters.get(&job_id) {
            Ok(broadcaster.subscribe())
        } else {
            Err(anyhow::anyhow!("Job {} not found", job_id))
        }
    }

    /// Detect available system resources
    async fn detect_system_resources(_config: &MLConfig) -> Result<Vec<ResourceAllocation>> {
        let mut resources = Vec::new();

        // Default configuration for resource allocation
        let worker_threads = 4u32;
        let default_device = "cpu"; // Default to CPU for safety
        let max_memory_gb = 8.0; // Default memory allocation

        // For now, create one resource allocation per worker thread
        // In a real implementation, this would detect actual GPU devices
        for worker_id in 0..worker_threads {
            let allocation = ResourceAllocation {
                gpu_id: if default_device == "cuda" {
                    Some(0)
                } else {
                    None
                },
                cpu_cores: num_cpus::get() as u32 / worker_threads,
                memory_gb: max_memory_gb / worker_threads as f64,
                worker_id,
            };
            resources.push(allocation);
        }

        info!("Detected {} resource allocations", resources.len());
        Ok(resources)
    }

    /// Spawn a worker task
    async fn spawn_worker(&self, worker_id: usize) -> Result<tokio::task::JoinHandle<()>> {
        let jobs = Arc::clone(&self.jobs);
        let job_receiver = Arc::clone(&self.job_receiver);
        let available_resources = Arc::clone(&self.available_resources);
        let resource_assignments = Arc::clone(&self.resource_assignments);
        let status_broadcasters = Arc::clone(&self.status_broadcasters);
        let database = Arc::clone(&self.database);
        let storage = Arc::clone(&self.storage);
        let config = self.config.clone();
        let cancellation_token = self.cancellation_token.child_token();

        let handle = tokio::spawn(async move {
            info!("Training worker {} started", worker_id);

            loop {
                // Wait for a job to process
                let job_id = {
                    let mut receiver = job_receiver.lock().await;
                    tokio::select! {
                        job_result = receiver.recv() => {
                            match job_result {
                                Some(id) => id,
                                None => {
                                    warn!("Worker {} job receiver closed", worker_id);
                                    break;
                                }
                            }
                        }
                        _ = cancellation_token.cancelled() => {
                            info!("Worker {} received cancellation signal", worker_id);
                            break;
                        }
                    }
                };

                // Process the job
                if let Err(e) = Self::process_job(
                    job_id,
                    worker_id,
                    &jobs,
                    &available_resources,
                    &resource_assignments,
                    &status_broadcasters,
                    &database,
                    &storage,
                    &config,
                )
                .await
                {
                    error!(
                        "Worker {} failed to process job {}: {}",
                        worker_id, job_id, e
                    );
                }
            }

            info!("Training worker {} stopped", worker_id);
        });

        Ok(handle)
    }

    /// Process a single training job
    async fn process_job(
        job_id: Uuid,
        worker_id: usize,
        jobs: &Arc<RwLock<HashMap<Uuid, TrainingJob>>>,
        available_resources: &Arc<Mutex<Vec<ResourceAllocation>>>,
        resource_assignments: &Arc<RwLock<HashMap<Uuid, ResourceAllocation>>>,
        status_broadcasters: &Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,
        database: &Arc<DatabaseManager>,
        storage: &Arc<ModelStorageManager>,
        _config: &MLConfig,
    ) -> Result<()> {
        info!("Worker {} processing job {}", worker_id, job_id);

        // Get job details
        let (job_config, model_type) = {
            let jobs_read = jobs.read().await;
            if let Some(job) = jobs_read.get(&job_id) {
                (job.config.clone(), job.model_type.clone())
            } else {
                return Err(anyhow::anyhow!("Job {} not found", job_id));
            }
        };

        // Allocate resources
        let resource_allocation = {
            let mut resources = available_resources.lock().await;
            if let Some(resource) = resources.pop() {
                resource
            } else {
                error!("No available resources for job {}", job_id);
                return Self::mark_job_failed(
                    job_id,
                    "No available resources".to_string(),
                    jobs,
                    database,
                )
                .await;
            }
        };

        // Assign resource to job
        {
            let mut assignments = resource_assignments.write().await;
            assignments.insert(job_id, resource_allocation.clone());
        }

        // Update job status to running
        {
            let mut jobs_write = jobs.write().await;
            if let Some(job) = jobs_write.get_mut(&job_id) {
                job.status = JobStatus::Running;
                job.started_at = Some(Utc::now());
            }
        }

        // Broadcast status update
        Self::send_status_update(
            job_id,
            JobStatus::Running,
            0.0,
            0,
            job_config.training_params.max_epochs as u32,
            HashMap::new(),
            "Training started".to_string(),
            None,
            &resource_allocation,
            status_broadcasters,
        )
        .await;

        // Execute training using existing ML infrastructure
        let training_result = Self::execute_training(
            job_id,
            job_config,
            model_type,
            &resource_allocation,
            status_broadcasters,
        )
        .await;

        // Handle training result
        match training_result {
            Ok(result) => {
                Self::handle_training_success(job_id, result, jobs, database, storage).await?;
            },
            Err(e) => {
                Self::mark_job_failed(job_id, e.to_string(), jobs, database).await?;
            },
        }

        // Release resources
        {
            let mut resources = available_resources.lock().await;
            resources.push(resource_allocation);
        }
        {
            let mut assignments = resource_assignments.write().await;
            assignments.remove(&job_id);
        }

        info!("Worker {} completed job {}", worker_id, job_id);
        Ok(())
    }

    /// Execute training using the existing ML training infrastructure
    async fn execute_training(
        job_id: Uuid,
        config: ProductionTrainingConfig,
        model_type: String,
        _resource_allocation: &ResourceAllocation,
        _status_broadcasters: &Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,
    ) -> Result<TrainingResult> {
        info!(
            "Starting training execution for job {} with model type {}",
            job_id, model_type
        );

        // Create the production training system
        let training_system = ProductionMLTrainingSystem::new(config.clone())
            .await
            .map_err(|e| anyhow::anyhow!("Failed to create training system: {:?}", e))?;

        // Load training data from configured source
        // Phase 1: Configuration established, Phase 2-3 will implement actual loading
        let (training_data, validation_data) = Self::load_training_data().await?;

        // Execute training with progress callbacks
        let result = training_system
            .train_model(training_data, Some(validation_data))
            .await
            .map_err(|e| anyhow::anyhow!("Training failed: {:?}", e))?;

        info!(
            "Training completed for job {} with final loss: {:.6}",
            job_id, result.final_train_loss
        );
        Ok(result)
    }

    /// Load training data from configured source with 225-feature extraction
    ///
    /// Loads real market data from DBN files and applies full 225-feature extraction
    /// (Wave C 201 features + Wave D 24 features) using FeatureExtractor.
    pub async fn load_training_data() -> Result<(
        Vec<(FinancialFeatures, Vec<f64>)>,
        Vec<(FinancialFeatures, Vec<f64>)>,
    )> {
        

        // Primary: Try to load real DBN market data
        let dbn_file_path = std::env::var("DBN_DATA_FILE").unwrap_or_else(|_| {
            "test_data/real/databento/ES.FUT_ohlcv-1m_2024-01-02.dbn".to_string()
        });

        if std::path::Path::new(&dbn_file_path).exists() {
            info!(
                "📊 Loading REAL market data from DBN file with 225-feature extraction: {}",
                dbn_file_path
            );

            match Self::load_training_data_with_225_features(&dbn_file_path, 0.8).await {
                Ok((training_data, validation_data)) => {
                    info!(
                        "✅ Loaded {} training samples, {} validation samples with 225 features (Wave C + Wave D)",
                        training_data.len(),
                        validation_data.len()
                    );
                    return Ok((training_data, validation_data));
                },
                Err(e) => {
                    warn!("Failed to load DBN data with 225-feature extraction: {}, falling back to database", e);
                },
            }
        } else {
            debug!("DBN file not found at {}, trying database", dbn_file_path);
        }

        // Fallback: Try database loading or mock data
        #[cfg(feature = "mock-data")]
        {
            error!("Using MOCK training data - NOT FOR PRODUCTION USE!");
            error!("Rebuild without --features mock-data for production");
            let training_data = Self::generate_mock_training_data()?;
            let validation_data = Self::generate_mock_validation_data()?;
            return Ok((training_data, validation_data));
        }

        #[cfg(not(feature = "mock-data"))]
        {
            use crate::data_config::DataSourceType;
            use crate::data_loader::HistoricalDataLoader;

            // Load data source configuration
            let data_config = TrainingDataSourceConfig::from_env()
                .map_err(|e| anyhow::anyhow!("Failed to load data source configuration: {}", e))?;

            data_config
                .validate()
                .map_err(|e| anyhow::anyhow!("Invalid data source configuration: {}", e))?;

            info!(
                "📊 Training data configuration loaded: {:?}",
                data_config.summary()
            );

            // Load data based on source type
            match data_config.source_type {
                DataSourceType::Historical | DataSourceType::Hybrid => {
                    info!("Loading historical training data from PostgreSQL");

                    let mut loader = HistoricalDataLoader::new(data_config)
                        .await
                        .map_err(|e| anyhow::anyhow!("Failed to create data loader: {}", e))?;

                    let (training_data, validation_data) = loader
                        .load_training_data()
                        .await
                        .map_err(|e| anyhow::anyhow!("Failed to load training data: {}", e))?;

                    info!(
                        "✅ Loaded {} training samples, {} validation samples from database",
                        training_data.len(),
                        validation_data.len()
                    );

                    Ok((training_data, validation_data))
                },
                DataSourceType::RealTime => Err(anyhow::anyhow!(
                    "❌ RealTime data source not yet implemented (Phase 3)\n\
                         \n\
                         📋 Supported data sources:\n\
                         - DBN: Real market data files (Phase 2 ✅)\n\
                         - Historical: PostgreSQL database (Phase 2 ✅)\n\
                         - RealTime: Live streaming data (Phase 3 pending)\n\
                         - Hybrid: Historical + RealTime (Phase 3 pending)\n\
                         - Parquet: S3 parquet files (Phase 4 pending)\n\
                         \n\
                         🔧 Set DBN_DATA_FILE=/path/to/file.dbn for real market data\n\
                         Set DATA_SOURCE_TYPE=historical to use database loading\n\
                         Set DATABASE_URL to your PostgreSQL instance"
                )),
                DataSourceType::Parquet => Err(anyhow::anyhow!(
                    "❌ Parquet data source not yet implemented (Phase 4)\n\
                         \n\
                         📋 Supported data sources:\n\
                         - DBN: Real market data files (Phase 2 ✅)\n\
                         - Historical: PostgreSQL database (Phase 2 ✅)\n\
                         - Parquet: S3 parquet files (Phase 4 pending)\n\
                         \n\
                         🔧 Set DBN_DATA_FILE=/path/to/file.dbn for real market data\n\
                         Set DATA_SOURCE_TYPE=historical to use database loading\n\
                         Set DATABASE_URL to your PostgreSQL instance"
                )),
            }
        }
    }

    /// Load training data from DBN file with 225-feature extraction (Wave C + Wave D)
    ///
    /// Uses FeatureExtractor to produce all 225 features per bar:
    /// - Wave C (indices 0-200): 201 features
    /// - Wave D (indices 201-224): 24 regime detection features
    async fn load_training_data_with_225_features(
        dbn_file_path: &str,
        train_split: f64,
    ) -> Result<(
        Vec<(FinancialFeatures, Vec<f64>)>,
        Vec<(FinancialFeatures, Vec<f64>)>,
    )> {
        use dbn::decode::{DbnDecoder, DecodeRecordRef};
        use dbn::OhlcvMsg;
        use ml::features::extraction::{FeatureExtractor, OHLCVBar};
        use anyhow::Context;
        use chrono::TimeZone;

        // Load OHLCV bars from DBN file (reuse existing loader logic)
        debug!("Loading DBN file for 225-feature extraction: {}", dbn_file_path);

        let mut decoder = DbnDecoder::from_file(dbn_file_path)
            .context(format!("Failed to create DBN decoder for file: {}", dbn_file_path))?;

        let mut ohlcv_bars = Vec::new();

        // Decode all OHLCV records from DBN file
        while let Some(record_ref) = decoder.decode_record_ref()
            .context("Failed to decode DBN record")?
        {
            if let Some(ohlcv_msg) = record_ref.get::<OhlcvMsg>() {
                // Convert timestamp from nanoseconds
                let ts_nanos = ohlcv_msg.hd.ts_event as i64;
                let secs = ts_nanos / 1_000_000_000;
                let nanos = (ts_nanos % 1_000_000_000) as u32;
                let timestamp = chrono::Utc
                    .timestamp_opt(secs, nanos)
                    .single()
                    .ok_or_else(|| anyhow::anyhow!("Invalid timestamp: {}", ts_nanos))?;

                // Convert DBN fixed-point prices to f64 (9 decimal places)
                let dbn_to_f64 = |price: i64| price as f64 / 1_000_000_000.0;

                ohlcv_bars.push(OHLCVBar {
                    timestamp,
                    open: dbn_to_f64(ohlcv_msg.open),
                    high: dbn_to_f64(ohlcv_msg.high),
                    low: dbn_to_f64(ohlcv_msg.low),
                    close: dbn_to_f64(ohlcv_msg.close),
                    volume: ohlcv_msg.volume as f64,
                });
            }
        }

        info!("Loaded {} OHLCV bars from DBN file", ohlcv_bars.len());

        // Extract 225 features using FeatureExtractor (Wave C + Wave D)
        const WARMUP_PERIOD: usize = 50;

        if ohlcv_bars.len() < WARMUP_PERIOD {
            return Err(anyhow::anyhow!(
                "Insufficient data: {} bars provided, {} required for warmup",
                ohlcv_bars.len(),
                WARMUP_PERIOD
            ));
        }

        let mut extractor = FeatureExtractor::new();
        let mut feature_vectors_225 = Vec::with_capacity(ohlcv_bars.len() - WARMUP_PERIOD);

        // Feed bars sequentially to build rolling windows
        for (i, bar) in ohlcv_bars.iter().enumerate() {
            extractor.update(bar).context(format!("Feature extraction failed at bar {}", i))?;

            // Start extracting features after warmup
            if i >= WARMUP_PERIOD {
                let features_225 = extractor.extract_current_features()
                    .context(format!("Failed to extract features at bar {}", i))?;
                feature_vectors_225.push(features_225);
            }
        }

        info!(
            "Extracted {} feature vectors (225 dimensions each, Wave C + Wave D)",
            feature_vectors_225.len()
        );

        // Convert [f64; 225] to (FinancialFeatures, Vec<f64>) format for compatibility.
        //
        // This conversion exists because UnifiedTrainable::train() accepts
        // Vec<(FinancialFeatures, Vec<f64>)>. The FinancialFeatures fields below are
        // stubs — only the Vec<f64> component is consumed by the underlying models.
        //
        // To remove this layer: change UnifiedTrainable::train() to accept &[Vec<f64>]
        // directly and update all 4 model adapters (DQN, PPO, TFT, Mamba2).
        let training_samples: Vec<(FinancialFeatures, Vec<f64>)> = feature_vectors_225
            .iter()
            .zip(ohlcv_bars.iter().skip(WARMUP_PERIOD))
            .map(|(features_225, bar)| {
                // Create stub FinancialFeatures for compatibility
                let close_price = common::Price::new(bar.close).unwrap_or_default();
                let financial_features = FinancialFeatures {
                    timestamp: bar.timestamp,
                    prices: vec![close_price],
                    volumes: vec![bar.volume as i64],
                    technical_indicators: std::collections::HashMap::new(),
                    microstructure: ml::training_pipeline::MicrostructureFeatures {
                        spread_bps: 0,
                        imbalance: 0.0,
                        trade_intensity: 0.0,
                        vwap: close_price,
                    },
                    risk_metrics: ml::training_pipeline::RiskFeatures {
                        var_5pct: 0.0,
                        expected_shortfall: 0.0,
                        max_drawdown: 0.0,
                        sharpe_ratio: 0.0,
                    },
                };

                // Convert [f64; 225] to Vec<f64>
                (financial_features, features_225.to_vec())
            })
            .collect();

        // Split into train/validation (preserve time-series ordering)
        let split_idx = (training_samples.len() as f64 * train_split) as usize;
        let train_data = training_samples[..split_idx].to_vec();
        let val_data = training_samples[split_idx..].to_vec();

        info!(
            "✅ Created {} training samples, {} validation samples (225 features each)",
            train_data.len(),
            val_data.len()
        );

        Ok((train_data, val_data))
    }

    /// Generate mock training data for demonstration
    #[cfg(feature = "mock-data")]
    fn generate_mock_training_data() -> Result<Vec<(FinancialFeatures, Vec<f64>)>> {
        // This is a simplified mock implementation
        // In production, this would load real market data
        let mut data = Vec::new();

        for i in 0..1000 {
            let price = 100.0 + (i as f64 * 0.01);
            let features = FinancialFeatures {
                prices: vec![
                    common::Price::from_f64(price).unwrap_or_default()
                ],
                volumes: vec![1000 + i as i64],
                technical_indicators: [("rsi".to_string(), 0.5 + 0.3 * (i as f64 / 1000.0).sin())]
                    .iter()
                    .cloned()
                    .collect(),
                microstructure: ml::training_pipeline::MicrostructureFeatures {
                    spread_bps: 10,
                    imbalance: 0.1 * (i as f64 / 100.0).sin(),
                    trade_intensity: 2.5,
                    vwap: common::Price::from_f64(price * 0.9995).unwrap_or_default(),
                },
                risk_metrics: ml::training_pipeline::RiskFeatures {
                    var_5pct: -0.02,
                    expected_shortfall: -0.03,
                    max_drawdown: -0.05,
                    sharpe_ratio: 1.2,
                },
                timestamp: chrono::Utc::now(),
            };

            let targets = vec![price * 1.001]; // Predict small price increase
            data.push((features, targets));
        }

        Ok(data)
    }

    /// Generate mock validation data
    #[cfg(feature = "mock-data")]
    fn generate_mock_validation_data() -> Result<Vec<(FinancialFeatures, Vec<f64>)>> {
        // Similar to training data but smaller
        Self::generate_mock_training_data().map(|mut data| {
            data.truncate(200);
            data
        })
    }

    /// Handle successful training completion
    async fn handle_training_success(
        job_id: Uuid,
        result: TrainingResult,
        jobs: &Arc<RwLock<HashMap<Uuid, TrainingJob>>>,
        _database: &Arc<DatabaseManager>,
        storage: &Arc<ModelStorageManager>,
    ) -> Result<()> {
        // Serialize training result to model data (simplified)
        let model_data = format!(
            "{{\"final_train_loss\":{},\"final_val_loss\":{}}}",
            result.final_train_loss, result.final_val_loss
        )
        .into_bytes();

        // Attempt to store model artifact with S3 upload
        let model_path = match storage.store_model(job_id, &model_data).await {
            Ok(path) => {
                info!(
                    "Successfully uploaded model for job {} to storage: {}",
                    job_id, path
                );
                path
            },
            Err(e) => {
                error!("Failed to upload model for job {}: {}", job_id, e);
                // Fallback to local path for tracking
                let fallback_path = format!("models/{}.bin", job_id);
                warn!("Using fallback path for job {}: {}", job_id, fallback_path);
                fallback_path
            },
        };

        // Extract accuracy from training result metrics history
        let accuracy = result
            .metrics_history
            .last()
            .map(|m| m.prediction_accuracy)
            .unwrap_or(0.0);

        // Create model metadata for tracking and versioning
        let model_metadata = {
            let job_guard = jobs.read().await;
            if let Some(job) = job_guard.get(&job_id) {
                config::ModelRegistryEntry {
                    id: job_id,
                    name: job.model_type.clone(),
                    version: format!("v{}", chrono::Utc::now().format("%Y%m%d_%H%M%S")),
                    created_at: chrono::Utc::now(),
                    updated_at: chrono::Utc::now(),
                    training_metrics: config::TrainingMetrics {
                        accuracy,
                        loss: result.final_train_loss,
                        validation_accuracy: accuracy, // Same as training accuracy for now
                        validation_loss: result.final_val_loss,
                        epochs: result.epochs_trained as u32,
                        training_time_seconds: result.training_duration.num_seconds() as f64,
                    },
                    architecture: config::ModelArchitecture {
                        model_type: job.model_type.clone(),
                        input_dim: job.config.model_config.input_dim,
                        output_dim: job.config.model_config.output_dim,
                        hidden_layers: job.config.model_config.hidden_dims.clone(),
                        activation: job.config.model_config.activation.clone(),
                        optimizer: "adamw".to_string(), // Standard optimizer for production ML
                        learning_rate: job.config.training_params.learning_rate,
                    },
                }
            } else {
                return Err(anyhow::anyhow!(
                    "Job {} not found for metadata creation",
                    job_id
                ));
            }
        };

        info!(
            "Created model metadata for job {} with version {}",
            job_id, model_metadata.version
        );

        // Update job status
        {
            let mut jobs_write = jobs.write().await;
            if let Some(job) = jobs_write.get_mut(&job_id) {
                job.status = JobStatus::Completed;
                job.completed_at = Some(Utc::now());
                job.progress_percentage = 100.0;
                job.model_artifact_path = Some(model_path.clone());
                job.metrics
                    .insert("final_train_loss".to_string(), result.final_train_loss);
                job.metrics
                    .insert("final_val_loss".to_string(), result.final_val_loss);
                job.metrics.insert(
                    "model_version".to_string(),
                    model_metadata.version.parse().unwrap_or(0.0),
                );
                job.metrics
                    .insert("model_size_bytes".to_string(), model_data.len() as f64);
            }
        }

        info!("Training job {} completed successfully", job_id);
        Ok(())
    }

    /// Mark job as failed
    async fn mark_job_failed(
        job_id: Uuid,
        error_message: String,
        jobs: &Arc<RwLock<HashMap<Uuid, TrainingJob>>>,
        _database: &Arc<DatabaseManager>,
    ) -> Result<()> {
        {
            let mut jobs_write = jobs.write().await;
            if let Some(job) = jobs_write.get_mut(&job_id) {
                job.status = JobStatus::Failed;
                job.completed_at = Some(Utc::now());
                job.error_message = Some(error_message.clone());
            }
        }

        error!("Training job {} failed: {}", job_id, error_message);
        Ok(())
    }

    /// Send status update
    async fn send_status_update(
        job_id: Uuid,
        status: JobStatus,
        progress: f32,
        current_epoch: u32,
        total_epochs: u32,
        metrics: HashMap<String, f64>,
        message: String,
        financial_metrics: Option<ml::training_pipeline::FinancialPerformanceMetrics>,
        resource_allocation: &ResourceAllocation,
        status_broadcasters: &Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,
    ) {
        let update = TrainingStatusUpdate {
            job_id,
            status,
            progress_percentage: progress,
            current_epoch,
            total_epochs,
            metrics,
            message,
            timestamp: Utc::now(),
            financial_metrics,
            resource_usage: ResourceUsage {
                cpu_usage_percent: 75.0, // Mock values
                memory_usage_gb: resource_allocation.memory_gb as f32,
                gpu_usage_percent: resource_allocation.gpu_id.map(|_| 80.0),
                gpu_memory_usage_gb: resource_allocation.gpu_id.map(|_| 4.0),
                active_workers: 1,
            },
        };

        let broadcasters = status_broadcasters.read().await;
        if let Some(broadcaster) = broadcasters.get(&job_id) {
            match broadcaster.send(update.clone()) {
                Ok(_) => {
                    debug!("Status update sent successfully for job {}", job_id);
                },
                Err(broadcast::error::SendError(_)) => {
                    warn!(
                        "Failed to send status update for job {} - all receivers dropped. \
                         This indicates clients have disconnected.",
                        job_id
                    );
                    // Note: The broadcast channel automatically handles the case where
                    // receivers are lagging behind. Slow receivers will miss old messages
                    // but will still receive new ones, providing natural back-pressure.
                },
            }
        } else {
            debug!("No broadcaster found for job {}", job_id);
        }
    }

    /// Broadcast status update for a job
    async fn broadcast_status_update(&self, job_id: Uuid, message: String) {
        let jobs = self.jobs.read().await;
        if let Some(job) = jobs.get(&job_id) {
            Self::send_status_update(
                job_id,
                job.status.clone(),
                job.progress_percentage,
                job.current_epoch,
                job.total_epochs,
                job.metrics.clone(),
                message,
                None,
                &ResourceAllocation {
                    gpu_id: None,
                    cpu_cores: 1,
                    memory_gb: 1.0,
                    worker_id: 0,
                },
                &self.status_broadcasters,
            )
            .await;
        }
    }

    /// Get job model type
    async fn get_job_model_type(&self, job_id: Uuid) -> Option<String> {
        let jobs = self.jobs.read().await;
        jobs.get(&job_id).map(|job| job.model_type.clone())
    }

    /// Spawn resource monitor
    async fn spawn_resource_monitor(&self) -> Result<tokio::task::JoinHandle<()>> {
        let interval = Duration::from_secs(30);

        let handle = tokio::spawn(async move {
            let mut interval_timer = tokio::time::interval(interval);

            loop {
                interval_timer.tick().await;

                // Monitor system resources
                // This is a simplified implementation
                debug!("Resource monitor tick - system resources OK");
            }
        });

        Ok(handle)
    }

    /// Periodically send snapshot status updates for all active jobs
    ///
    /// This provides a fallback mechanism when streaming updates are overwhelmed
    async fn spawn_snapshot_broadcaster(&self) -> Result<tokio::task::JoinHandle<()>> {
        let jobs = Arc::clone(&self.jobs);
        let status_broadcasters = Arc::clone(&self.status_broadcasters);
        let interval = Duration::from_secs(30); // Default snapshot interval
        let cancellation_token = self.cancellation_token.child_token();

        let handle = tokio::spawn(async move {
            let mut interval_timer = tokio::time::interval(interval);

            loop {
                tokio::select! {
                    _ = interval_timer.tick() => {
                        // Send snapshot updates for all active jobs
                        let jobs_read = jobs.read().await;
                        for job in jobs_read.values() {
                            if matches!(job.status, JobStatus::Running | JobStatus::Pending) {
                                Self::send_snapshot_update(job, &status_broadcasters).await;
                            }
                        }

                        // Clean up disconnected broadcasters periodically
                        Self::cleanup_disconnected_broadcasters_static(&status_broadcasters).await;
                    }
                    _ = cancellation_token.cancelled() => {
                        info!("Snapshot broadcaster received cancellation signal");
                        break;
                    }
                }
            }
        });

        Ok(handle)
    }

    /// Send a snapshot status update for a job
    async fn send_snapshot_update(
        job: &TrainingJob,
        status_broadcasters: &Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,
    ) {
        let snapshot_update = TrainingStatusUpdate {
            job_id: job.id,
            status: job.status.clone(),
            progress_percentage: job.progress_percentage,
            current_epoch: job.current_epoch,
            total_epochs: job.total_epochs,
            metrics: job.metrics.clone(),
            message: "Snapshot status update".to_string(),
            timestamp: Utc::now(),
            financial_metrics: None,
            resource_usage: ResourceUsage {
                cpu_usage_percent: 0.0, // Snapshot doesn't have real-time resource data
                memory_usage_gb: 0.0,
                gpu_usage_percent: None,
                gpu_memory_usage_gb: None,
                active_workers: 0,
            },
        };

        let broadcasters = status_broadcasters.read().await;
        if let Some(broadcaster) = broadcasters.get(&job.id) {
            // For snapshots, we use send (broadcast channels don't have try_send)
            if let Err(e) = broadcaster.send(snapshot_update) {
                tracing::warn!("Channel send failed (receiver dropped?): {:?}", e);
            }
        }
    }

    /// Static version of cleanup for use in spawned tasks
    async fn cleanup_disconnected_broadcasters_static(
        status_broadcasters: &Arc<RwLock<HashMap<Uuid, broadcast::Sender<TrainingStatusUpdate>>>>,
    ) {
        let mut broadcasters = status_broadcasters.write().await;
        let mut to_remove = Vec::new();

        for (job_id, broadcaster) in broadcasters.iter() {
            if broadcaster.receiver_count() == 0 {
                to_remove.push(*job_id);
            }
        }

        for job_id in to_remove {
            broadcasters.remove(&job_id);
            debug!("Cleaned up broadcaster for job {}", job_id);
        }
    }
}

/// Shutdown the orchestrator
impl Drop for TrainingOrchestrator {
    fn drop(&mut self) {
        info!("Shutting down training orchestrator");

        // Cancel all worker tasks
        for handle in self.worker_handles.drain(..) {
            handle.abort();
        }
    }
}
