foxhunt/services/ml_training_service/src/job_spawner.rs

//! Job Spawner - Multi-Model Training Orchestration
//!
//! This module provides production-ready job spawning and orchestration for ML training.
//! It manages batch jobs that spawn multiple child training jobs across different models
//! and assets, with support for sequential execution, progress tracking, and rollback.
//!
//! ## Architecture
//!
//! - **BatchJob**: Parent job that orchestrates multiple child jobs
//! - **ChildJob**: Individual training job for a specific model and asset
//! - **Sequential Execution**: Jobs are executed DQN → PPO → MAMBA-2 → TFT to avoid GPU OOM
//! - **Database Persistence**: All jobs stored in PostgreSQL with transaction support
//! - **Rollback Safety**: Failed batch creation rolls back all changes
//!
//! ## Usage
//!
//! ```rust,no_run
//! use ml_training_service::job_spawner::{JobSpawner, Asset, ModelType};
//! use sqlx::PgPool;
//! use std::path::PathBuf;
//!
//! async fn example(pool: PgPool) -> anyhow::Result<()> {
//!     let spawner = JobSpawner::new(pool);
//!     
//!     let assets = vec![
//!         Asset {
//!             symbol: "ES.FUT".to_string(),
//!             data_file: PathBuf::from("/data/ES_FUT_180d.parquet"),
//!         },
//!     ];
//!     
//!     let models = vec![
//!         ModelType::DQN,
//!         ModelType::PPO,
//!         ModelType::MAMBA,
//!         ModelType::TFT,
//!     ];
//!     
//!     let batch = spawner.spawn_batch(assets, models).await?;
//!     println!("Created batch: {}", batch.batch_id);
//!     
//!     Ok(())
//! }
//! ```

use anyhow::{anyhow, Context, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::{PgPool, Postgres, Transaction};
use std::path::PathBuf;
use tracing::{debug, info, warn};
use uuid::Uuid;

// Use canonical ModelType from common crate
pub use common::model_types::ModelType;

/// Trading asset with data file
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Asset {
    /// Asset symbol (e.g., "ES.FUT", "NQ.FUT")
    pub symbol: String,
    /// Path to training data file
    pub data_file: PathBuf,
}

/// Parent batch job that orchestrates multiple child jobs
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchJob {
    /// Unique batch job identifier
    pub batch_id: Uuid,
    /// Assets to train on
    pub assets: Vec<Asset>,
    /// Models to train
    pub models: Vec<ModelType>,
    /// Batch status
    pub status: String,
    /// Creation timestamp
    pub created_at: DateTime<Utc>,
}

/// Individual child training job
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
pub struct ChildJob {
    /// Unique job identifier
    pub id: Uuid,
    /// Parent batch identifier
    pub batch_id: Uuid,
    /// Model type
    pub model_type: String,
    /// Job status
    pub status: String,
    /// Creation timestamp
    pub created_at: DateTime<Utc>,
    /// Job configuration JSON
    pub config_json: serde_json::Value,
}

/// Batch status summary
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchStatus {
    /// Batch identifier
    pub batch_id: Uuid,
    /// Overall status
    pub status: String,
    /// Total number of jobs
    pub total_jobs: i32,
    /// Pending jobs count
    pub pending_jobs: i32,
    /// Running jobs count
    pub running_jobs: i32,
    /// Completed jobs count
    pub completed_jobs: i32,
    /// Failed jobs count
    pub failed_jobs: i32,
    /// Overall progress percentage
    pub overall_progress: f64,
}

/// Job spawner for batch training orchestration
pub struct JobSpawner {
    db_pool: PgPool,
}

impl JobSpawner {
    /// Create a new job spawner
    ///
    /// # Arguments
    ///
    /// * `db_pool` - PostgreSQL connection pool
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use ml_training_service::job_spawner::JobSpawner;
    /// use sqlx::PgPool;
    ///
    /// async fn example(pool: PgPool) {
    ///     let spawner = JobSpawner::new(pool);
    /// }
    /// ```
    pub fn new(db_pool: PgPool) -> Self {
        Self { db_pool }
    }

    /// Spawn a batch job with multiple child training jobs
    ///
    /// Creates a parent batch job and spawns child jobs for each asset × model combination.
    /// Jobs are ordered sequentially (DQN → PPO → MAMBA-2 → TFT) to avoid GPU OOM.
    /// All operations are atomic - if any step fails, the entire batch is rolled back.
    ///
    /// # Arguments
    ///
    /// * `assets` - List of trading assets to train on
    /// * `models` - List of models to train (in execution order)
    ///
    /// # Returns
    ///
    /// * `Ok(BatchJob)` - Created batch job with all metadata
    /// * `Err` - Database error or validation failure
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Assets list is empty
    /// - Models list is empty
    /// - Database transaction fails
    /// - Child job creation fails
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use ml_training_service::job_spawner::{JobSpawner, Asset, ModelType};
    /// use sqlx::PgPool;
    /// use std::path::PathBuf;
    ///
    /// async fn example(pool: PgPool) -> anyhow::Result<()> {
    ///     let spawner = JobSpawner::new(pool);
    ///     
    ///     let assets = vec![
    ///         Asset {
    ///             symbol: "ES.FUT".to_string(),
    ///             data_file: PathBuf::from("/data/ES.parquet"),
    ///         },
    ///     ];
    ///     
    ///     let models = vec![ModelType::DQN, ModelType::PPO];
    ///     
    ///     let batch = spawner.spawn_batch(assets, models).await?;
    ///     println!("Batch ID: {}", batch.batch_id);
    ///     
    ///     Ok(())
    /// }
    /// ```
    pub async fn spawn_batch(
        &self,
        assets: Vec<Asset>,
        models: Vec<ModelType>,
    ) -> Result<BatchJob> {
        // Validation
        if assets.is_empty() {
            return Err(anyhow!("Assets list cannot be empty"));
        }
        if models.is_empty() {
            return Err(anyhow!("Models list cannot be empty"));
        }

        info!(
            "Spawning batch job: {} assets × {} models = {} jobs",
            assets.len(),
            models.len(),
            assets.len() * models.len()
        );

        // Start database transaction for atomicity
        let mut tx = self
            .db_pool
            .begin()
            .await
            .context("Failed to begin transaction")?;

        // Create parent batch job
        let batch_id = Uuid::new_v4();
        let created_at = Utc::now();
        let total_jobs = (assets.len() * models.len()) as i32;

        let batch = self
            .create_batch_job(&mut tx, batch_id, &assets, &models, total_jobs, created_at)
            .await
            .context("Failed to create batch job")?;

        // Create child jobs for each asset × model combination
        let mut created_count = 0;
        for asset in &assets {
            for model in &models {
                self.create_child_job(&mut tx, batch_id, asset, model, created_at)
                    .await
                    .with_context(|| {
                        format!(
                            "Failed to create child job for {} / {:?}",
                            asset.symbol, model
                        )
                    })?;
                created_count += 1;
            }
        }

        // Commit transaction
        tx.commit()
            .await
            .context("Failed to commit batch job transaction")?;

        info!(
            "Successfully created batch {} with {} child jobs",
            batch_id, created_count
        );

        Ok(batch)
    }

    /// Create the parent batch job record
    async fn create_batch_job(
        &self,
        tx: &mut Transaction<'_, Postgres>,
        batch_id: Uuid,
        assets: &[Asset],
        models: &[ModelType],
        total_jobs: i32,
        created_at: DateTime<Utc>,
    ) -> Result<BatchJob> {
        // Serialize assets and models to JSON for storage
        let assets_json = serde_json::to_value(assets)
            .context("Failed to serialize assets")?;
        let models_json = serde_json::to_value(models)
            .context("Failed to serialize models")?;

        // Build batch name
        let asset_symbols: Vec<_> = assets.iter().map(|a| a.symbol.as_str()).collect();
        let batch_name = format!(
            "Batch: {} ({})",
            asset_symbols.join(", "),
            models.len()
        );

        // Build config JSON
        let config = serde_json::json!({
            "assets": assets_json,
            "models": models_json,
        });

        // Serialize config to string for JSONB column
        let config_str = serde_json::to_string(&config)
            .context("Failed to serialize batch config")?;

        // Insert batch job
        sqlx::query(
            r#"
            INSERT INTO batch_jobs (
                id, name, description, status, total_jobs,
                pending_jobs, config_json, created_at
            )
            VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8)
            "#,
        )
        .bind(batch_id)
        .bind(&batch_name)
        .bind("Multi-model training batch")
        .bind("Pending")
        .bind(total_jobs)
        .bind(total_jobs) // All jobs start as pending
        .bind(&config_str)
        .bind(created_at)
        .execute(&mut **tx)
        .await
        .context("Failed to insert batch job")?;

        debug!("Created batch job: {}", batch_id);

        Ok(BatchJob {
            batch_id,
            assets: assets.to_vec(),
            models: models.to_vec(),
            status: "Pending".to_string(),
            created_at,
        })
    }

    /// Create a child training job
    async fn create_child_job(
        &self,
        tx: &mut Transaction<'_, Postgres>,
        batch_id: Uuid,
        asset: &Asset,
        model: &ModelType,
        created_at: DateTime<Utc>,
    ) -> Result<()> {
        let job_id = Uuid::new_v4();
        let model_str = model.to_db_string();
        let model_weight = model.weight();

        // Build job configuration
        let config = serde_json::json!({
            "asset": asset.symbol,
            "data_file": asset.data_file.to_string_lossy(),
            "model_type": model_str,
        });

        let config_str = serde_json::to_string(&config)
            .context("Failed to serialize job config")?;

        // Insert child job
        sqlx::query(
            r#"
            INSERT INTO child_jobs (
                id, batch_id, model_type, model_weight,
                status, config_json, created_at
            )
            VALUES ($1, $2, $3, $4, $5, $6::jsonb, $7)
            "#,
        )
        .bind(job_id)
        .bind(batch_id)
        .bind(model_str)
        .bind(model_weight)
        .bind("Pending")
        .bind(config_str)
        .bind(created_at)
        .execute(&mut **tx)
        .await
        .context("Failed to insert child job")?;

        debug!(
            "Created child job {} for {} / {}",
            job_id, asset.symbol, model_str
        );

        Ok(())
    }

    /// Get the next pending job in FIFO order
    ///
    /// Returns the oldest pending job across all batches.
    /// This ensures jobs are processed in the order they were created.
    ///
    /// # Returns
    ///
    /// * `Ok(Some(ChildJob))` - Next pending job
    /// * `Ok(None)` - No pending jobs in queue
    /// * `Err` - Database error
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use ml_training_service::job_spawner::JobSpawner;
    /// use sqlx::PgPool;
    ///
    /// async fn example(pool: PgPool) -> anyhow::Result<()> {
    ///     let spawner = JobSpawner::new(pool);
    ///     
    ///     if let Some(job) = spawner.get_next_pending_job().await? {
    ///         println!("Processing job: {}", job.id);
    ///     } else {
    ///         println!("No pending jobs");
    ///     }
    ///     
    ///     Ok(())
    /// }
    /// ```
    pub async fn get_next_pending_job(&self) -> Result<Option<ChildJob>> {
        let job = sqlx::query_as::<_, ChildJob>(
            r#"
            SELECT id, batch_id, model_type, status, created_at, config_json
            FROM child_jobs
            WHERE status = 'Pending'
            ORDER BY created_at ASC
            LIMIT 1
            "#,
        )
        .fetch_optional(&self.db_pool)
        .await
        .context("Failed to fetch next pending job")?;

        if let Some(ref j) = job {
            debug!("Next pending job: {} ({})", j.id, j.model_type);
        } else {
            debug!("No pending jobs in queue");
        }

        Ok(job)
    }

    /// Get batch status and progress
    ///
    /// # Arguments
    ///
    /// * `batch_id` - Batch job identifier
    ///
    /// # Returns
    ///
    /// * `Ok(Some(BatchStatus))` - Batch status summary
    /// * `Ok(None)` - Batch not found
    /// * `Err` - Database error
    pub async fn get_batch_status(&self, batch_id: Uuid) -> Result<Option<BatchStatus>> {
        let status = sqlx::query_as::<_, (Uuid, String, i32, i32, i32, i32, i32, f64)>(
            r#"
            SELECT 
                id, status, total_jobs, pending_jobs, running_jobs,
                completed_jobs, failed_jobs, overall_progress
            FROM batch_jobs
            WHERE id = $1
            "#,
        )
        .bind(batch_id)
        .fetch_optional(&self.db_pool)
        .await
        .context("Failed to fetch batch status")?;

        Ok(status.map(|(id, status, total, pending, running, completed, failed, progress)| {
            BatchStatus {
                batch_id: id,
                status,
                total_jobs: total,
                pending_jobs: pending,
                running_jobs: running,
                completed_jobs: completed,
                failed_jobs: failed,
                overall_progress: progress,
            }
        }))
    }

    /// Update child job status
    ///
    /// # Arguments
    ///
    /// * `job_id` - Child job identifier
    /// * `new_status` - New status ("Running", "Completed", "Failed", etc.)
    ///
    /// # Returns
    ///
    /// * `Ok(())` - Status updated successfully
    /// * `Err` - Database error or job not found
    pub async fn update_job_status(&self, job_id: Uuid, new_status: &str) -> Result<()> {
        let rows_affected = sqlx::query(
            r#"
            UPDATE child_jobs
            SET status = $1, updated_at = NOW()
            WHERE id = $2
            "#,
        )
        .bind(new_status)
        .bind(job_id)
        .execute(&self.db_pool)
        .await
        .context("Failed to update job status")?
        .rows_affected();

        if rows_affected == 0 {
            warn!("Job not found: {}", job_id);
            return Err(anyhow!("Job not found: {}", job_id));
        }

        info!("Updated job {} status to {}", job_id, new_status);
        Ok(())
    }

    /// Fetch a single child job by its ID
    ///
    /// # Arguments
    ///
    /// * `job_id` - Child job identifier
    ///
    /// # Returns
    ///
    /// * `Ok(Some(ChildJob))` - The child job
    /// * `Ok(None)` - Job not found
    /// * `Err` - Database error
    pub async fn get_job_by_id(&self, job_id: Uuid) -> Result<Option<ChildJob>> {
        let job = sqlx::query_as::<_, ChildJob>(
            r#"
            SELECT id, batch_id, model_type, status, created_at, config_json
            FROM child_jobs
            WHERE id = $1
            "#,
        )
        .bind(job_id)
        .fetch_optional(&self.db_pool)
        .await
        .context("Failed to fetch job by id")?;

        Ok(job)
    }

    /// Get all child jobs for a batch
    ///
    /// # Arguments
    ///
    /// * `batch_id` - Batch job identifier
    ///
    /// # Returns
    ///
    /// * `Ok(Vec<ChildJob>)` - List of child jobs
    /// * `Err` - Database error
    pub async fn get_batch_jobs(&self, batch_id: Uuid) -> Result<Vec<ChildJob>> {
        let jobs = sqlx::query_as::<_, ChildJob>(
            r#"
            SELECT id, batch_id, model_type, status, created_at, config_json
            FROM child_jobs
            WHERE batch_id = $1
            ORDER BY created_at
            "#,
        )
        .bind(batch_id)
        .fetch_all(&self.db_pool)
        .await
        .context("Failed to fetch batch jobs")?;

        Ok(jobs)
    }
}

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

    #[test]
    fn test_model_type_to_db_string() {
        assert_eq!(ModelType::DQN.to_db_string(), "DQN");
        assert_eq!(ModelType::PPO.to_db_string(), "PPO");
        assert_eq!(ModelType::MAMBA.to_db_string(), "MAMBA-2");
        assert_eq!(ModelType::TFT.to_db_string(), "TFT");
    }

    #[test]
    fn test_model_type_weight() {
        assert_eq!(ModelType::DQN.weight(), 0.25);
        assert_eq!(ModelType::PPO.weight(), 0.25);
        assert_eq!(ModelType::MAMBA.weight(), 0.25);
        assert_eq!(ModelType::TFT.weight(), 0.25);
    }
}