netabase_store 0.0.8

use crate::config::FileConfig;
use crate::error::NetabaseError;
use crate::traits::backend_store::{BackendStore, PathBasedBackend};
use crate::traits::convert::ToIVec;
use crate::traits::definition::NetabaseDefinitionTrait;
use crate::traits::model::NetabaseModelTrait;
use std::marker::PhantomData;
use std::path::Path;
use std::str::FromStr;
use strum::{IntoDiscriminant, IntoEnumIterator};

use super::transaction::SledTransactionalTree;
use super::tree::SledStoreTree;
use super::types::SecondaryKeyOp;

/// Type-safe wrapper around a [sled](https://docs.rs/sled) embedded database.
///
/// `SledStore` provides a high-performance, type-safe interface to the underlying sled database,
/// using model discriminants as tree names and ensuring all operations are compile-time checked.
///
/// Sled is a modern embedded database that offers:
/// - High performance with lock-free operations
/// - Crash-safe persistence
/// - Zero-copy reads
/// - Atomic operations
///
/// # Type Parameters
///
/// * `D` - A type implementing `NetabaseDefinitionTrait`, typically generated by the
///   `#[netabase_definition_module]` macro
///
/// # Examples
///
/// ## Basic Usage
///
/// ```rust
/// use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
/// use netabase_store::databases::sled_store::SledStore;
/// use netabase_store::traits::tree::NetabaseTreeSync;
/// use netabase_store::traits::model::NetabaseModelTrait;
///
/// // Define your schema
/// #[netabase_definition_module(BlogDefinition, BlogKeys)]
/// mod blog {
///     use super::*;
///     use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
///     #[derive(NetabaseModel, Clone, Debug, PartialEq,
///              bincode::Encode, bincode::Decode,
///              serde::Serialize, serde::Deserialize)]
///     #[netabase(BlogDefinition)]
///     pub struct User {
///         #[primary_key]
///         pub id: u64,
///         pub name: String,
///         #[secondary_key]
///         pub email: String,
///     }
/// }
///
/// use blog::*;
///
/// // Create a temporary database for testing
/// let store = SledStore::<BlogDefinition>::temp().unwrap();
///
/// // Open a type-safe tree for the User model
/// let user_tree = store.open_tree::<User>();
///
/// // Insert a user
/// let alice = User {
///     id: 1,
///     name: "Alice".to_string(),
///     email: "alice@example.com".to_string(),
/// };
/// user_tree.put(alice.clone()).unwrap();
///
/// // Retrieve by primary key
/// let retrieved = user_tree.get(alice.primary_key()).unwrap();
/// assert_eq!(retrieved, Some(alice));
/// ```
///
/// ## Persistent Storage
///
/// ```rust
/// # use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
/// # use netabase_store::databases::sled_store::SledStore;
/// #
/// # #[netabase_definition_module(MyDefinition, MyKeys)]
/// # mod my_models {
/// #     use super::*;
/// #     use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
/// #     #[derive(NetabaseModel, Clone, Debug, bincode::Encode, bincode::Decode,
/// #              serde::Serialize, serde::Deserialize)]
/// #     #[netabase(MyDefinition)]
/// #     pub struct MyModel { #[primary_key] pub id: u64 }
/// # }
/// # use my_models::*;
/// // Open a persistent database at a specific path
/// let store = SledStore::<MyDefinition>::new("./my_app_data").unwrap();
///
/// // Data persists across application restarts
/// // store.flush().unwrap(); // Force flush to disk if needed
/// ```
///
/// # Performance Tips
///
/// - Use `flush()` sparingly - sled automatically persists data
/// - Batch multiple operations when possible
/// - Use `temp()` for testing to avoid disk I/O
/// - Secondary key queries are O(m) where m is the number of matches
/// - Primary key access is O(log n) where n is total records
///
/// # See Also
///
/// - [`SledStoreTree`] - Type-safe tree interface for a specific model
/// - [`RedbStore`](crate::databases::redb_store::RedbStore) - Alternative backend with ACID guarantees
pub struct SledStore<D>
where
    D: NetabaseDefinitionTrait,
    <D as IntoDiscriminant>::Discriminant: crate::traits::definition::NetabaseDiscriminant,
{
    pub(crate) db: sled::Db,
    pub trees: Vec<D::Discriminant>,
}

impl<D> SledStore<D>
where
    D: NetabaseDefinitionTrait,
    <D as IntoDiscriminant>::Discriminant: crate::traits::definition::NetabaseDiscriminant,
{
    /// Get direct access to the underlying sled database.
    ///
    /// This provides low-level access to the sled database for advanced operations
    /// that are not covered by the type-safe interface.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// # use netabase_store::databases::sled_store::SledStore;
    /// # #[netabase_definition_module(MyDef, MyKeys)]
    /// # mod my { use super::*;
    /// #   use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// #   #[derive(NetabaseModel, Clone, Debug, bincode::Encode, bincode::Decode, serde::Serialize, serde::Deserialize)]
    /// #   #[netabase(MyDef)]
    /// #   pub struct MyModel { #[primary_key] pub id: u64 }
    /// # }
    /// # use my::*;
    /// let store = SledStore::<MyDef>::temp().unwrap();
    ///
    /// // Access the underlying sled database
    /// let sled_db = store.db();
    ///
    /// // Perform low-level operations
    /// let size_on_disk = sled_db.size_on_disk().unwrap();
    /// println!("Database size: {} bytes", size_on_disk);
    /// ```
    pub fn db(&self) -> &sled::Db {
        &self.db
    }
}

impl<D> SledStore<D>
where
    D: NetabaseDefinitionTrait + ToIVec,
    <D as IntoDiscriminant>::Discriminant: crate::traits::definition::NetabaseDiscriminant,
{
    /// Open a new `SledStore` at the given filesystem path.
    ///
    /// Creates or opens a persistent sled database at the specified path. The database
    /// will automatically create the directory if it doesn't exist.
    ///
    /// # Arguments
    ///
    /// * `path` - The filesystem path where the database should be stored
    ///
    /// # Returns
    ///
    /// * `Ok(SledStore<D>)` - Successfully opened database
    /// * `Err(NetabaseError)` - Failed to open database (I/O error, permissions, etc.)
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// # use netabase_store::databases::sled_store::SledStore;
    /// # #[netabase_definition_module(AppDef, AppKeys)]
    /// # mod app { use super::*;
    /// #   use netabase_store::{NetabaseModel, netabase};
    /// #   #[derive(NetabaseModel, Clone, Debug, bincode::Encode, bincode::Decode, serde::Serialize, serde::Deserialize)]
    /// #   #[netabase(AppDef)]
    /// #   pub struct User { #[primary_key] pub id: u64, pub name: String }
    /// # }
    /// # use app::*;
    /// // Open a persistent database
    /// let store = SledStore::<AppDef>::new("./app_database").unwrap();
    ///
    /// // Data persists across application restarts
    /// let user_tree = store.open_tree::<User>();
    /// ```
    ///
    /// ```rust
    /// # use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// # use netabase_store::databases::sled_store::SledStore;
    /// # #[netabase_definition_module(TestDef, TestKeys)]
    /// # mod test { use super::*;
    /// #   use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// #   #[derive(NetabaseModel, Clone, Debug, bincode::Encode, bincode::Decode, serde::Serialize, serde::Deserialize)]
    /// #   #[netabase(TestDef)]
    /// #   pub struct Item { #[primary_key] pub id: u64 }
    /// # }
    /// # use test::*;
    /// // Use a temporary directory for testing
    /// let temp_dir = tempfile::tempdir().unwrap();
    /// let db_path = temp_dir.path().join("test_db");
    /// let store = SledStore::<TestDef>::new(&db_path).unwrap();
    /// ```
    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, NetabaseError> {
        let db = sled::open(path)?;
        Ok(Self {
            db,
            trees: D::Discriminant::iter().collect(),
        })
    }

    /// Create a temporary in-memory `SledStore` for testing.
    ///
    /// Creates a temporary database that exists only in memory and is automatically
    /// cleaned up when dropped. This is ideal for unit tests and examples where
    /// persistence is not needed.
    ///
    /// # Returns
    ///
    /// * `Ok(SledStore<D>)` - Successfully created temporary database
    /// * `Err(NetabaseError)` - Failed to create database (rare, usually indicates system issues)
    ///
    /// # Examples
    ///
    /// ```rust
    /// use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// use netabase_store::databases::sled_store::SledStore;
    /// use netabase_store::traits::tree::NetabaseTreeSync;
    /// use netabase_store::traits::model::NetabaseModelTrait;
    ///
    /// #[netabase_definition_module(TestDef, TestKeys)]
    /// mod test {
    ///     use super::*;
    ///     use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    ///     #[derive(NetabaseModel, Clone, Debug, PartialEq,
    ///              bincode::Encode, bincode::Decode,
    ///              serde::Serialize, serde::Deserialize)]
    ///     #[netabase(TestDef)]
    ///     pub struct Product {
    ///         #[primary_key]
    ///         pub id: String,
    ///         pub name: String,
    ///         pub price: f64,
    ///     }
    /// }
    /// use test::*;
    ///
    /// // Create temporary database
    /// let store = SledStore::<TestDef>::temp().unwrap();
    /// let product_tree = store.open_tree::<Product>();
    ///
    /// // Use normally
    /// let product = Product {
    ///     id: "prod-1".to_string(),
    ///     name: "Widget".to_string(),
    ///     price: 29.99,
    /// };
    /// product_tree.put(product.clone()).unwrap();
    ///
    /// // Verify
    /// let retrieved = product_tree.get(product.primary_key()).unwrap();
    /// assert_eq!(retrieved, Some(product));
    /// // Database is automatically cleaned up when `store` is dropped
    /// ```
    pub fn temp() -> Result<Self, NetabaseError> {
        let config = sled::Config::new().temporary(true);
        let db = config.open()?;
        Ok(Self {
            db,
            trees: D::Discriminant::iter().collect(),
        })
    }

    /// Open a type-safe tree for a specific model type.
    ///
    /// Returns a `SledStoreTree` that provides CRUD operations for a single model type.
    /// Each model type gets its own tree within the database, identified by its discriminant.
    ///
    /// The returned tree is type-safe and will only accept/return instances of the specified model.
    ///
    /// # Type Parameters
    ///
    /// * `M` - The model type, must implement `NetabaseModelTrait<D>`
    ///
    /// # Returns
    ///
    /// A `SledStoreTree<D, M>` ready for CRUD operations on the specified model type.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// use netabase_store::databases::sled_store::SledStore;
    /// use netabase_store::traits::tree::NetabaseTreeSync;
    ///
    /// #[netabase_definition_module(ShopDef, ShopKeys)]
    /// mod shop {
    ///     use super::*;
    ///
    /// use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    ///
    ///     #[derive(NetabaseModel, Clone, Debug, PartialEq,
    ///              bincode::Encode, bincode::Decode,
    ///              serde::Serialize, serde::Deserialize)]
    ///     #[netabase(ShopDef)]
    ///     pub struct Customer {
    ///         #[primary_key]
    ///         pub id: u64,
    ///         pub name: String,
    ///         #[secondary_key]
    ///         pub email: String,
    ///     }
    ///
    ///     #[derive(NetabaseModel, Clone, Debug, PartialEq,
    ///              bincode::Encode, bincode::Decode,
    ///              serde::Serialize, serde::Deserialize)]
    ///     #[netabase(ShopDef)]
    ///     pub struct Order {
    ///         #[primary_key]
    ///         pub id: u64,
    ///         pub customer_id: u64,
    ///         pub total: f64,
    ///     }
    /// }
    /// use shop::*;
    ///
    /// let store = SledStore::<ShopDef>::temp().unwrap();
    ///
    /// // Open separate trees for each model
    /// let customer_tree = store.open_tree::<Customer>();
    /// let order_tree = store.open_tree::<Order>();
    ///
    /// // Each tree is independent and type-safe
    /// let customer = Customer {
    ///     id: 1,
    ///     name: "Alice".to_string(),
    ///     email: "alice@example.com".to_string(),
    /// };
    /// customer_tree.put(customer).unwrap();
    ///
    /// let order = Order {
    ///     id: 100,
    ///     customer_id: 1,
    ///     total: 99.99,
    /// };
    /// order_tree.put(order).unwrap();
    ///
    /// // Verify independent storage
    /// assert_eq!(customer_tree.len(), 1);
    /// assert_eq!(order_tree.len(), 1);
    /// ```
    ///
    /// The returned tree is tied to the lifetime of this store, ensuring it cannot outlive the database.
    pub fn open_tree<M>(&self) -> SledStoreTree<'_, D, M>
    where
        M: NetabaseModelTrait<D> + TryFrom<D> + Into<D>,
        D: TryFrom<M> + ToIVec,
    {
        SledStoreTree::new(&self.db, M::DISCRIMINANT)
    }

    pub fn open_t(&self, _key: D::Keys) {
        // TODO: Implement this method properly
        // let _inner = key.inner();
        // let _ = self.open_tree::<M>();
    }

    /// Get all model discriminants (tree names) in the database schema.
    ///
    /// Returns a vector of all discriminants defined in the definition enum.
    /// Each discriminant corresponds to a model type and its associated tree.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// # use netabase_store::databases::sled_store::SledStore;
    /// #[netabase_definition_module(MultiModelDef, MultiModelKeys)]
    /// mod models {
    ///     use super::*;
    /// use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    ///     #[derive(NetabaseModel, Clone, Debug, bincode::Encode, bincode::Decode,
    ///              serde::Serialize, serde::Deserialize)]
    ///     #[netabase(MultiModelDef)]
    ///     pub struct User { #[primary_key] pub id: u64 }
    ///
    ///     #[derive(NetabaseModel, Clone, Debug, bincode::Encode, bincode::Decode,
    ///              serde::Serialize, serde::Deserialize)]
    ///     #[netabase(MultiModelDef)]
    ///     pub struct Post { #[primary_key] pub id: u64 }
    /// }
    /// use models::*;
    ///
    /// let store = SledStore::<MultiModelDef>::temp().unwrap();
    /// let tree_names = store.tree_names();
    ///
    /// // Will contain discriminants for both User and Post
    /// assert_eq!(tree_names.len(), 2);
    /// ```
    pub fn tree_names(&self) -> Vec<D::Discriminant> {
        D::Discriminant::iter().collect()
    }

    /// Flush all pending writes to disk.
    ///
    /// Forces sled to flush all buffered writes to persistent storage.
    /// This ensures data durability but may impact performance if called too frequently.
    ///
    /// Sled automatically flushes data periodically, so manual flushes are typically
    /// only needed when:
    /// - You need to guarantee data is on disk before continuing
    /// - Before shutting down the application
    /// - After critical writes that must not be lost
    ///
    /// # Returns
    ///
    /// * `Ok(usize)` - Number of bytes written to disk
    /// * `Err(NetabaseError)` - I/O error during flush
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// # use netabase_store::databases::sled_store::SledStore;
    /// # #[netabase_definition_module(AppDef, AppKeys)]
    /// # mod app { use super::*;
    /// #   use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// #   #[derive(NetabaseModel, Clone, Debug, bincode::Encode, bincode::Decode,
    /// #            serde::Serialize, serde::Deserialize)]
    /// #   #[netabase(AppDef)]
    /// #   pub struct Config { #[primary_key] pub key: String, pub value: String }
    /// # }
    /// # use app::*;
    /// # let temp_dir = tempfile::tempdir().unwrap();
    /// # let db_path = temp_dir.path().join("db");
    /// let store = SledStore::<AppDef>::new(&db_path).unwrap();
    /// let config_tree = store.open_tree::<Config>();
    ///
    /// // Write critical configuration
    /// let config = Config {
    ///     key: "api_key".to_string(),
    ///     value: "secret123".to_string(),
    /// };
    /// config_tree.put(config).unwrap();
    ///
    /// // Ensure it's written to disk before continuing
    /// let bytes_flushed = store.flush().unwrap();
    /// println!("Flushed {} bytes to disk", bytes_flushed);
    /// ```
    pub fn flush(&self) -> Result<usize, NetabaseError> {
        Ok(self.db.flush()?)
    }

    /// Execute a transaction on a single model tree.
    ///
    /// Provides ACID transaction semantics for operations on a single model type.
    /// The closure receives a `SledTransactionalTree` that provides atomic operations.
    /// All operations either succeed together or fail together.
    ///
    /// The transaction closure may be called multiple times if there are conflicts,
    /// so it should be idempotent and avoid side effects.
    ///
    /// # Type Parameters
    ///
    /// * `M` - The model type to operate on
    /// * `F` - The transaction closure
    /// * `R` - The return type
    ///
    /// # Arguments
    ///
    /// * `f` - Transaction closure that performs operations on the transactional tree
    ///
    /// # Returns
    ///
    /// * `Ok(R)` - Transaction succeeded, returns result from closure
    /// * `Err(NetabaseError)` - Transaction failed (conflict, I/O error, etc.)
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// # use netabase_store::databases::sled_store::SledStore;
    /// # use netabase_store::traits::model::NetabaseModelTrait;
    /// # #[netabase_definition_module(BankDef, BankKeys)]
    /// # mod bank {
    /// #     use super::*;
    /// #     use netabase_store::{netabase_definition_module, NetabaseModel, netabase};
    /// #     #[derive(NetabaseModel, Clone, Debug, PartialEq, bincode::Encode, bincode::Decode,
    /// #              serde::Serialize, serde::Deserialize)]
    /// #     #[netabase(BankDef)]
    /// #     pub struct Account {
    /// #         #[primary_key] pub id: u64,
    /// #         pub balance: i64,
    /// #     }
    /// # }
    /// # use bank::*;
    /// # use netabase_store::traits::tree::NetabaseTreeSync;
    /// let store = SledStore::<BankDef>::temp().unwrap();
    /// let tree = store.open_tree::<Account>();
    ///
    /// // Set up initial accounts
    /// tree.put(Account { id: 1, balance: 1000 }).unwrap();
    /// tree.put(Account { id: 2, balance: 500 }).unwrap();
    ///
    /// // Atomic transfer between accounts
    /// store.transaction::<Account, _, _>(|txn_tree| {
    ///     // Get both accounts
    ///     let mut from_account = txn_tree.get(Account { id: 1, balance: 0 }.primary_key())?
    ///         .ok_or_else(|| "Account not found")?;
    ///     let mut to_account = txn_tree.get(Account { id: 2, balance: 0 }.primary_key())?
    ///         .ok_or_else(|| "Account not found")?;
    ///
    ///     // Transfer $100
    ///     from_account.balance -= 100;
    ///     to_account.balance += 100;
    ///
    ///     // Update both atomically
    ///     txn_tree.put(from_account)?;
    ///     txn_tree.put(to_account)?;
    ///
    ///     Ok(())
    /// }).unwrap();
    /// ```
    pub fn transaction<M, F, R>(&self, f: F) -> Result<R, NetabaseError>
    where
        M: NetabaseModelTrait<D> + TryFrom<D> + Into<D>,
        D: TryFrom<M>,
        F: Fn(&SledTransactionalTree<D, M>) -> Result<R, Box<dyn std::error::Error>>,
    {
        let tree = self.db.open_tree(M::DISCRIMINANT.to_string())?;
        let sec_tree_name = format!("{}_secondary", M::discriminant_name());
        let secondary_tree = self.db.open_tree(sec_tree_name)?;

        // Collect secondary key operations during transaction
        let pending_secondary_ops = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
        let pending_ops_clone = pending_secondary_ops.clone();

        let result = tree
            .transaction(|txn_tree| {
                let wrapper = SledTransactionalTree {
                    tree: txn_tree.clone(),
                    secondary_tree: secondary_tree.clone(),
                    pending_secondary_keys: Some(pending_ops_clone.clone()),
                    _phantom_d: PhantomData,
                    _phantom_m: PhantomData,
                };

                f(&wrapper).map_err(sled::transaction::ConflictableTransactionError::Abort)
            })
            .map_err(|e| match e {
                sled::transaction::TransactionError::Abort(e) => {
                    NetabaseError::Transaction(format!("Transaction aborted: {}", e))
                }
                sled::transaction::TransactionError::Storage(e) => e.into(),
            })?;

        // After transaction commits, apply all pending secondary key operations
        let ops = pending_secondary_ops.lock().unwrap();
        for op in ops.iter() {
            match op {
                SecondaryKeyOp::Insert(key) => {
                    secondary_tree.insert(key, &[] as &[u8])?;
                }
                SecondaryKeyOp::Remove(key) => {
                    secondary_tree.remove(key)?;
                }
            }
        }

        Ok(result)
    }
}

// ============================================================================
// BackendStore Trait Implementation
// ============================================================================

impl<D> BackendStore<D> for SledStore<D>
where
    D: NetabaseDefinitionTrait,
    <D as IntoDiscriminant>::Discriminant: Clone
        + Copy
        + std::fmt::Debug
        + std::fmt::Display
        + PartialEq
        + Eq
        + std::hash::Hash
        + strum::IntoEnumIterator
        + 'static
        + FromStr,
{
    type Config = FileConfig;

    fn new(config: Self::Config) -> Result<Self, NetabaseError> {
        // Remove existing database if truncate is true
        if config.truncate && config.path.exists() {
            std::fs::remove_dir_all(&config.path)?;
        }

        let db = if config.create_if_missing {
            sled::open(&config.path)?
        } else {
            sled::Config::new()
                .path(&config.path)
                .cache_capacity((config.cache_size_mb * 1024 * 1024) as u64)
                .use_compression(true)
                .open()?
        };

        Ok(Self {
            db,
            trees: D::Discriminant::iter().collect(),
        })
    }

    fn open(config: Self::Config) -> Result<Self, NetabaseError> {
        let db = sled::Config::new()
            .path(&config.path)
            .cache_capacity((config.cache_size_mb * 1024 * 1024) as u64)
            .use_compression(true)
            .open()?;

        Ok(Self {
            db,
            trees: D::Discriminant::iter().collect(),
        })
    }

    fn temp() -> Result<Self, NetabaseError> {
        let config = FileConfig::temp();
        <Self as BackendStore<D>>::new(config)
    }
}

impl<D> PathBasedBackend<D> for SledStore<D>
where
    D: NetabaseDefinitionTrait,
    <D as IntoDiscriminant>::Discriminant: Clone
        + Copy
        + std::fmt::Debug
        + std::fmt::Display
        + PartialEq
        + Eq
        + std::hash::Hash
        + strum::IntoEnumIterator
        + 'static
        + FromStr,
{
    fn at_path<P: AsRef<Path>>(path: P) -> Result<Self, NetabaseError> {
        let config = FileConfig::new(path.as_ref());
        <Self as BackendStore<D>>::open(config)
    }
}