foxhunt/services/ml_training_service/src/job_queue.rs

//! Job Queue for ML Training Service
//!
//! This module provides a priority-based job queue with GPU resource management,
//! job cancellation support, and Redis persistence for crash recovery.
//!
//! ## Architecture
//! - Priority queue: DQN/PPO (High) > MAMBA-2/TFT (Medium) > TLOB/LIQUID (Low)
//! - GPU resource management: Semaphore limits concurrent GPU jobs (typically 1)
//! - Redis persistence: Queue state persists for crash recovery
//! - Concurrent-safe: All operations are thread-safe via Arc<Mutex<>> and channels

use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use ml::training_pipeline::ProductionTrainingConfig;
use redis::{AsyncCommands, Client as RedisClient};
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
use std::sync::Arc;
use tokio::sync::{Mutex, Semaphore, SemaphorePermit};
use tracing::{debug, info};
use uuid::Uuid;

/// Job priority levels
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub enum JobPriority {
    Low = 0,
    Medium = 1,
    High = 2,
}

impl JobPriority {
    /// Determine priority from model type
    pub fn from_model_type(model_type: &str) -> Self {
        match model_type {
            "DQN" | "PPO" => JobPriority::High,
            "MAMBA_2" | "TFT" => JobPriority::Medium,
            "TLOB" | "LIQUID" | _ => JobPriority::Low,
        }
    }
}

/// Queued training job with priority
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueuedJob {
    pub job_id: Uuid,
    pub model_type: String,
    pub config: ProductionTrainingConfig,
    pub description: String,
    pub tags: HashMap<String, String>,
    pub priority: JobPriority,
    pub enqueued_at: DateTime<Utc>,
}

impl PartialEq for QueuedJob {
    fn eq(&self, other: &Self) -> bool {
        self.job_id == other.job_id
    }
}

impl Eq for QueuedJob {}

impl PartialOrd for QueuedJob {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for QueuedJob {
    fn cmp(&self, other: &Self) -> Ordering {
        // Higher priority jobs come first
        // If priorities are equal, earlier jobs come first (FIFO within priority)
        match self.priority.cmp(&other.priority) {
            Ordering::Equal => other.enqueued_at.cmp(&self.enqueued_at), // Earlier timestamp = higher priority
            other => other, // Reverse order so higher priority comes first in max-heap
        }
    }
}

/// Job status information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JobStatusInfo {
    pub job_id: Uuid,
    pub model_type: String,
    pub status: String,
    pub priority: JobPriority,
    pub enqueued_at: DateTime<Utc>,
}

/// Queue metrics
#[derive(Debug, Clone)]
pub struct QueueMetrics {
    pub queued_jobs: usize,
    pub processing_jobs: usize,
    pub available_gpu_slots: usize,
    pub total_gpu_slots: usize,
}

/// Job Queue with priority, GPU management, and Redis persistence
#[derive(Clone)]
pub struct JobQueue {
    /// Inner state wrapped in Arc<Mutex<>> for thread-safety
    inner: Arc<Mutex<JobQueueInner>>,
    /// GPU resource semaphore (limits concurrent GPU jobs)
    gpu_semaphore: Arc<Semaphore>,
    /// Total GPU slots available
    total_gpu_slots: usize,
    /// Redis client for persistence (optional)
    redis_client: Option<Arc<RedisClient>>,
    /// Redis namespace for this queue
    redis_namespace: String,
}

/// Internal job queue state
struct JobQueueInner {
    /// Priority queue (max-heap)
    queue: BinaryHeap<QueuedJob>,
    /// Job lookup by ID
    jobs: HashMap<Uuid, QueuedJob>,
    /// Maximum queue capacity
    capacity: usize,
    /// Number of jobs currently processing
    processing_count: usize,
}

impl JobQueue {
    /// Create a new job queue without Redis persistence
    pub async fn new(capacity: usize, gpu_slots: usize) -> Result<Self> {
        info!(
            "Creating job queue with capacity={}, gpu_slots={}",
            capacity, gpu_slots
        );

        Ok(Self {
            inner: Arc::new(Mutex::new(JobQueueInner {
                queue: BinaryHeap::new(),
                jobs: HashMap::new(),
                capacity,
                processing_count: 0,
            })),
            gpu_semaphore: Arc::new(Semaphore::new(gpu_slots)),
            total_gpu_slots: gpu_slots,
            redis_client: None,
            redis_namespace: "ml_training_queue".to_string(),
        })
    }

    /// Create a new job queue with Redis persistence
    pub async fn with_redis(capacity: usize, gpu_slots: usize, redis_url: &str) -> Result<Self> {
        let redis_client = RedisClient::open(redis_url).context("Failed to connect to Redis")?;

        info!("Creating job queue with Redis persistence at {}", redis_url);

        Ok(Self {
            inner: Arc::new(Mutex::new(JobQueueInner {
                queue: BinaryHeap::new(),
                jobs: HashMap::new(),
                capacity,
                processing_count: 0,
            })),
            gpu_semaphore: Arc::new(Semaphore::new(gpu_slots)),
            total_gpu_slots: gpu_slots,
            redis_client: Some(Arc::new(redis_client)),
            redis_namespace: "ml_training_queue".to_string(),
        })
    }

    /// Create a new job queue with Redis persistence and custom namespace
    pub async fn with_redis_namespace(
        capacity: usize,
        gpu_slots: usize,
        redis_url: &str,
        namespace: &str,
    ) -> Result<Self> {
        let redis_client = RedisClient::open(redis_url).context("Failed to connect to Redis")?;

        info!(
            "Creating job queue with Redis persistence (namespace: {})",
            namespace
        );

        Ok(Self {
            inner: Arc::new(Mutex::new(JobQueueInner {
                queue: BinaryHeap::new(),
                jobs: HashMap::new(),
                capacity,
                processing_count: 0,
            })),
            gpu_semaphore: Arc::new(Semaphore::new(gpu_slots)),
            total_gpu_slots: gpu_slots,
            redis_client: Some(Arc::new(redis_client)),
            redis_namespace: namespace.to_string(),
        })
    }

    /// Enqueue a new training job
    pub async fn enqueue(
        &self,
        job_id: Uuid,
        model_type: String,
        config: ProductionTrainingConfig,
        description: String,
        tags: HashMap<String, String>,
    ) -> Result<()> {
        let mut inner = self.inner.lock().await;

        // Check capacity
        if inner.jobs.len() >= inner.capacity {
            return Err(anyhow::anyhow!(
                "Queue is at capacity ({}/{})",
                inner.jobs.len(),
                inner.capacity
            ));
        }

        let priority = JobPriority::from_model_type(&model_type);
        let job = QueuedJob {
            job_id,
            model_type: model_type.clone(),
            config,
            description,
            tags,
            priority,
            enqueued_at: Utc::now(),
        };

        debug!(
            "Enqueuing job {} (model={}, priority={:?})",
            job_id, model_type, priority
        );

        inner.queue.push(job.clone());
        inner.jobs.insert(job_id, job);

        Ok(())
    }

    /// Dequeue the highest priority job
    pub async fn dequeue(&self) -> Result<Option<QueuedJob>> {
        let mut inner = self.inner.lock().await;

        if let Some(job) = inner.queue.pop() {
            inner.jobs.remove(&job.job_id);
            inner.processing_count += 1;

            debug!(
                "Dequeued job {} (model={}, priority={:?})",
                job.job_id, job.model_type, job.priority
            );

            Ok(Some(job))
        } else {
            Ok(None)
        }
    }

    /// Cancel a job by ID
    pub async fn cancel_job(&self, job_id: Uuid) -> Result<bool> {
        let mut inner = self.inner.lock().await;

        if inner.jobs.remove(&job_id).is_some() {
            // Job was in the queue, need to rebuild heap without this job
            let jobs: Vec<QueuedJob> = inner.queue.drain().filter(|j| j.job_id != job_id).collect();

            inner.queue = BinaryHeap::from(jobs);

            info!("Cancelled job {}", job_id);
            Ok(true)
        } else {
            debug!("Job {} not found in queue for cancellation", job_id);
            Ok(false)
        }
    }

    /// Get status of a specific job
    pub async fn get_job_status(&self, job_id: Uuid) -> Result<Option<JobStatusInfo>> {
        let inner = self.inner.lock().await;

        if let Some(job) = inner.jobs.get(&job_id) {
            Ok(Some(JobStatusInfo {
                job_id: job.job_id,
                model_type: job.model_type.clone(),
                status: "queued".to_string(),
                priority: job.priority,
                enqueued_at: job.enqueued_at,
            }))
        } else {
            Ok(None)
        }
    }

    /// List all jobs in the queue
    pub async fn list_jobs(&self) -> Result<Vec<JobStatusInfo>> {
        let inner = self.inner.lock().await;

        let jobs: Vec<JobStatusInfo> = inner
            .jobs
            .values()
            .map(|job| JobStatusInfo {
                job_id: job.job_id,
                model_type: job.model_type.clone(),
                status: "queued".to_string(),
                priority: job.priority,
                enqueued_at: job.enqueued_at,
            })
            .collect();

        Ok(jobs)
    }

    /// Get queue metrics
    pub async fn get_metrics(&self) -> Result<QueueMetrics> {
        let inner = self.inner.lock().await;

        Ok(QueueMetrics {
            queued_jobs: inner.jobs.len(),
            processing_jobs: inner.processing_count,
            available_gpu_slots: self.gpu_semaphore.available_permits(),
            total_gpu_slots: self.total_gpu_slots,
        })
    }

    /// Acquire GPU permit (blocks until available)
    pub async fn acquire_gpu_permit(&self) -> Result<SemaphorePermit<'_>> {
        debug!("Acquiring GPU permit...");
        let permit = self
            .gpu_semaphore
            .acquire()
            .await
            .context("Failed to acquire GPU permit")?;
        debug!("GPU permit acquired");
        Ok(permit)
    }

    /// Persist queue state to Redis
    pub async fn persist_to_redis(&self) -> Result<()> {
        let redis_client = match &self.redis_client {
            Some(client) => client,
            None => {
                debug!("Redis persistence not configured");
                return Ok(());
            },
        };

        let inner = self.inner.lock().await;
        let mut conn = redis_client
            .get_multiplexed_async_connection()
            .await
            .context("Failed to get Redis connection")?;

        // Serialize all jobs
        let jobs: Vec<QueuedJob> = inner.jobs.values().cloned().collect();
        let serialized = serde_json::to_string(&jobs).context("Failed to serialize jobs")?;

        // Store in Redis with namespace
        let key = format!("{}:jobs", self.redis_namespace);
        conn.set::<_, _, ()>(&key, serialized)
            .await
            .context("Failed to store jobs in Redis")?;

        debug!(
            "Persisted {} jobs to Redis (namespace: {})",
            jobs.len(),
            self.redis_namespace
        );

        Ok(())
    }

    /// Restore queue state from Redis
    pub async fn restore_from_redis(&self) -> Result<()> {
        let redis_client = match &self.redis_client {
            Some(client) => client,
            None => {
                debug!("Redis persistence not configured");
                return Ok(());
            },
        };

        let mut conn = redis_client
            .get_multiplexed_async_connection()
            .await
            .context("Failed to get Redis connection")?;

        // Retrieve from Redis
        let key = format!("{}:jobs", self.redis_namespace);
        let serialized: Option<String> = conn
            .get(&key)
            .await
            .context("Failed to retrieve jobs from Redis")?;

        if let Some(data) = serialized {
            let jobs: Vec<QueuedJob> =
                serde_json::from_str(&data).context("Failed to deserialize jobs")?;

            let mut inner = self.inner.lock().await;

            // Clear existing state
            inner.queue.clear();
            inner.jobs.clear();

            // Restore jobs
            for job in jobs {
                inner.queue.push(job.clone());
                inner.jobs.insert(job.job_id, job);
            }

            info!(
                "Restored {} jobs from Redis (namespace: {})",
                inner.jobs.len(),
                self.redis_namespace
            );
        } else {
            info!("No jobs found in Redis to restore");
        }

        Ok(())
    }

    /// Mark a job as completed processing (decrements processing count)
    pub async fn mark_completed(&self, _job_id: Uuid) -> Result<()> {
        let mut inner = self.inner.lock().await;
        if inner.processing_count > 0 {
            inner.processing_count -= 1;
        }
        Ok(())
    }
}

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

    #[test]
    fn test_job_priority_ordering() {
        // Test priority enum ordering
        assert!(JobPriority::High > JobPriority::Medium);
        assert!(JobPriority::Medium > JobPriority::Low);
    }

    #[test]
    fn test_job_priority_from_model_type() {
        assert_eq!(JobPriority::from_model_type("DQN"), JobPriority::High);
        assert_eq!(JobPriority::from_model_type("PPO"), JobPriority::High);
        assert_eq!(JobPriority::from_model_type("MAMBA_2"), JobPriority::Medium);
        assert_eq!(JobPriority::from_model_type("TFT"), JobPriority::Medium);
        assert_eq!(JobPriority::from_model_type("TLOB"), JobPriority::Low);
        assert_eq!(JobPriority::from_model_type("LIQUID"), JobPriority::Low);
        assert_eq!(JobPriority::from_model_type("UNKNOWN"), JobPriority::Low);
    }

    #[test]
    fn test_queued_job_ordering() {
        let config = ProductionTrainingConfig::default();
        let now = Utc::now();

        let job_high = QueuedJob {
            job_id: Uuid::new_v4(),
            model_type: "DQN".to_string(),
            config: config.clone(),
            description: "High priority".to_string(),
            tags: HashMap::new(),
            priority: JobPriority::High,
            enqueued_at: now,
        };

        let job_medium = QueuedJob {
            job_id: Uuid::new_v4(),
            model_type: "MAMBA_2".to_string(),
            config: config.clone(),
            description: "Medium priority".to_string(),
            tags: HashMap::new(),
            priority: JobPriority::Medium,
            enqueued_at: now,
        };

        let job_low = QueuedJob {
            job_id: Uuid::new_v4(),
            model_type: "TLOB".to_string(),
            config: config.clone(),
            description: "Low priority".to_string(),
            tags: HashMap::new(),
            priority: JobPriority::Low,
            enqueued_at: now,
        };

        // Higher priority jobs should be "greater" (come first in max-heap)
        assert!(job_high > job_medium);
        assert!(job_medium > job_low);
        assert!(job_high > job_low);
    }

    #[test]
    fn test_queued_job_fifo_within_priority() {
        use chrono::Duration;

        let config = ProductionTrainingConfig::default();
        let now = Utc::now();

        let job_early = QueuedJob {
            job_id: Uuid::new_v4(),
            model_type: "DQN".to_string(),
            config: config.clone(),
            description: "Earlier job".to_string(),
            tags: HashMap::new(),
            priority: JobPriority::High,
            enqueued_at: now,
        };

        let job_later = QueuedJob {
            job_id: Uuid::new_v4(),
            model_type: "DQN".to_string(),
            config: config.clone(),
            description: "Later job".to_string(),
            tags: HashMap::new(),
            priority: JobPriority::High,
            enqueued_at: now + Duration::seconds(10),
        };

        // Earlier job should come first within same priority
        assert!(job_early > job_later);
    }
}