pilgrimage/broker/

data_replication.rs

//! Data Replication System
//!
//! Implements distributed data replication with strong consistency guarantees,
//! automatic failover, and conflict resolution for production environments.

use crate::message::message::Message;
/// Node identifier type
pub type NodeId = String;

/// Node state enumeration
#[derive(Debug, Clone, PartialEq)]
pub enum NodeState {
    Active,
    Inactive,
    Syncing,
    Failed,
}

use crate::network::error::{NetworkError, NetworkResult};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use tokio::sync::{mpsc, RwLock, Mutex};
use tokio::time::{interval, timeout};

/// Replication configuration
#[derive(Debug, Clone)]
pub struct ReplicationConfig {
    /// Replication factor (number of replicas)
    pub replication_factor: usize,
    /// Minimum in-sync replicas for writes
    pub min_in_sync_replicas: usize,
    /// Acknowledgment timeout for replication
    pub ack_timeout: Duration,
    /// Maximum log lag for in-sync replicas
    pub max_lag_ms: u64,
    /// Batch size for replication
    pub batch_size: usize,
    /// Replication retry attempts
    pub max_retries: usize,
}

impl Default for ReplicationConfig {
    fn default() -> Self {
        Self {
            replication_factor: 3,
            min_in_sync_replicas: 2,
            ack_timeout: Duration::from_secs(5),
            max_lag_ms: 1000,
            batch_size: 100,
            max_retries: 3,
        }
    }
}

/// Replication log entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReplicationLogEntry {
    /// Log sequence number
    pub lsn: u64,
    /// Topic name
    pub topic: String,
    /// Partition ID
    pub partition_id: u32,
    /// Message data
    pub message: Message,
    /// Timestamp
    pub timestamp: u64,
    /// Checksum for integrity
    pub checksum: u64,
}

/// Replica state
#[derive(Debug, Clone)]
pub struct ReplicaState {
    /// Node ID of the replica
    pub node_id: NodeId,
    /// Last synced LSN
    pub last_synced_lsn: u64,
    /// Last heartbeat timestamp
    pub last_heartbeat: SystemTime,
    /// In-sync status
    pub in_sync: bool,
    /// Lag in milliseconds
    pub lag_ms: u64,
}

/// Replication manager for distributed data synchronization
pub struct ReplicationManager {
    /// Configuration
    config: ReplicationConfig,
    /// Current node ID
    node_id: NodeId,
    /// Replication log
    replication_log: Arc<RwLock<VecDeque<ReplicationLogEntry>>>,
    /// Current LSN (Log Sequence Number)
    current_lsn: Arc<Mutex<u64>>,
    /// Replica states
    replica_states: Arc<RwLock<HashMap<NodeId, ReplicaState>>>,
    /// In-sync replica set
    in_sync_replicas: Arc<RwLock<Vec<NodeId>>>,
    /// Pending acknowledgments
    pending_acks: Arc<RwLock<HashMap<u64, PendingAck>>>,
    /// Replication channels
    replication_tx: mpsc::UnboundedSender<ReplicationMessage>,
    /// Statistics
    stats: Arc<RwLock<ReplicationStats>>,
}

/// Pending acknowledgment tracking
#[derive(Debug)]
#[allow(dead_code)]
struct PendingAck {
    lsn: u64,
    required_acks: usize,
    received_acks: Vec<NodeId>,
    created_at: SystemTime,
    sender: Option<tokio::sync::oneshot::Sender<NetworkResult<()>>>,
}

/// Replication message types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ReplicationMessage {
    /// Replicate log entries to followers
    ReplicateEntries {
        entries: Vec<ReplicationLogEntry>,
        leader_commit: u64,
    },
    /// Acknowledgment from follower
    ReplicationAck {
        node_id: NodeId,
        last_synced_lsn: u64,
        success: bool,
    },
    /// Heartbeat for replica liveness
    Heartbeat {
        node_id: NodeId,
        last_lsn: u64,
    },
    /// Request for log compaction
    CompactionRequest {
        before_lsn: u64,
    },
}

/// Replication statistics
#[derive(Debug, Default)]
pub struct ReplicationStats {
    pub total_replicated_entries: u64,
    pub failed_replications: u64,
    pub average_replication_lag_ms: f64,
    pub in_sync_replica_count: usize,
    pub total_acks_received: u64,
    pub replication_errors: u64,
}

impl ReplicationManager {
    /// Create new replication manager
    pub async fn new(config: ReplicationConfig, node_id: NodeId) -> NetworkResult<Self> {
        let (replication_tx, replication_rx) = mpsc::unbounded_channel();

        let manager = Self {
            config: config.clone(),
            node_id: node_id.clone(),
            replication_log: Arc::new(RwLock::new(VecDeque::new())),
            current_lsn: Arc::new(Mutex::new(0)),
            replica_states: Arc::new(RwLock::new(HashMap::new())),
            in_sync_replicas: Arc::new(RwLock::new(Vec::new())),
            pending_acks: Arc::new(RwLock::new(HashMap::new())),
            replication_tx: replication_tx.clone(),
            stats: Arc::new(RwLock::new(ReplicationStats::default())),
        };

        // Start replication message handler
        let manager_clone = manager.clone();
        tokio::spawn(async move {
            manager_clone.handle_replication_messages(replication_rx).await;
        });

        // Start heartbeat and cleanup tasks
        manager.start_background_tasks().await;

        println!("✅ Replication manager initialized for node: {}", node_id);
        Ok(manager)
    }

    /// Replicate message to all replicas
    pub async fn replicate_message(
        &self,
        _topic: &str,
        _partition_id: u32,
        message: Message,
    ) -> NetworkResult<u64> {
        // Generate LSN
        let lsn = self.next_lsn().await;

        // Create replication log entry
            let entry = ReplicationLogEntry {
                lsn,
                topic: message.topic_id.clone(),
                partition_id: message.partition_id as u32,
                message: message.clone(),
                timestamp: std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap()
                    .as_secs(),
                checksum: self.calculate_checksum(&message.content.as_bytes()),
            };        // Add to local log
        self.replication_log.write().await.push_back(entry.clone());

        // Get current in-sync replicas
        let in_sync_replicas = self.in_sync_replicas.read().await.clone();

        if in_sync_replicas.len() < self.config.min_in_sync_replicas {
            return Err(NetworkError::ReplicationError(
                "Insufficient in-sync replicas".to_string()
            ));
        }

        // Create pending acknowledgment
        let (ack_tx, ack_rx) = tokio::sync::oneshot::channel();
        {
            let mut pending_acks = self.pending_acks.write().await;
            pending_acks.insert(lsn, PendingAck {
                lsn,
                required_acks: self.config.min_in_sync_replicas,
                received_acks: Vec::new(),
                created_at: SystemTime::now(),
                sender: Some(ack_tx),
            });
        }

        // Send replication message to all replicas
        let replication_msg = ReplicationMessage::ReplicateEntries {
            entries: vec![entry],
            leader_commit: lsn,
        };

        for replica_id in &in_sync_replicas {
            if *replica_id != self.node_id {
                // In a real implementation, this would send over network
                self.replication_tx.send(replication_msg.clone())
                    .map_err(|e| NetworkError::ReplicationError(e.to_string()))?;
            }
        }

        // Wait for acknowledgments with timeout
        match timeout(self.config.ack_timeout, ack_rx).await {
            Ok(Ok(Ok(()))) => {
                self.stats.write().await.total_replicated_entries += 1;
                Ok(lsn)
            }
            Ok(Ok(Err(e))) => Err(NetworkError::ReplicationError(
                format!("Replication acknowledgment error: {}", e)
            )),
            Ok(Err(_)) => Err(NetworkError::ReplicationError(
                "Replication acknowledgment channel closed".to_string()
            )),
            Err(_) => {
                // Remove pending ack on timeout
                self.pending_acks.write().await.remove(&lsn);
                self.stats.write().await.failed_replications += 1;
                Err(NetworkError::ReplicationError(
                    "Replication acknowledgment timeout".to_string()
                ))
            }
        }
    }

    /// Add replica to the cluster
    pub async fn add_replica(&self, node_id: NodeId) -> NetworkResult<()> {
        let replica_state = ReplicaState {
            node_id: node_id.clone(),
            last_synced_lsn: 0,
            last_heartbeat: SystemTime::now(),
            in_sync: false,
            lag_ms: 0,
        };

        self.replica_states.write().await.insert(node_id.clone(), replica_state);

        println!("📥 Added replica: {}", node_id);
        Ok(())
    }

    /// Remove replica from the cluster
    pub async fn remove_replica(&self, node_id: &NodeId) -> NetworkResult<()> {
        self.replica_states.write().await.remove(node_id);

        // Remove from in-sync replicas
        let mut in_sync_replicas = self.in_sync_replicas.write().await;
        in_sync_replicas.retain(|id| id != node_id);

        println!("📤 Removed replica: {}", node_id);
        Ok(())
    }

    /// Handle replication acknowledgment
    pub async fn handle_replication_ack(
        &self,
        node_id: NodeId,
        lsn: u64,
        success: bool,
    ) -> NetworkResult<()> {
        // Update replica state
        if let Some(replica_state) = self.replica_states.write().await.get_mut(&node_id) {
            if success {
                replica_state.last_synced_lsn = lsn;
                replica_state.last_heartbeat = SystemTime::now();
                replica_state.lag_ms = 0; // Would calculate actual lag in real implementation
            }
        }

        // Check pending acknowledgments
        let mut pending_acks = self.pending_acks.write().await;
        if let Some(pending_ack) = pending_acks.get_mut(&lsn) {
            if success && !pending_ack.received_acks.contains(&node_id) {
                pending_ack.received_acks.push(node_id.clone());

                // Check if we have enough acknowledgments
                if pending_ack.received_acks.len() >= pending_ack.required_acks {
                    if let Some(sender) = pending_ack.sender.take() {
                        let _ = sender.send(Ok(()));
                    }
                    pending_acks.remove(&lsn);
                    self.stats.write().await.total_acks_received += 1;
                }
            }
        }

        Ok(())
    }

    /// Get replication statistics
    pub async fn get_statistics(&self) -> ReplicationStats {
        let stats = self.stats.read().await;
        let mut result = ReplicationStats {
            total_replicated_entries: stats.total_replicated_entries,
            failed_replications: stats.failed_replications,
            average_replication_lag_ms: stats.average_replication_lag_ms,
            in_sync_replica_count: 0, // Will be set below
            total_acks_received: stats.total_acks_received,
            replication_errors: stats.replication_errors,
        };
        result.in_sync_replica_count = self.in_sync_replicas.read().await.len();

        // Calculate average lag
        let replica_states = self.replica_states.read().await;
        if !replica_states.is_empty() {
            let total_lag: u64 = replica_states.values().map(|s| s.lag_ms).sum();
            result.average_replication_lag_ms = total_lag as f64 / replica_states.len() as f64;
        }

        result
    }

    /// Get next LSN
    async fn next_lsn(&self) -> u64 {
        let mut current_lsn = self.current_lsn.lock().await;
        *current_lsn += 1;
        *current_lsn
    }

    /// Calculate byte checksum
    fn calculate_checksum(&self, data: &[u8]) -> u64 {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();
        data.hash(&mut hasher);
        hasher.finish()
    }

    /// Calculate message checksum
    #[allow(dead_code)]
    fn calculate_message_checksum(&self, message: &Message) -> u64 {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();
        message.id.hash(&mut hasher);
        message.content.hash(&mut hasher);
        hasher.finish()
    }

    /// Handle replication messages
    async fn handle_replication_messages(&self, mut rx: mpsc::UnboundedReceiver<ReplicationMessage>) {
        while let Some(message) = rx.recv().await {
            match message {
                ReplicationMessage::ReplicateEntries { entries, leader_commit: _ } => {
                    // Handle replication from leader
                    for entry in entries {
                        // Add to local log
                        self.replication_log.write().await.push_back(entry.clone());

                        // Send acknowledgment
                        let ack_msg = ReplicationMessage::ReplicationAck {
                            node_id: self.node_id.clone(),
                            last_synced_lsn: entry.lsn,
                            success: true,
                        };

                        let _ = self.replication_tx.send(ack_msg);
                    }
                }
                ReplicationMessage::ReplicationAck { node_id, last_synced_lsn, success } => {
                    let _ = self.handle_replication_ack(node_id, last_synced_lsn, success).await;
                }
                ReplicationMessage::Heartbeat { node_id, last_lsn } => {
                    // Update replica heartbeat
                    if let Some(replica_state) = self.replica_states.write().await.get_mut(&node_id) {
                        replica_state.last_heartbeat = SystemTime::now();
                        replica_state.last_synced_lsn = last_lsn;
                    }
                }
                ReplicationMessage::CompactionRequest { before_lsn } => {
                    // Handle log compaction
                    self.compact_log(before_lsn).await;
                }
            }
        }
    }

    /// Start background tasks
    async fn start_background_tasks(&self) {
        let manager = self.clone();

        // Heartbeat task
        tokio::spawn(async move {
            let mut interval = interval(Duration::from_secs(1));
            loop {
                interval.tick().await;
                manager.send_heartbeats().await;
                manager.update_in_sync_replicas().await;
                manager.cleanup_expired_acks().await;
            }
        });
    }

    /// Send heartbeats to all replicas
    async fn send_heartbeats(&self) {
        let current_lsn = *self.current_lsn.lock().await;
        let heartbeat_msg = ReplicationMessage::Heartbeat {
            node_id: self.node_id.clone(),
            last_lsn: current_lsn,
        };

        let _ = self.replication_tx.send(heartbeat_msg);
    }

    /// Update in-sync replica set
    async fn update_in_sync_replicas(&self) {
        let now = SystemTime::now();
        let mut new_in_sync_replicas = Vec::new();

        let replica_states = self.replica_states.read().await;
        for (node_id, state) in replica_states.iter() {
            // Check if replica is in sync based on lag and heartbeat
            let heartbeat_age = now.duration_since(state.last_heartbeat)
                .unwrap_or_default()
                .as_millis() as u64;

            if heartbeat_age < self.config.max_lag_ms && state.lag_ms < self.config.max_lag_ms {
                new_in_sync_replicas.push(node_id.clone());
            }
        }

        *self.in_sync_replicas.write().await = new_in_sync_replicas;
    }

    /// Clean up expired acknowledgments
    async fn cleanup_expired_acks(&self) {
        let now = SystemTime::now();
        let mut pending_acks = self.pending_acks.write().await;

        pending_acks.retain(|_, ack| {
            let age = now.duration_since(ack.created_at).unwrap_or_default();
            age < self.config.ack_timeout * 2
        });
    }

    /// Compact replication log
    async fn compact_log(&self, before_lsn: u64) {
        let mut log = self.replication_log.write().await;
        while let Some(entry) = log.front() {
            if entry.lsn < before_lsn {
                log.pop_front();
            } else {
                break;
            }
        }

        println!("🗜️ Compacted replication log up to LSN: {}", before_lsn);
    }
}

impl Clone for ReplicationManager {
    fn clone(&self) -> Self {
        Self {
            config: self.config.clone(),
            node_id: self.node_id.clone(),
            replication_log: Arc::clone(&self.replication_log),
            current_lsn: Arc::clone(&self.current_lsn),
            replica_states: Arc::clone(&self.replica_states),
            in_sync_replicas: Arc::clone(&self.in_sync_replicas),
            pending_acks: Arc::clone(&self.pending_acks),
            replication_tx: self.replication_tx.clone(),
            stats: Arc::clone(&self.stats),
        }
    }
}

// Add ReplicationError to NetworkError enum
impl NetworkError {
    pub fn replication_error(msg: String) -> Self {
        NetworkError::ReplicationError(msg)
    }
}

// NetworkError Display implementation is handled by thiserror derive macro

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

    #[tokio::test]
    async fn test_replication_manager_creation() {
        let config = ReplicationConfig::default();
        let node_id = "test-node-1".to_string();

        let manager = ReplicationManager::new(config, node_id).await;
        assert!(manager.is_ok());
    }

    #[tokio::test]
    async fn test_replica_management() {
        let config = ReplicationConfig::default();
        let manager = ReplicationManager::new(config, "leader".to_string()).await.unwrap();

        // Add replica
        let result = manager.add_replica("replica-1".to_string()).await;
        assert!(result.is_ok());

        // Remove replica
        let result = manager.remove_replica(&"replica-1".to_string()).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_statistics() {
        let config = ReplicationConfig::default();
        let manager = ReplicationManager::new(config, "test-node".to_string()).await.unwrap();

        let stats = manager.get_statistics().await;
        assert_eq!(stats.total_replicated_entries, 0);
        assert_eq!(stats.in_sync_replica_count, 0);
    }
}