mecha10_core/

state.rs

//! Persistent State Storage Abstraction for Nodes
//!
//! This module provides a unified interface for nodes to persist and restore
//! state across restarts, crashes, or deployments.
//!
//! # Features
//!
//! - Multiple storage backends (in-memory, filesystem, Redis, PostgreSQL)
//! - Type-safe state serialization/deserialization
//! - Versioned state with migration support
//! - Atomic updates with compare-and-swap
//! - State snapshots and restoration
//! - TTL support for temporary state
//! - Namespace isolation per node
//!
//! # Usage
//!
//! ```rust
//! use mecha10::prelude::*;
//!
//! #[derive(Debug, Clone, Serialize, Deserialize)]
//! struct MyNodeState {
//!     counter: u64,
//!     last_position: (f64, f64),
//!     calibration_data: Vec<f32>,
//! }
//!
//! async fn example(ctx: &Context) -> Result<()> {
//!     let state_mgr = ctx.state_manager();
//!
//!     // Save state
//!     let state = MyNodeState {
//!         counter: 42,
//!         last_position: (1.5, 2.3),
//!         calibration_data: vec![0.1, 0.2, 0.3],
//!     };
//!     state_mgr.save("my_node", &state).await?;
//!
//!     // Load state
//!     if let Some(restored) = state_mgr.load::<MyNodeState>("my_node").await? {
//!         println!("Restored state: {:?}", restored);
//!     }
//!
//!     Ok(())
//! }
//! ```

use crate::error::{Mecha10Error, Result};
use async_trait::async_trait;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use tokio::fs;
use tokio::sync::RwLock;

// ============================================================================
// Core Types
// ============================================================================

/// Versioned state wrapper for migration support
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VersionedState<T> {
    /// Schema version for migration
    pub version: u32,

    /// Actual state data
    pub data: T,

    /// Timestamp when state was saved
    pub timestamp: u64,

    /// Optional TTL in seconds
    pub ttl_seconds: Option<u64>,
}

impl<T> VersionedState<T> {
    /// Create new versioned state
    pub fn new(data: T, version: u32) -> Self {
        Self {
            version,
            data,
            timestamp: now_micros(),
            ttl_seconds: None,
        }
    }

    /// Create new versioned state with TTL
    pub fn with_ttl(data: T, version: u32, ttl: Duration) -> Self {
        Self {
            version,
            data,
            timestamp: now_micros(),
            ttl_seconds: Some(ttl.as_secs()),
        }
    }

    /// Check if state has expired based on TTL
    pub fn is_expired(&self) -> bool {
        if let Some(ttl_seconds) = self.ttl_seconds {
            let age_seconds = (now_micros() - self.timestamp) / 1_000_000;
            age_seconds > ttl_seconds
        } else {
            false
        }
    }

    /// Get age of state in seconds
    pub fn age_seconds(&self) -> u64 {
        (now_micros() - self.timestamp) / 1_000_000
    }
}

/// State snapshot metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateSnapshot {
    /// Snapshot identifier
    pub id: String,

    /// Node ID this snapshot belongs to
    pub node_id: String,

    /// Timestamp when snapshot was taken
    pub timestamp: u64,

    /// Optional description
    pub description: Option<String>,

    /// Size in bytes
    pub size_bytes: usize,
}

/// Options for state operations
#[derive(Debug, Clone)]
pub struct StateOptions {
    /// Version for schema migration
    pub version: u32,

    /// Optional TTL
    pub ttl: Option<Duration>,

    /// Whether to create backup before overwrite
    pub create_backup: bool,

    /// Whether to compress state
    pub compress: bool,
}

impl Default for StateOptions {
    fn default() -> Self {
        Self {
            version: 1,
            ttl: None,
            create_backup: false,
            compress: false,
        }
    }
}

// ============================================================================
// StateManager Trait
// ============================================================================

/// Persistent state storage abstraction
#[async_trait]
pub trait StateManager: Send + Sync {
    /// Save state for a node
    async fn save<T: Serialize + Send + Sync>(&self, node_id: &str, state: &T) -> Result<()>;

    /// Save state with options
    async fn save_with_options<T: Serialize + Send + Sync>(
        &self,
        node_id: &str,
        state: &T,
        options: StateOptions,
    ) -> Result<()>;

    /// Load state for a node
    async fn load<T: DeserializeOwned + Send + Sync>(&self, node_id: &str) -> Result<Option<T>>;

    /// Load versioned state with migration support
    async fn load_versioned<T: DeserializeOwned + Send + Sync>(
        &self,
        node_id: &str,
    ) -> Result<Option<VersionedState<T>>>;

    /// Delete state for a node
    async fn delete(&self, node_id: &str) -> Result<()>;

    /// Check if state exists for a node
    async fn exists(&self, node_id: &str) -> Result<bool>;

    /// List all node IDs with stored state
    async fn list_nodes(&self) -> Result<Vec<String>>;

    /// Get state size in bytes
    async fn size(&self, node_id: &str) -> Result<Option<usize>>;

    /// Create a snapshot of current state
    async fn create_snapshot(&self, node_id: &str, description: Option<String>) -> Result<StateSnapshot>;

    /// Restore from a snapshot
    async fn restore_snapshot(&self, snapshot_id: &str) -> Result<()>;

    /// List snapshots for a node
    async fn list_snapshots(&self, node_id: &str) -> Result<Vec<StateSnapshot>>;

    /// Delete a snapshot
    async fn delete_snapshot(&self, snapshot_id: &str) -> Result<()>;

    /// Atomic compare-and-swap operation
    async fn compare_and_swap<T: Serialize + DeserializeOwned + Send + Sync>(
        &self,
        node_id: &str,
        expected: &T,
        new: &T,
    ) -> Result<bool>;

    /// Clear all expired states (based on TTL)
    async fn clear_expired(&self) -> Result<usize>;
}

// ============================================================================
// In-Memory Backend
// ============================================================================

/// In-memory state storage (not persistent across restarts)
#[derive(Clone)]
pub struct MemoryStateManager {
    states: Arc<RwLock<HashMap<String, Vec<u8>>>>,
    #[allow(clippy::type_complexity)]
    snapshots: Arc<RwLock<HashMap<String, (Vec<u8>, StateSnapshot)>>>,
}

impl MemoryStateManager {
    /// Create new in-memory state manager
    pub fn new() -> Self {
        Self {
            states: Arc::new(RwLock::new(HashMap::new())),
            snapshots: Arc::new(RwLock::new(HashMap::new())),
        }
    }
}

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

#[async_trait]
impl StateManager for MemoryStateManager {
    async fn save<T: Serialize + Send + Sync>(&self, node_id: &str, state: &T) -> Result<()> {
        self.save_with_options(node_id, state, StateOptions::default()).await
    }

    async fn save_with_options<T: Serialize + Send + Sync>(
        &self,
        node_id: &str,
        state: &T,
        options: StateOptions,
    ) -> Result<()> {
        let versioned = if options.ttl.is_some() {
            VersionedState::with_ttl(state, options.version, options.ttl.unwrap())
        } else {
            VersionedState::new(state, options.version)
        };

        let data = serde_json::to_vec(&versioned).map_err(|e| Mecha10Error::SerializationError {
            message: format!("Failed to serialize state: {}", e),
            suggestion: "Ensure state type implements Serialize".to_string(),
        })?;

        let mut states = self.states.write().await;
        states.insert(node_id.to_string(), data);

        Ok(())
    }

    async fn load<T: DeserializeOwned + Send + Sync>(&self, node_id: &str) -> Result<Option<T>> {
        if let Some(versioned) = self.load_versioned::<T>(node_id).await? {
            Ok(Some(versioned.data))
        } else {
            Ok(None)
        }
    }

    async fn load_versioned<T: DeserializeOwned + Send + Sync>(
        &self,
        node_id: &str,
    ) -> Result<Option<VersionedState<T>>> {
        let states = self.states.read().await;

        if let Some(data) = states.get(node_id) {
            let versioned: VersionedState<T> =
                serde_json::from_slice(data).map_err(|e| Mecha10Error::SerializationError {
                    message: format!("Failed to deserialize state: {}", e),
                    suggestion: "State may be corrupted or type mismatch".to_string(),
                })?;

            // Check TTL expiration
            if versioned.is_expired() {
                return Ok(None);
            }

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

    async fn delete(&self, node_id: &str) -> Result<()> {
        let mut states = self.states.write().await;
        states.remove(node_id);
        Ok(())
    }

    async fn exists(&self, node_id: &str) -> Result<bool> {
        let states = self.states.read().await;
        Ok(states.contains_key(node_id))
    }

    async fn list_nodes(&self) -> Result<Vec<String>> {
        let states = self.states.read().await;
        Ok(states.keys().cloned().collect())
    }

    async fn size(&self, node_id: &str) -> Result<Option<usize>> {
        let states = self.states.read().await;
        Ok(states.get(node_id).map(|data| data.len()))
    }

    async fn create_snapshot(&self, node_id: &str, description: Option<String>) -> Result<StateSnapshot> {
        let states = self.states.read().await;

        let data = states
            .get(node_id)
            .ok_or_else(|| Mecha10Error::Configuration(format!("No state for node: {}", node_id)))?
            .clone();

        let snapshot = StateSnapshot {
            id: uuid::Uuid::new_v4().to_string(),
            node_id: node_id.to_string(),
            timestamp: now_micros(),
            description,
            size_bytes: data.len(),
        };

        let mut snapshots = self.snapshots.write().await;
        snapshots.insert(snapshot.id.clone(), (data, snapshot.clone()));

        Ok(snapshot)
    }

    async fn restore_snapshot(&self, snapshot_id: &str) -> Result<()> {
        let snapshots = self.snapshots.read().await;

        let (data, snapshot) = snapshots
            .get(snapshot_id)
            .ok_or_else(|| Mecha10Error::Configuration(format!("Snapshot not found: {}", snapshot_id)))?;

        let mut states = self.states.write().await;
        states.insert(snapshot.node_id.clone(), data.clone());

        Ok(())
    }

    async fn list_snapshots(&self, node_id: &str) -> Result<Vec<StateSnapshot>> {
        let snapshots = self.snapshots.read().await;

        let mut result: Vec<StateSnapshot> = snapshots
            .values()
            .filter_map(|(_, snapshot)| {
                if snapshot.node_id == node_id {
                    Some(snapshot.clone())
                } else {
                    None
                }
            })
            .collect();

        result.sort_by_key(|s| s.timestamp);
        result.reverse();

        Ok(result)
    }

    async fn delete_snapshot(&self, snapshot_id: &str) -> Result<()> {
        let mut snapshots = self.snapshots.write().await;
        snapshots.remove(snapshot_id);
        Ok(())
    }

    async fn compare_and_swap<T: Serialize + DeserializeOwned + Send + Sync>(
        &self,
        node_id: &str,
        expected: &T,
        new: &T,
    ) -> Result<bool> {
        let mut states = self.states.write().await;

        if let Some(data) = states.get(node_id) {
            let current: VersionedState<T> =
                serde_json::from_slice(data).map_err(|e| Mecha10Error::SerializationError {
                    message: format!("Failed to deserialize for CAS: {}", e),
                    suggestion: "State may be corrupted".to_string(),
                })?;

            let expected_json = serde_json::to_string(expected).unwrap();
            let current_json = serde_json::to_string(&current.data).unwrap();

            if expected_json == current_json {
                // Match - perform swap
                let versioned = VersionedState::new(new, current.version + 1);
                let new_data = serde_json::to_vec(&versioned).unwrap();
                states.insert(node_id.to_string(), new_data);
                Ok(true)
            } else {
                // No match
                Ok(false)
            }
        } else {
            // No existing state
            Ok(false)
        }
    }

    async fn clear_expired(&self) -> Result<usize> {
        let mut states = self.states.write().await;
        let mut expired_keys = Vec::new();

        for (node_id, data) in states.iter() {
            if let Ok(versioned) = serde_json::from_slice::<VersionedState<serde_json::Value>>(data) {
                if versioned.is_expired() {
                    expired_keys.push(node_id.clone());
                }
            }
        }

        let count = expired_keys.len();
        for key in expired_keys {
            states.remove(&key);
        }

        Ok(count)
    }
}

// ============================================================================
// Filesystem Backend
// ============================================================================

/// Filesystem-based state storage
#[derive(Clone)]
pub struct FilesystemStateManager {
    base_path: PathBuf,
}

impl FilesystemStateManager {
    /// Create new filesystem state manager
    pub async fn new<P: AsRef<Path>>(base_path: P) -> Result<Self> {
        let base_path = base_path.as_ref().to_path_buf();

        // Create directory if it doesn't exist
        fs::create_dir_all(&base_path)
            .await
            .map_err(|e| Mecha10Error::IoError {
                message: format!("Failed to create state directory: {}", e),
                suggestion: "Check directory permissions".to_string(),
            })?;

        Ok(Self { base_path })
    }

    fn state_path(&self, node_id: &str) -> PathBuf {
        self.base_path.join(format!("{}.state.json", node_id))
    }

    fn snapshot_path(&self, snapshot_id: &str) -> PathBuf {
        self.base_path
            .join("snapshots")
            .join(format!("{}.snapshot.json", snapshot_id))
    }

    fn snapshot_meta_path(&self, snapshot_id: &str) -> PathBuf {
        self.base_path
            .join("snapshots")
            .join(format!("{}.meta.json", snapshot_id))
    }
}

#[async_trait]
impl StateManager for FilesystemStateManager {
    async fn save<T: Serialize + Send + Sync>(&self, node_id: &str, state: &T) -> Result<()> {
        self.save_with_options(node_id, state, StateOptions::default()).await
    }

    async fn save_with_options<T: Serialize + Send + Sync>(
        &self,
        node_id: &str,
        state: &T,
        options: StateOptions,
    ) -> Result<()> {
        let versioned = if options.ttl.is_some() {
            VersionedState::with_ttl(state, options.version, options.ttl.unwrap())
        } else {
            VersionedState::new(state, options.version)
        };

        let path = self.state_path(node_id);

        // Create backup if requested
        if options.create_backup && path.exists() {
            let backup_path = path.with_extension("state.json.backup");
            fs::copy(&path, backup_path).await.ok();
        }

        let json = serde_json::to_string_pretty(&versioned).map_err(|e| Mecha10Error::SerializationError {
            message: format!("Failed to serialize state: {}", e),
            suggestion: "Ensure state type implements Serialize".to_string(),
        })?;

        fs::write(&path, json).await.map_err(|e| Mecha10Error::IoError {
            message: format!("Failed to write state file: {}", e),
            suggestion: "Check file permissions".to_string(),
        })?;

        Ok(())
    }

    async fn load<T: DeserializeOwned + Send + Sync>(&self, node_id: &str) -> Result<Option<T>> {
        if let Some(versioned) = self.load_versioned::<T>(node_id).await? {
            Ok(Some(versioned.data))
        } else {
            Ok(None)
        }
    }

    async fn load_versioned<T: DeserializeOwned + Send + Sync>(
        &self,
        node_id: &str,
    ) -> Result<Option<VersionedState<T>>> {
        let path = self.state_path(node_id);

        if !path.exists() {
            return Ok(None);
        }

        let json = fs::read_to_string(&path).await.map_err(|e| Mecha10Error::IoError {
            message: format!("Failed to read state file: {}", e),
            suggestion: "Check file permissions".to_string(),
        })?;

        let versioned: VersionedState<T> =
            serde_json::from_str(&json).map_err(|e| Mecha10Error::SerializationError {
                message: format!("Failed to deserialize state: {}", e),
                suggestion: "State may be corrupted or type mismatch".to_string(),
            })?;

        // Check TTL expiration
        if versioned.is_expired() {
            return Ok(None);
        }

        Ok(Some(versioned))
    }

    async fn delete(&self, node_id: &str) -> Result<()> {
        let path = self.state_path(node_id);

        if path.exists() {
            fs::remove_file(&path).await.map_err(|e| Mecha10Error::IoError {
                message: format!("Failed to delete state file: {}", e),
                suggestion: "Check file permissions".to_string(),
            })?;
        }

        Ok(())
    }

    async fn exists(&self, node_id: &str) -> Result<bool> {
        Ok(self.state_path(node_id).exists())
    }

    async fn list_nodes(&self) -> Result<Vec<String>> {
        let mut nodes = Vec::new();
        let mut entries = fs::read_dir(&self.base_path).await.map_err(|e| Mecha10Error::IoError {
            message: format!("Failed to read state directory: {}", e),
            suggestion: "Check directory permissions".to_string(),
        })?;

        while let Some(entry) = entries.next_entry().await.map_err(|e| Mecha10Error::IoError {
            message: format!("Failed to read directory entry: {}", e),
            suggestion: "Check directory permissions".to_string(),
        })? {
            let path = entry.path();
            if let Some(filename) = path.file_name() {
                let filename = filename.to_string_lossy();
                if filename.ends_with(".state.json") {
                    let node_id = filename.trim_end_matches(".state.json").to_string();
                    nodes.push(node_id);
                }
            }
        }

        Ok(nodes)
    }

    async fn size(&self, node_id: &str) -> Result<Option<usize>> {
        let path = self.state_path(node_id);

        if !path.exists() {
            return Ok(None);
        }

        let metadata = fs::metadata(&path).await.map_err(|e| Mecha10Error::IoError {
            message: format!("Failed to get file metadata: {}", e),
            suggestion: "Check file permissions".to_string(),
        })?;

        Ok(Some(metadata.len() as usize))
    }

    async fn create_snapshot(&self, node_id: &str, description: Option<String>) -> Result<StateSnapshot> {
        let state_path = self.state_path(node_id);

        if !state_path.exists() {
            return Err(Mecha10Error::Configuration(format!("No state for node: {}", node_id)));
        }

        let snapshot_id = uuid::Uuid::new_v4().to_string();

        // Ensure snapshots directory exists
        let snapshots_dir = self.base_path.join("snapshots");
        fs::create_dir_all(&snapshots_dir)
            .await
            .map_err(|e| Mecha10Error::IoError {
                message: format!("Failed to create snapshots directory: {}", e),
                suggestion: "Check directory permissions".to_string(),
            })?;

        let snapshot_path = self.snapshot_path(&snapshot_id);
        fs::copy(&state_path, &snapshot_path)
            .await
            .map_err(|e| Mecha10Error::IoError {
                message: format!("Failed to create snapshot: {}", e),
                suggestion: "Check file permissions".to_string(),
            })?;

        let size_bytes = self.size(node_id).await?.unwrap_or(0);

        let snapshot = StateSnapshot {
            id: snapshot_id.clone(),
            node_id: node_id.to_string(),
            timestamp: now_micros(),
            description,
            size_bytes,
        };

        // Save snapshot metadata
        let meta_path = self.snapshot_meta_path(&snapshot_id);
        let meta_json = serde_json::to_string_pretty(&snapshot).unwrap();
        fs::write(&meta_path, meta_json)
            .await
            .map_err(|e| Mecha10Error::IoError {
                message: format!("Failed to write snapshot metadata: {}", e),
                suggestion: "Check file permissions".to_string(),
            })?;

        Ok(snapshot)
    }

    async fn restore_snapshot(&self, snapshot_id: &str) -> Result<()> {
        let snapshot_path = self.snapshot_path(snapshot_id);
        let meta_path = self.snapshot_meta_path(snapshot_id);

        if !snapshot_path.exists() || !meta_path.exists() {
            return Err(Mecha10Error::Configuration(format!(
                "Snapshot not found: {}",
                snapshot_id
            )));
        }

        let meta_json = fs::read_to_string(&meta_path)
            .await
            .map_err(|e| Mecha10Error::IoError {
                message: format!("Failed to read snapshot metadata: {}", e),
                suggestion: "Check file permissions".to_string(),
            })?;

        let snapshot: StateSnapshot =
            serde_json::from_str(&meta_json).map_err(|e| Mecha10Error::SerializationError {
                message: format!("Failed to deserialize snapshot metadata: {}", e),
                suggestion: "Snapshot metadata may be corrupted".to_string(),
            })?;

        let state_path = self.state_path(&snapshot.node_id);
        fs::copy(&snapshot_path, &state_path)
            .await
            .map_err(|e| Mecha10Error::IoError {
                message: format!("Failed to restore snapshot: {}", e),
                suggestion: "Check file permissions".to_string(),
            })?;

        Ok(())
    }

    async fn list_snapshots(&self, node_id: &str) -> Result<Vec<StateSnapshot>> {
        let snapshots_dir = self.base_path.join("snapshots");

        if !snapshots_dir.exists() {
            return Ok(Vec::new());
        }

        let mut snapshots = Vec::new();
        let mut entries = fs::read_dir(&snapshots_dir).await.map_err(|e| Mecha10Error::IoError {
            message: format!("Failed to read snapshots directory: {}", e),
            suggestion: "Check directory permissions".to_string(),
        })?;

        while let Some(entry) = entries.next_entry().await.map_err(|e| Mecha10Error::IoError {
            message: format!("Failed to read directory entry: {}", e),
            suggestion: "Check directory permissions".to_string(),
        })? {
            let path = entry.path();
            if let Some(filename) = path.file_name() {
                let filename = filename.to_string_lossy();
                if filename.ends_with(".meta.json") {
                    let json = fs::read_to_string(&path).await.ok();
                    if let Some(json) = json {
                        if let Ok(snapshot) = serde_json::from_str::<StateSnapshot>(&json) {
                            if snapshot.node_id == node_id {
                                snapshots.push(snapshot);
                            }
                        }
                    }
                }
            }
        }

        snapshots.sort_by_key(|s| s.timestamp);
        snapshots.reverse();

        Ok(snapshots)
    }

    async fn delete_snapshot(&self, snapshot_id: &str) -> Result<()> {
        let snapshot_path = self.snapshot_path(snapshot_id);
        let meta_path = self.snapshot_meta_path(snapshot_id);

        if snapshot_path.exists() {
            fs::remove_file(&snapshot_path).await.ok();
        }

        if meta_path.exists() {
            fs::remove_file(&meta_path).await.ok();
        }

        Ok(())
    }

    async fn compare_and_swap<T: Serialize + DeserializeOwned + Send + Sync>(
        &self,
        node_id: &str,
        expected: &T,
        new: &T,
    ) -> Result<bool> {
        let path = self.state_path(node_id);

        if !path.exists() {
            return Ok(false);
        }

        let json = fs::read_to_string(&path).await.map_err(|e| Mecha10Error::IoError {
            message: format!("Failed to read state file: {}", e),
            suggestion: "Check file permissions".to_string(),
        })?;

        let current: VersionedState<T> = serde_json::from_str(&json).map_err(|e| Mecha10Error::SerializationError {
            message: format!("Failed to deserialize for CAS: {}", e),
            suggestion: "State may be corrupted".to_string(),
        })?;

        let expected_json = serde_json::to_string(expected).unwrap();
        let current_json = serde_json::to_string(&current.data).unwrap();

        if expected_json == current_json {
            // Match - perform swap
            let versioned = VersionedState::new(new, current.version + 1);
            let new_json = serde_json::to_string_pretty(&versioned).unwrap();
            fs::write(&path, new_json).await.map_err(|e| Mecha10Error::IoError {
                message: format!("Failed to write state file: {}", e),
                suggestion: "Check file permissions".to_string(),
            })?;
            Ok(true)
        } else {
            // No match
            Ok(false)
        }
    }

    async fn clear_expired(&self) -> Result<usize> {
        let nodes = self.list_nodes().await?;
        let mut count = 0;

        for node_id in nodes {
            if let Ok(Some(versioned)) = self.load_versioned::<serde_json::Value>(&node_id).await {
                if versioned.is_expired() {
                    self.delete(&node_id).await?;
                    count += 1;
                }
            }
        }

        Ok(count)
    }
}

// ============================================================================
// Concrete StateManager enum for trait object compatibility
// ============================================================================

/// Concrete enum to wrap StateManager implementations
/// This allows using StateManager without dyn (which isn't object-safe due to generic methods)
#[derive(Clone)]
pub enum ConcreteStateManager {
    Memory(MemoryStateManager),
    Filesystem(FilesystemStateManager),
}

impl ConcreteStateManager {
    pub async fn save<T: Serialize + Send + Sync>(&self, node_id: &str, state: &T) -> Result<()> {
        match self {
            Self::Memory(m) => m.save(node_id, state).await,
            Self::Filesystem(f) => f.save(node_id, state).await,
        }
    }

    pub async fn save_with_options<T: Serialize + Send + Sync>(
        &self,
        node_id: &str,
        state: &T,
        options: StateOptions,
    ) -> Result<()> {
        match self {
            Self::Memory(m) => m.save_with_options(node_id, state, options).await,
            Self::Filesystem(f) => f.save_with_options(node_id, state, options).await,
        }
    }

    pub async fn load<T: DeserializeOwned + Send + Sync>(&self, node_id: &str) -> Result<Option<T>> {
        match self {
            Self::Memory(m) => m.load(node_id).await,
            Self::Filesystem(f) => f.load(node_id).await,
        }
    }

    pub async fn load_versioned<T: DeserializeOwned + Send + Sync>(
        &self,
        node_id: &str,
    ) -> Result<Option<VersionedState<T>>> {
        match self {
            Self::Memory(m) => m.load_versioned(node_id).await,
            Self::Filesystem(f) => f.load_versioned(node_id).await,
        }
    }

    pub async fn delete(&self, node_id: &str) -> Result<()> {
        match self {
            Self::Memory(m) => m.delete(node_id).await,
            Self::Filesystem(f) => f.delete(node_id).await,
        }
    }

    pub async fn exists(&self, node_id: &str) -> Result<bool> {
        match self {
            Self::Memory(m) => m.exists(node_id).await,
            Self::Filesystem(f) => f.exists(node_id).await,
        }
    }

    pub async fn list_nodes(&self) -> Result<Vec<String>> {
        match self {
            Self::Memory(m) => m.list_nodes().await,
            Self::Filesystem(f) => f.list_nodes().await,
        }
    }

    pub async fn size(&self, node_id: &str) -> Result<Option<usize>> {
        match self {
            Self::Memory(m) => m.size(node_id).await,
            Self::Filesystem(f) => f.size(node_id).await,
        }
    }

    pub async fn create_snapshot(&self, node_id: &str, description: Option<String>) -> Result<StateSnapshot> {
        match self {
            Self::Memory(m) => m.create_snapshot(node_id, description).await,
            Self::Filesystem(f) => f.create_snapshot(node_id, description).await,
        }
    }

    pub async fn restore_snapshot(&self, snapshot_id: &str) -> Result<()> {
        match self {
            Self::Memory(m) => m.restore_snapshot(snapshot_id).await,
            Self::Filesystem(f) => f.restore_snapshot(snapshot_id).await,
        }
    }

    pub async fn list_snapshots(&self, node_id: &str) -> Result<Vec<StateSnapshot>> {
        match self {
            Self::Memory(m) => m.list_snapshots(node_id).await,
            Self::Filesystem(f) => f.list_snapshots(node_id).await,
        }
    }

    pub async fn delete_snapshot(&self, snapshot_id: &str) -> Result<()> {
        match self {
            Self::Memory(m) => m.delete_snapshot(snapshot_id).await,
            Self::Filesystem(f) => f.delete_snapshot(snapshot_id).await,
        }
    }

    pub async fn compare_and_swap<T: Serialize + DeserializeOwned + Send + Sync>(
        &self,
        node_id: &str,
        expected: &T,
        new_value: &T,
    ) -> Result<bool> {
        match self {
            Self::Memory(m) => m.compare_and_swap(node_id, expected, new_value).await,
            Self::Filesystem(f) => f.compare_and_swap(node_id, expected, new_value).await,
        }
    }
}

// ============================================================================
// Helper Functions
// ============================================================================

fn now_micros() -> u64 {
    SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_micros() as u64
}

// ============================================================================
// Tests
// ============================================================================