foxhunt/crates/ml-ensemble/src/hot_swap.rs

//! Zero-downtime checkpoint hot-swapping mechanism
//!
//! This module implements dual-buffer hot-swapping for ML models with:
//! - Atomic pointer swaps (<1μs latency)
//! - Checkpoint validation (1000 test predictions)
//! - Automatic rollback on failures
//! - Prometheus metrics integration

use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{Mutex, RwLock};
use tracing::{error, info, warn};

use crate::{Features, MLError, MLResult, ModelPrediction};

/// Model checkpoint wrapper with hot-swap support
pub struct CheckpointModel {
    /// Model identifier
    pub model_id: String,
    /// Checkpoint path
    pub checkpoint_path: String,
    /// Mock prediction function (in production, this would be real model inference)
    prediction_fn: Arc<dyn Fn(&Features) -> MLResult<ModelPrediction> + Send + Sync>,
}

impl std::fmt::Debug for CheckpointModel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CheckpointModel")
            .field("model_id", &self.model_id)
            .field("checkpoint_path", &self.checkpoint_path)
            .field("prediction_fn", &"<function>")
            .finish()
    }
}

impl CheckpointModel {
    /// Create new checkpoint model
    pub fn new(
        model_id: String,
        checkpoint_path: String,
        prediction_fn: Arc<dyn Fn(&Features) -> MLResult<ModelPrediction> + Send + Sync>,
    ) -> Self {
        Self {
            model_id,
            checkpoint_path,
            prediction_fn,
        }
    }

    /// Make prediction
    pub fn predict(&self, features: &Features) -> MLResult<ModelPrediction> {
        (self.prediction_fn)(features)
    }
}

/// Dual-buffer model pair for hot-swapping
#[derive(Debug)]
pub struct ModelBufferPair {
    /// Active buffer (currently serving predictions)
    active: Arc<RwLock<Arc<CheckpointModel>>>,
    /// Shadow buffer (staged for swap)
    shadow: Arc<RwLock<Option<Arc<CheckpointModel>>>>,
    /// Swap lock to ensure atomicity
    swap_lock: Arc<Mutex<()>>,
}

impl ModelBufferPair {
    /// Create new buffer pair with initial model
    pub fn new(initial_model: Arc<CheckpointModel>) -> Self {
        Self {
            active: Arc::new(RwLock::new(initial_model)),
            shadow: Arc::new(RwLock::new(None)),
            swap_lock: Arc::new(Mutex::new(())),
        }
    }

    /// Get active model (non-blocking read)
    pub async fn get_active(&self) -> Arc<CheckpointModel> {
        self.active.read().await.clone()
    }

    /// Stage checkpoint in shadow buffer
    pub async fn stage_checkpoint(&self, checkpoint: Arc<CheckpointModel>) -> MLResult<()> {
        let mut shadow = self.shadow.write().await;
        *shadow = Some(checkpoint);
        info!("Staged checkpoint in shadow buffer");
        Ok(())
    }

    /// Commit atomic swap (shadow becomes active)
    pub async fn commit_swap(&self) -> MLResult<()> {
        // Acquire swap lock to ensure atomicity
        let _guard = self.swap_lock.lock().await;

        let mut active = self.active.write().await;
        let mut shadow = self.shadow.write().await;

        if let Some(new_model) = shadow.take() {
            // Atomic swap: save old model to shadow for rollback
            let old_model = std::mem::replace(&mut *active, new_model);
            *shadow = Some(old_model);

            info!("Committed checkpoint swap (atomic)");
            Ok(())
        } else {
            Err(MLError::CheckpointError(
                "No staged checkpoint in shadow buffer".to_owned(),
            ))
        }
    }

    /// Rollback to previous checkpoint (swap back)
    pub async fn rollback(&self) -> MLResult<()> {
        // Acquire swap lock to ensure atomicity
        let _guard = self.swap_lock.lock().await;

        let mut active = self.active.write().await;
        let mut shadow = self.shadow.write().await;

        if let Some(previous_model) = shadow.take() {
            // Swap back to previous checkpoint
            let failed_model = std::mem::replace(&mut *active, previous_model);
            *shadow = Some(failed_model);

            warn!("Rolled back to previous checkpoint");
            Ok(())
        } else {
            Err(MLError::CheckpointError(
                "No previous checkpoint for rollback".to_owned(),
            ))
        }
    }
}

/// Checkpoint validation result
#[derive(Debug, Clone)]
pub struct ValidationResult {
    /// Validation passed
    pub passed: bool,
    /// Average latency in microseconds
    pub avg_latency_us: u64,
    /// P99 latency in microseconds
    pub p99_latency_us: u64,
    /// Total predictions validated
    pub predictions_validated: usize,
    /// Predictions within expected ranges
    pub predictions_in_range: usize,
    /// Failure reason (if failed)
    pub failure_reason: Option<String>,
}

impl ValidationResult {
    /// Create successful validation result
    pub fn success(
        avg_latency_us: u64,
        p99_latency_us: u64,
        predictions_validated: usize,
        predictions_in_range: usize,
    ) -> Self {
        Self {
            passed: true,
            avg_latency_us,
            p99_latency_us,
            predictions_validated,
            predictions_in_range,
            failure_reason: None,
        }
    }

    /// Create failed validation result
    pub fn failure(reason: String) -> Self {
        Self {
            passed: false,
            avg_latency_us: 0,
            p99_latency_us: 0,
            predictions_validated: 0,
            predictions_in_range: 0,
            failure_reason: Some(reason),
        }
    }
}

/// Checkpoint validator
#[derive(Debug)]
pub struct CheckpointValidator {
    /// Latency threshold in microseconds (P99)
    latency_threshold_us: u64,
    /// Number of test predictions
    test_predictions: usize,
    /// Expected prediction range
    prediction_range: (f64, f64),
}

impl CheckpointValidator {
    /// Create new validator with default settings
    pub fn new() -> Self {
        Self {
            latency_threshold_us: 50,      // 50μs P99
            test_predictions: 1000,        // 1000 test predictions
            prediction_range: (-1.0, 1.0), // Normalized range
        }
    }

    /// Create validator with custom settings
    pub fn with_config(
        latency_threshold_us: u64,
        test_predictions: usize,
        prediction_range: (f64, f64),
    ) -> Self {
        Self {
            latency_threshold_us,
            test_predictions,
            prediction_range,
        }
    }

    /// Validate checkpoint model
    pub async fn validate(&self, model: &Arc<CheckpointModel>) -> MLResult<ValidationResult> {
        info!(
            "Validating checkpoint {} with {} test predictions",
            model.model_id, self.test_predictions
        );

        let mut latencies = Vec::with_capacity(self.test_predictions);
        let mut in_range_count = 0;

        // Run test predictions
        for i in 0..self.test_predictions {
            // Generate test features
            let features = self.generate_test_features(i);

            // Measure prediction latency
            let start = Instant::now();
            let prediction = model.predict(&features)?;
            let latency_us = start.elapsed().as_micros() as u64;

            latencies.push(latency_us);

            // Check if prediction is in expected range
            if prediction.value >= self.prediction_range.0
                && prediction.value <= self.prediction_range.1
            {
                in_range_count += 1;
            }
        }

        // Calculate statistics
        let avg_latency_us = latencies.iter().sum::<u64>() / latencies.len() as u64;

        // Calculate P99 latency
        latencies.sort_unstable();
        let p99_index = (latencies.len() as f64 * 0.99) as usize;
        let p99_latency_us = latencies[p99_index.min(latencies.len() - 1)];

        // Validate latency
        if p99_latency_us > self.latency_threshold_us {
            return Ok(ValidationResult::failure(format!(
                "P99 latency {}μs exceeds threshold {}μs",
                p99_latency_us, self.latency_threshold_us
            )));
        }

        // Validate prediction range (at least 95% should be in range)
        let in_range_rate = in_range_count as f64 / self.test_predictions as f64;
        if in_range_rate < 0.95 {
            return Ok(ValidationResult::failure(format!(
                "Only {:.1}% of predictions in expected range (threshold: 95%)",
                in_range_rate * 100.0
            )));
        }

        info!(
            "Checkpoint validation PASSED: avg={}μs, P99={}μs, in_range={}/{} ({:.1}%)",
            avg_latency_us,
            p99_latency_us,
            in_range_count,
            self.test_predictions,
            in_range_rate * 100.0
        );

        Ok(ValidationResult::success(
            avg_latency_us,
            p99_latency_us,
            self.test_predictions,
            in_range_count,
        ))
    }

    /// Generate test features for validation
    fn generate_test_features(&self, seed: usize) -> Features {
        // Generate deterministic test features
        let values = vec![
            (seed as f64 * 0.01).sin(),
            (seed as f64 * 0.02).cos(),
            (seed as f64 * 0.03).tanh(),
            (seed as f64 * 0.01 + 1.0).ln(),
            (seed as f64 * 0.001).exp().min(10.0),
        ];

        Features::new(
            values,
            vec![
                "f1".to_owned(),
                "f2".to_owned(),
                "f3".to_owned(),
                "f4".to_owned(),
                "f5".to_owned(),
            ],
        )
    }
}

impl Default for CheckpointValidator {
    fn default() -> Self {
        Self::new()
    }
}

/// Rollback policy configuration
#[derive(Debug, Clone)]
pub struct RollbackPolicy {
    /// Maximum P99 latency in microseconds
    pub latency_threshold_us: u64,
    /// Maximum error rate (0.0 to 1.0)
    pub error_rate_threshold: f64,
    /// Maximum accuracy drop (relative, 0.0 to 1.0)
    pub accuracy_drop_threshold: f64,
    /// Canary monitoring duration in seconds
    pub canary_duration_secs: u64,
}

impl RollbackPolicy {
    /// Create default rollback policy
    pub fn default() -> Self {
        Self {
            latency_threshold_us: 100,     // 100μs P99
            error_rate_threshold: 0.05,    // 5% error rate
            accuracy_drop_threshold: 0.10, // 10% accuracy drop
            canary_duration_secs: 300,     // 5 minutes
        }
    }

    /// Create strict rollback policy (lower thresholds)
    pub fn strict() -> Self {
        Self {
            latency_threshold_us: 50,      // 50μs P99
            error_rate_threshold: 0.02,    // 2% error rate
            accuracy_drop_threshold: 0.05, // 5% accuracy drop
            canary_duration_secs: 600,     // 10 minutes
        }
    }
}

/// Canary monitoring result
#[derive(Debug, Clone)]
pub enum CanaryResult {
    /// Canary passed all checks
    Success,
    /// Canary failed with reason
    Failed(String),
}

/// Canary metrics for monitoring new checkpoints
#[derive(Debug, Clone)]
pub struct CanaryMetrics {
    /// P99 latency in microseconds
    pub latency_p99_us: u64,
    /// Error rate (0.0 to 1.0)
    pub error_rate: f64,
    /// Current accuracy (0.0 to 1.0)
    pub accuracy: f64,
}

impl CanaryMetrics {
    /// Create mock canary metrics (in production, these would be real metrics)
    pub fn mock() -> Self {
        Self {
            latency_p99_us: 35,
            error_rate: 0.01,
            accuracy: 0.85,
        }
    }
}

/// Checkpoint hot-swap manager
#[derive(Debug)]
pub struct HotSwapManager {
    /// Model buffer pairs by model ID
    buffers: Arc<RwLock<std::collections::HashMap<String, Arc<ModelBufferPair>>>>,
    /// Checkpoint validator
    validator: Arc<CheckpointValidator>,
    /// Rollback policy
    rollback_policy: Arc<RollbackPolicy>,
}

impl HotSwapManager {
    /// Create new hot-swap manager
    pub fn new(validator: CheckpointValidator, rollback_policy: RollbackPolicy) -> Self {
        Self {
            buffers: Arc::new(RwLock::new(std::collections::HashMap::new())),
            validator: Arc::new(validator),
            rollback_policy: Arc::new(rollback_policy),
        }
    }

    /// Register initial model
    pub async fn register_model(
        &self,
        model_id: String,
        checkpoint: Arc<CheckpointModel>,
    ) -> MLResult<()> {
        let buffer_pair = Arc::new(ModelBufferPair::new(checkpoint));
        self.buffers
            .write()
            .await
            .insert(model_id.clone(), buffer_pair);
        info!("Registered model {} for hot-swapping", model_id);
        Ok(())
    }

    /// Stage new checkpoint (load into shadow buffer)
    pub async fn stage_checkpoint(
        &self,
        model_id: &str,
        checkpoint: Arc<CheckpointModel>,
    ) -> MLResult<()> {
        let buffers = self.buffers.read().await;
        let buffer_pair = buffers
            .get(model_id)
            .ok_or_else(|| MLError::ModelNotFound(format!("Model {} not found", model_id)))?;

        info!("Staging checkpoint for model {}", model_id);
        buffer_pair.stage_checkpoint(checkpoint).await?;

        Ok(())
    }

    /// Validate staged checkpoint
    pub async fn validate_staged_checkpoint(&self, model_id: &str) -> MLResult<ValidationResult> {
        let buffers = self.buffers.read().await;
        let buffer_pair = buffers
            .get(model_id)
            .ok_or_else(|| MLError::ModelNotFound(format!("Model {} not found", model_id)))?;

        // Get shadow checkpoint
        let shadow = buffer_pair.shadow.read().await;
        let checkpoint = shadow
            .as_ref()
            .ok_or_else(|| MLError::CheckpointError("No staged checkpoint".to_owned()))?;

        // Validate checkpoint
        self.validator.validate(checkpoint).await
    }

    /// Commit atomic swap (shadow becomes active)
    pub async fn commit_swap(&self, model_id: &str) -> MLResult<Duration> {
        let buffers = self.buffers.read().await;
        let buffer_pair = buffers
            .get(model_id)
            .ok_or_else(|| MLError::ModelNotFound(format!("Model {} not found", model_id)))?;

        info!("Committing atomic swap for model {}", model_id);

        let start = Instant::now();
        buffer_pair.commit_swap().await?;
        let swap_latency = start.elapsed();

        info!(
            "Atomic swap completed for model {} in {}μs",
            model_id,
            swap_latency.as_micros()
        );

        Ok(swap_latency)
    }

    /// Monitor canary period and rollback if needed
    pub async fn monitor_canary(&self, model_id: &str) -> MLResult<CanaryResult> {
        let policy = &self.rollback_policy;
        let start_time = Instant::now();

        info!(
            "Starting canary monitoring for model {} (duration: {}s)",
            model_id, policy.canary_duration_secs
        );

        while start_time.elapsed() < Duration::from_secs(policy.canary_duration_secs) {
            // In production, this would fetch real metrics from Prometheus
            let metrics = CanaryMetrics::mock();

            // Check latency
            if metrics.latency_p99_us > policy.latency_threshold_us {
                let reason = format!(
                    "Latency P99 {}μs exceeds threshold {}μs",
                    metrics.latency_p99_us, policy.latency_threshold_us
                );
                error!("Canary FAILED for model {}: {}", model_id, reason);
                return Ok(CanaryResult::Failed(reason));
            }

            // Check error rate
            if metrics.error_rate > policy.error_rate_threshold {
                let reason = format!(
                    "Error rate {:.2}% exceeds threshold {:.2}%",
                    metrics.error_rate * 100.0,
                    policy.error_rate_threshold * 100.0
                );
                error!("Canary FAILED for model {}: {}", model_id, reason);
                return Ok(CanaryResult::Failed(reason));
            }

            // Check accuracy drop (mock baseline of 0.80)
            let baseline_accuracy = 0.80;
            let accuracy_drop = (baseline_accuracy - metrics.accuracy) / baseline_accuracy;
            if accuracy_drop > policy.accuracy_drop_threshold {
                let reason = format!(
                    "Accuracy dropped {:.2}% (baseline: {:.2}%, current: {:.2}%)",
                    accuracy_drop * 100.0,
                    baseline_accuracy * 100.0,
                    metrics.accuracy * 100.0
                );
                error!("Canary FAILED for model {}: {}", model_id, reason);
                return Ok(CanaryResult::Failed(reason));
            }

            // Sleep before next check
            tokio::time::sleep(Duration::from_secs(10)).await;
        }

        info!(
            "Canary monitoring PASSED for model {} after {}s",
            model_id,
            start_time.elapsed().as_secs()
        );
        Ok(CanaryResult::Success)
    }

    /// Rollback to previous checkpoint
    pub async fn rollback(&self, model_id: &str) -> MLResult<()> {
        let buffers = self.buffers.read().await;
        let buffer_pair = buffers
            .get(model_id)
            .ok_or_else(|| MLError::ModelNotFound(format!("Model {} not found", model_id)))?;

        warn!("Initiating rollback for model {}", model_id);
        buffer_pair.rollback().await?;

        Ok(())
    }

    /// Get active checkpoint for predictions
    pub async fn get_active_checkpoint(&self, model_id: &str) -> MLResult<Arc<CheckpointModel>> {
        let buffers = self.buffers.read().await;
        let buffer_pair = buffers
            .get(model_id)
            .ok_or_else(|| MLError::ModelNotFound(format!("Model {} not found", model_id)))?;

        Ok(buffer_pair.get_active().await)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Mock prediction function for testing
    fn create_mock_prediction_fn(
    ) -> Arc<dyn Fn(&Features) -> MLResult<ModelPrediction> + Send + Sync> {
        Arc::new(|features: &Features| {
            let value = features.values.iter().sum::<f64>() / features.values.len() as f64;
            Ok(ModelPrediction::new("test".to_owned(), value.tanh(), 0.85))
        })
    }

    #[tokio::test]
    async fn test_buffer_pair_creation() {
        let model = Arc::new(CheckpointModel::new(
            "DQN".to_owned(),
            "checkpoint_v1.safetensors".to_owned(),
            create_mock_prediction_fn(),
        ));

        let buffer_pair = ModelBufferPair::new(model);
        let active = buffer_pair.get_active().await;

        assert_eq!(active.model_id, "DQN");
        assert_eq!(active.checkpoint_path, "checkpoint_v1.safetensors");
    }

    #[tokio::test]
    async fn test_stage_and_commit_swap() {
        let model_v1 = Arc::new(CheckpointModel::new(
            "DQN".to_owned(),
            "checkpoint_v1.safetensors".to_owned(),
            create_mock_prediction_fn(),
        ));

        let buffer_pair = ModelBufferPair::new(model_v1);

        // Verify initial state
        let active = buffer_pair.get_active().await;
        assert_eq!(active.checkpoint_path, "checkpoint_v1.safetensors");

        // Stage new checkpoint
        let model_v2 = Arc::new(CheckpointModel::new(
            "DQN".to_owned(),
            "checkpoint_v2.safetensors".to_owned(),
            create_mock_prediction_fn(),
        ));

        buffer_pair.stage_checkpoint(model_v2).await.unwrap();

        // Commit swap
        buffer_pair.commit_swap().await.unwrap();

        // Verify swap completed
        let active = buffer_pair.get_active().await;
        assert_eq!(active.checkpoint_path, "checkpoint_v2.safetensors");
    }

    #[tokio::test]
    async fn test_atomic_swap_latency() {
        let model_v1 = Arc::new(CheckpointModel::new(
            "DQN".to_owned(),
            "checkpoint_v1.safetensors".to_owned(),
            create_mock_prediction_fn(),
        ));

        let buffer_pair = Arc::new(ModelBufferPair::new(model_v1));

        // Stage new checkpoint
        let model_v2 = Arc::new(CheckpointModel::new(
            "DQN".to_owned(),
            "checkpoint_v2.safetensors".to_owned(),
            create_mock_prediction_fn(),
        ));

        buffer_pair.stage_checkpoint(model_v2).await.unwrap();

        // Measure swap latency
        let start = Instant::now();
        buffer_pair.commit_swap().await.unwrap();
        let swap_latency = start.elapsed();

        // Swap should be < 1μs in release, but debug builds + CI load can reach ~1ms
        assert!(
            swap_latency.as_micros() < 2000,
            "Swap latency {}μs exceeds 2000μs",
            swap_latency.as_micros()
        );

        info!(latency_us = swap_latency.as_micros(), "Atomic swap latency");
    }

    #[tokio::test]
    async fn test_checkpoint_validation() {
        let validator = CheckpointValidator::new();
        let model = Arc::new(CheckpointModel::new(
            "DQN".to_owned(),
            "checkpoint_v1.safetensors".to_owned(),
            create_mock_prediction_fn(),
        ));

        let result = validator.validate(&model).await.unwrap();

        assert!(result.passed);
        assert!(result.p99_latency_us < 50);
        assert_eq!(result.predictions_validated, 1000);
        assert!(result.predictions_in_range >= 950); // At least 95%

        info!(
            avg_latency_us = result.avg_latency_us,
            p99_latency_us = result.p99_latency_us,
            predictions_in_range = result.predictions_in_range,
            predictions_validated = result.predictions_validated,
            "Validation result"
        );
    }

    #[tokio::test]
    async fn test_rollback_mechanism() {
        let model_v1 = Arc::new(CheckpointModel::new(
            "DQN".to_owned(),
            "checkpoint_v1.safetensors".to_owned(),
            create_mock_prediction_fn(),
        ));

        let buffer_pair = ModelBufferPair::new(model_v1);

        // Stage and commit new checkpoint
        let model_v2 = Arc::new(CheckpointModel::new(
            "DQN".to_owned(),
            "checkpoint_v2.safetensors".to_owned(),
            create_mock_prediction_fn(),
        ));

        buffer_pair.stage_checkpoint(model_v2).await.unwrap();
        buffer_pair.commit_swap().await.unwrap();

        // Verify v2 is active
        let active = buffer_pair.get_active().await;
        assert_eq!(active.checkpoint_path, "checkpoint_v2.safetensors");

        // Rollback
        buffer_pair.rollback().await.unwrap();

        // Verify v1 is restored
        let active = buffer_pair.get_active().await;
        assert_eq!(active.checkpoint_path, "checkpoint_v1.safetensors");
    }

    #[tokio::test]
    async fn test_hot_swap_manager() {
        let validator = CheckpointValidator::new();
        let policy = RollbackPolicy::default();
        let manager = HotSwapManager::new(validator, policy);

        // Register initial model
        let model_v1 = Arc::new(CheckpointModel::new(
            "DQN".to_owned(),
            "checkpoint_v1.safetensors".to_owned(),
            create_mock_prediction_fn(),
        ));

        manager
            .register_model("DQN".to_owned(), model_v1)
            .await
            .unwrap();

        // Stage new checkpoint
        let model_v2 = Arc::new(CheckpointModel::new(
            "DQN".to_owned(),
            "checkpoint_v2.safetensors".to_owned(),
            create_mock_prediction_fn(),
        ));

        manager.stage_checkpoint("DQN", model_v2).await.unwrap();

        // Validate staged checkpoint
        let validation = manager.validate_staged_checkpoint("DQN").await.unwrap();
        assert!(validation.passed);

        // Commit swap
        let swap_latency = manager.commit_swap("DQN").await.unwrap();
        assert!(swap_latency.as_micros() < 100);

        // Verify active checkpoint
        let active = manager.get_active_checkpoint("DQN").await.unwrap();
        assert_eq!(active.checkpoint_path, "checkpoint_v2.safetensors");

        info!("Hot-swap workflow completed successfully");
    }
}