rustytorch/crates/training/rtx-auto/src/rollback.rs

//! Rollback management system for safe optimization application and recovery.

use crate::{
    error::{AutoError, AutoResult},
    proposal::Proposal,
};
use rtx_runtime::Runtime;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tracing::{debug, info, warn};
use uuid::Uuid;

/// Types of checkpoints that can be created.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CheckpointType {
    /// Manual checkpoint created by user
    Manual,
    /// Automatic checkpoint before optimization
    BeforeOptimization,
    /// Automatic checkpoint during adaptive optimization
    Automatic,
    /// Checkpoint before applying experimental changes
    Experimental,
}

/// A checkpoint containing system state that can be restored.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Checkpoint {
    id: String,
    checkpoint_type: CheckpointType,
    description: String,
    timestamp: u64,
    state: HashMap<String, Vec<f32>>, // Simplified state representation
}

impl Checkpoint {
    /// Create a new checkpoint.
    pub fn new(
        checkpoint_type: CheckpointType,
        state: HashMap<String, Vec<f32>>,
        description: String,
    ) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            checkpoint_type,
            description,
            timestamp: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs(),
            state,
        }
    }

    /// Get checkpoint ID.
    pub fn id(&self) -> &str {
        &self.id
    }

    /// Get checkpoint type.
    pub fn checkpoint_type(&self) -> CheckpointType {
        self.checkpoint_type
    }

    /// Get checkpoint description.
    pub fn description(&self) -> &str {
        &self.description
    }

    /// Get checkpoint timestamp.
    pub fn timestamp(&self) -> u64 {
        self.timestamp
    }

    /// Get checkpoint state.
    pub fn state(&self) -> &HashMap<String, Vec<f32>> {
        &self.state
    }
}

/// Manages checkpoints and rollback operations for autonomous optimization.
pub struct RollbackManager {
    runtime: Arc<Runtime>,
    checkpoints: HashMap<String, Checkpoint>,
    current_state: HashMap<String, Vec<f32>>,
    rollback_thresholds: HashMap<String, f32>,
    auto_checkpoint_enabled: bool,
}

impl RollbackManager {
    /// Create a new rollback manager.
    pub fn new(runtime: Arc<Runtime>) -> AutoResult<Self> {
        Ok(Self {
            runtime,
            checkpoints: HashMap::new(),
            current_state: HashMap::new(),
            rollback_thresholds: HashMap::new(),
            auto_checkpoint_enabled: true,
        })
    }

    /// Create a new checkpoint with the given state.
    pub async fn create_checkpoint(
        &self,
        checkpoint_type: CheckpointType,
        state: HashMap<String, Vec<f32>>,
        description: String,
    ) -> AutoResult<Checkpoint> {
        info!("Creating checkpoint: {}", description);
        Ok(Checkpoint::new(checkpoint_type, state, description))
    }

    /// Save a checkpoint to storage.
    pub async fn save_checkpoint(&mut self, checkpoint: &Checkpoint) -> AutoResult<()> {
        info!("Saving checkpoint {}", checkpoint.id());
        self.checkpoints
            .insert(checkpoint.id().to_string(), checkpoint.clone());
        Ok(())
    }

    /// Load a checkpoint from storage.
    pub async fn load_checkpoint(&self, checkpoint_id: &str) -> AutoResult<Checkpoint> {
        self.checkpoints
            .get(checkpoint_id)
            .cloned()
            .ok_or_else(|| AutoError::checkpoint_not_found(checkpoint_id))
    }

    /// Rollback to a specific checkpoint.
    pub async fn rollback_to_checkpoint(&mut self, checkpoint_id: &str) -> AutoResult<()> {
        info!("Rolling back to checkpoint {}", checkpoint_id);

        let checkpoint = self.load_checkpoint(checkpoint_id).await?;
        self.current_state = checkpoint.state().clone();

        debug!(
            "Rollback successful, state restored to checkpoint {}",
            checkpoint_id
        );
        Ok(())
    }

    /// Update the current system state.
    pub async fn update_current_state(
        &mut self,
        new_state: HashMap<String, Vec<f32>>,
    ) -> AutoResult<()> {
        self.current_state = new_state;
        Ok(())
    }

    /// Get the current system state.
    pub async fn get_current_state(&self) -> AutoResult<HashMap<String, Vec<f32>>> {
        Ok(self.current_state.clone())
    }

    /// Set automatic rollback threshold for a metric.
    pub async fn set_rollback_threshold(&mut self, metric: &str, threshold: f32) -> AutoResult<()> {
        info!("Setting rollback threshold for {}: {}", metric, threshold);
        self.rollback_thresholds
            .insert(metric.to_string(), threshold);
        Ok(())
    }

    /// Check if automatic rollback should be triggered.
    pub async fn should_trigger_automatic_rollback(
        &self,
        current_metrics: &HashMap<String, Vec<f32>>,
    ) -> AutoResult<bool> {
        for (metric, threshold) in &self.rollback_thresholds {
            if let Some(values) = current_metrics.get(metric)
                && let Some(&current_value) = values.first()
                && current_value < *threshold
            {
                warn!(
                    "Metric {} ({}) below threshold ({}), triggering rollback",
                    metric, current_value, threshold
                );
                return Ok(true);
            }
        }
        Ok(false)
    }

    /// Perform automatic rollback to the most recent checkpoint.
    pub async fn automatic_rollback(&mut self) -> AutoResult<()> {
        info!("Performing automatic rollback");

        // Find most recent automatic checkpoint
        let mut latest_checkpoint_id = None;
        let mut latest_timestamp = 0u64;

        for checkpoint in self.checkpoints.values() {
            if checkpoint.checkpoint_type() == CheckpointType::Automatic
                && checkpoint.timestamp() > latest_timestamp
            {
                latest_timestamp = checkpoint.timestamp();
                latest_checkpoint_id = Some(checkpoint.id().to_string());
            }
        }

        if let Some(checkpoint_id) = latest_checkpoint_id {
            self.rollback_to_checkpoint(&checkpoint_id).await?;
            info!(
                "Automatic rollback completed to checkpoint {}",
                checkpoint_id
            );
        } else {
            return Err(AutoError::rollback_failed("No automatic checkpoint found"));
        }

        Ok(())
    }

    /// Clean up old checkpoints, keeping only the most recent N.
    pub async fn cleanup_old_checkpoints(&mut self, keep_count: usize) -> AutoResult<()> {
        info!("Cleaning up old checkpoints, keeping {}", keep_count);

        if self.checkpoints.len() <= keep_count {
            return Ok(());
        }

        // Sort checkpoints by timestamp
        let mut checkpoints: Vec<_> = self.checkpoints.values().cloned().collect();
        checkpoints.sort_by(|a, b| b.timestamp().cmp(&a.timestamp()));

        // Keep only the most recent ones
        let to_keep: Vec<_> = checkpoints.into_iter().take(keep_count).collect();

        self.checkpoints.clear();
        for checkpoint in to_keep {
            self.checkpoints
                .insert(checkpoint.id().to_string(), checkpoint);
        }

        Ok(())
    }

    /// List all available checkpoints.
    pub async fn list_checkpoints(&self) -> AutoResult<Vec<Checkpoint>> {
        let mut checkpoints: Vec<_> = self.checkpoints.values().cloned().collect();
        checkpoints.sort_by(|a, b| b.timestamp().cmp(&a.timestamp()));
        Ok(checkpoints)
    }

    /// Create a checkpoint specifically for a proposal application.
    pub async fn create_checkpoint_for_proposal(
        &self,
        proposal: &Proposal,
        current_state: HashMap<String, Vec<f32>>,
    ) -> AutoResult<Checkpoint> {
        let description = format!("Before applying proposal: {}", proposal.description());
        self.create_checkpoint(
            CheckpointType::BeforeOptimization,
            current_state,
            description,
        )
        .await
    }

    /// Validate the results of applying a proposal.
    pub async fn validate_proposal_result(
        &self,
        _proposal: &Proposal,
        new_state: &HashMap<String, Vec<f32>>,
    ) -> AutoResult<bool> {
        // Simple validation - in practice would be more sophisticated
        for (metric, values) in new_state {
            if values.is_empty() {
                warn!("Empty values for metric {}", metric);
                return Ok(false);
            }

            // Check for reasonable values
            for &value in values {
                if value.is_nan() || value.is_infinite() {
                    warn!("Invalid value {} for metric {}", value, metric);
                    return Ok(false);
                }
            }
        }

        Ok(true)
    }
}

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

    #[tokio::test]
    async fn test_checkpoint_creation() {
        let runtime = Runtime::new().unwrap();
        let manager = RollbackManager::new(Arc::new(runtime)).unwrap();

        let mut state = HashMap::new();
        state.insert("accuracy".to_string(), vec![0.95]);

        let checkpoint = manager
            .create_checkpoint(CheckpointType::Manual, state, "Test checkpoint".to_string())
            .await
            .unwrap();

        assert_eq!(checkpoint.checkpoint_type(), CheckpointType::Manual);
        assert!(!checkpoint.id().is_empty());
    }
}