mecha10-core 0.6.2

//! `StateManager` implementation for [`super::MemoryStateManager`].

use super::{MemoryStateManager, StateManager, StateOptions, StateSnapshot, VersionedState};
use crate::error::{Mecha10Error, Result};
use crate::prelude::now_micros;
use async_trait::async_trait;
use serde::{de::DeserializeOwned, Serialize};

#[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 let Some(ttl) = options.ttl {
            VersionedState::with_ttl(state, options.version, ttl)
        } 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)
    }
}