NAC_Blockchain/nac-upgrade-framework/src/snapshot.rs

//! 快照和回滚模块

use crate::version::Version;
use serde::{Deserialize, Serialize};
use sha3::{Digest, Sha3_384};
use std::fmt;

/// 快照ID
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct SnapshotId(pub String);

impl SnapshotId {
    pub fn new(id: String) -> Self {
        Self(id)
    }

    /// 从状态数据生成快照ID（使用SHA3-384哈希）
    pub fn generate(module_name: &str, version: &Version, state_data: &[u8]) -> Self {
        let mut hasher = Sha3_384::new();
        hasher.update(module_name.as_bytes());
        hasher.update(version.to_string().as_bytes());
        hasher.update(state_data);
        let hash = hasher.finalize();
        Self(hex::encode(hash))
    }
}

impl fmt::Display for SnapshotId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// 状态快照
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Snapshot {
    /// 快照ID
    pub snapshot_id: SnapshotId,
    /// 模块名称
    pub module_name: String,
    /// 版本
    pub version: Version,
    /// 状态数据（序列化后的二进制）
    pub state_data: Vec<u8>,
    /// 创建时间
    pub created_at: chrono::DateTime<chrono::Utc>,
    /// 元数据（可选）
    pub metadata: std::collections::HashMap<String, String>,
}

impl Snapshot {
    pub fn new(module_name: String, version: Version, state_data: Vec<u8>) -> Self {
        let snapshot_id = SnapshotId::generate(&module_name, &version, &state_data);
        Self {
            snapshot_id,
            module_name,
            version,
            state_data,
            created_at: chrono::Utc::now(),
            metadata: std::collections::HashMap::new(),
        }
    }

    /// 添加元数据
    pub fn add_metadata(&mut self, key: String, value: String) {
        self.metadata.insert(key, value);
    }

    /// 获取元数据
    pub fn get_metadata(&self, key: &str) -> Option<&String> {
        self.metadata.get(key)
    }

    /// 验证快照完整性
    pub fn verify_integrity(&self) -> bool {
        let expected_id = SnapshotId::generate(&self.module_name, &self.version, &self.state_data);
        self.snapshot_id == expected_id
    }

    /// 获取快照大小（字节）
    pub fn size(&self) -> usize {
        self.state_data.len()
    }
}

/// 快照管理器
#[derive(Debug, Default)]
pub struct SnapshotManager {
    snapshots: std::collections::HashMap<SnapshotId, Snapshot>,
}

impl SnapshotManager {
    pub fn new() -> Self {
        Self {
            snapshots: std::collections::HashMap::new(),
        }
    }

    /// 保存快照
    pub fn save(&mut self, snapshot: Snapshot) -> SnapshotId {
        let id = snapshot.snapshot_id.clone();
        self.snapshots.insert(id.clone(), snapshot);
        id
    }

    /// 获取快照
    pub fn get(&self, id: &SnapshotId) -> Option<&Snapshot> {
        self.snapshots.get(id)
    }

    /// 删除快照
    pub fn delete(&mut self, id: &SnapshotId) -> Option<Snapshot> {
        self.snapshots.remove(id)
    }

    /// 列出所有快照
    pub fn list(&self) -> Vec<&Snapshot> {
        self.snapshots.values().collect()
    }

    /// 列出指定模块的快照
    pub fn list_by_module(&self, module_name: &str) -> Vec<&Snapshot> {
        self.snapshots
            .values()
            .filter(|s| s.module_name == module_name)
            .collect()
    }

    /// 获取最新快照
    pub fn get_latest(&self, module_name: &str) -> Option<&Snapshot> {
        self.list_by_module(module_name)
            .into_iter()
            .max_by_key(|s| s.created_at)
    }

    /// 清理旧快照（保留最近N个）
    pub fn cleanup_old_snapshots(&mut self, module_name: &str, keep_count: usize) {
        let mut snapshots: Vec<_> = self.list_by_module(module_name)
            .into_iter()
            .map(|s| s.snapshot_id.clone())
            .collect();
        
        // 按创建时间排序
        snapshots.sort_by_key(|id| {
            self.snapshots.get(id).map(|s| s.created_at).unwrap_or_default()
        });

        // 删除旧快照
        if snapshots.len() > keep_count {
            for id in snapshots.iter().take(snapshots.len() - keep_count) {
                self.delete(id);
            }
        }
    }
}

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

    #[test]
    fn test_snapshot_id_generation() {
        let id1 = SnapshotId::generate("nac-nvm", &Version::new(1, 0, 0), b"state1");
        let id2 = SnapshotId::generate("nac-nvm", &Version::new(1, 0, 0), b"state1");
        let id3 = SnapshotId::generate("nac-nvm", &Version::new(1, 0, 0), b"state2");

        assert_eq!(id1, id2);
        assert_ne!(id1, id3);
    }

    #[test]
    fn test_snapshot_creation() {
        let snapshot = Snapshot::new(
            "nac-nvm".to_string(),
            Version::new(1, 0, 0),
            vec![1, 2, 3, 4],
        );

        assert_eq!(snapshot.module_name, "nac-nvm");
        assert_eq!(snapshot.version, Version::new(1, 0, 0));
        assert_eq!(snapshot.state_data, vec![1, 2, 3, 4]);
    }

    #[test]
    fn test_snapshot_integrity() {
        let snapshot = Snapshot::new(
            "nac-nvm".to_string(),
            Version::new(1, 0, 0),
            vec![1, 2, 3, 4],
        );

        assert!(snapshot.verify_integrity());
    }

    #[test]
    fn test_snapshot_metadata() {
        let mut snapshot = Snapshot::new(
            "nac-nvm".to_string(),
            Version::new(1, 0, 0),
            vec![1, 2, 3, 4],
        );

        snapshot.add_metadata("description".to_string(), "test snapshot".to_string());
        assert_eq!(snapshot.get_metadata("description"), Some(&"test snapshot".to_string()));
    }

    #[test]
    fn test_snapshot_manager() {
        let mut manager = SnapshotManager::new();

        let snapshot1 = Snapshot::new(
            "nac-nvm".to_string(),
            Version::new(1, 0, 0),
            vec![1, 2, 3],
        );
        let id1 = manager.save(snapshot1.clone());

        let snapshot2 = Snapshot::new(
            "nac-cbpp".to_string(),
            Version::new(1, 0, 0),
            vec![4, 5, 6],
        );
        let id2 = manager.save(snapshot2.clone());

        assert_eq!(manager.list().len(), 2);
        assert_eq!(manager.get(&id1).unwrap().module_name, "nac-nvm");
        assert_eq!(manager.get(&id2).unwrap().module_name, "nac-cbpp");
    }

    #[test]
    fn test_snapshot_manager_list_by_module() {
        let mut manager = SnapshotManager::new();

        let snapshot1 = Snapshot::new(
            "nac-nvm".to_string(),
            Version::new(1, 0, 0),
            vec![1, 2, 3],
        );
        manager.save(snapshot1);

        let snapshot2 = Snapshot::new(
            "nac-nvm".to_string(),
            Version::new(1, 1, 0),
            vec![4, 5, 6],
        );
        manager.save(snapshot2);

        let snapshot3 = Snapshot::new(
            "nac-cbpp".to_string(),
            Version::new(1, 0, 0),
            vec![7, 8, 9],
        );
        manager.save(snapshot3);

        let nvm_snapshots = manager.list_by_module("nac-nvm");
        assert_eq!(nvm_snapshots.len(), 2);
    }

    #[test]
    fn test_snapshot_manager_get_latest() {
        let mut manager = SnapshotManager::new();

        let snapshot1 = Snapshot::new(
            "nac-nvm".to_string(),
            Version::new(1, 0, 0),
            vec![1, 2, 3],
        );
        manager.save(snapshot1);

        std::thread::sleep(std::time::Duration::from_millis(10));

        let snapshot2 = Snapshot::new(
            "nac-nvm".to_string(),
            Version::new(1, 1, 0),
            vec![4, 5, 6],
        );
        manager.save(snapshot2.clone());

        let latest = manager.get_latest("nac-nvm").unwrap();
        assert_eq!(latest.version, Version::new(1, 1, 0));
    }

    #[test]
    fn test_snapshot_manager_delete() {
        let mut manager = SnapshotManager::new();

        let snapshot = Snapshot::new(
            "nac-nvm".to_string(),
            Version::new(1, 0, 0),
            vec![1, 2, 3],
        );
        let id = manager.save(snapshot);

        assert!(manager.get(&id).is_some());
        manager.delete(&id);
        assert!(manager.get(&id).is_none());
    }
}