trelent-hyok 0.1.12

use crate::dek::persistence::DEKPersisterTrait;
use crate::error::generator::PersistError;
use async_trait::async_trait;
use std::sync::Arc;

/// Custom persistence implementation for DEK storage.
///
/// This module enables integration with any storage backend by allowing
/// users to provide their own persistence functions. It supports:
///
/// - Custom storage implementations
/// - Async operations
/// - Type-safe error handling
/// - Flexible metadata handling
///
/// # Example
/// ```no_run
/// use std::collections::HashMap;
/// use std::sync::Arc;
/// use tokio::sync::Mutex;
/// use hyokashi::{CustomPersister, PersistError};
///
/// // Create an in-memory store
/// let store = Arc::new(Mutex::new(HashMap::new()));
/// let store_clone = store.clone();
///
/// // Define persistence functions
/// let persist = move |key: Vec<u8>, ctx: String| {
///     let store = store.clone();
///     Box::pin(async move {
///         let mut map = store.lock().await;
///         map.insert(ctx, key.clone());
///         Ok(key)
///     })
/// };
///
/// let fetch = move |ctx: String| {
///     let store = store_clone.clone();
///     Box::pin(async move {
///         let map = store.lock().await;
///         map.get(&ctx)
///            .cloned()
///            .ok_or_else(|| PersistError::Error("Key not found".into()))
///     })
/// };
///
/// let persister = CustomPersister::new(persist, fetch);
/// ```
pub struct CustomPersister {
    persist_fn: Arc<
        dyn (Fn(
            Vec<u8>,
            String
        ) -> futures::future::BoxFuture<'static, Result<Vec<u8>, PersistError>>) +
            Send +
            Sync
    >,
    fetch_fn: Arc<
        dyn (Fn(String) -> futures::future::BoxFuture<'static, Result<Vec<u8>, PersistError>>) +
            Send +
            Sync
    >,
}

impl CustomPersister {
    /// Creates a new `CustomPersister` with the provided persistence functions.
    ///
    /// # Arguments
    ///
    /// * `persist_fn` - Function for storing keys
    /// * `fetch_fn` - Function for retrieving keys
    ///
    /// # Example
    /// ```no_run
    /// use std::collections::HashMap;
    /// use std::sync::Arc;
    /// use tokio::sync::Mutex;
    ///
    /// // Create an in-memory store
    /// let store = Arc::new(Mutex::new(HashMap::new()));
    /// let store_clone = store.clone();
    ///
    /// // Define persist function
    /// let persist = move |key: Vec<u8>, ctx: String| {
    ///     let store = store.clone();
    ///     Box::pin(async move {
    ///         let mut map = store.lock().await;
    ///         map.insert(ctx, key.clone());
    ///         Ok(key)
    ///     })
    /// };
    ///
    /// // Define fetch function
    /// let fetch = move |ctx: String| {
    ///     let store = store_clone.clone();
    ///     Box::pin(async move {
    ///         let map = store.lock().await;
    ///         map.get(&ctx)
    ///            .cloned()
    ///            .ok_or_else(|| PersistError::Error("Key not found".into()))
    ///     })
    /// };
    ///
    /// let persister = CustomPersister::new(persist, fetch);
    /// ```
    pub fn new<P, F>(persist_fn: P, fetch_fn: F) -> Self
        where
            P: Fn(
                Vec<u8>,
                String
            ) -> futures::future::BoxFuture<'static, Result<Vec<u8>, PersistError>> +
                Send +
                Sync +
                'static,
            F: Fn(String) -> futures::future::BoxFuture<'static, Result<Vec<u8>, PersistError>> +
                Send +
                Sync +
                'static
    {
        CustomPersister {
            persist_fn: Arc::new(persist_fn),
            fetch_fn: Arc::new(fetch_fn),
        }
    }
}

#[async_trait]
impl DEKPersisterTrait for CustomPersister {
    /// Persists a dek using the user-provided persistence function.
    ///
    /// # Errors
    ///
    /// Returns a `PersistError` if the persistence function returns an error
    async fn persist(&self, dek: &Vec<u8>, context: String) -> Result<Vec<u8>, PersistError> {
        (self.persist_fn)(dek.clone(), context).await
    }
    /// Retrieves a dek using the user-provided fetch function.
    ///
    /// # Errors
    ///
    /// Returns a `PersistError` if the fetch function returns an error
    async fn fetch(&self, context: String) -> Result<Vec<u8>, PersistError> {
        (self.fetch_fn)(context).await
    }
}