guardian-db 0.19.0

use crate::address::Address;
use crate::cache::{Cache, Options};
use crate::data_store::{Datastore, Key, Order, Query, ResultItem, Results};
use crate::guardian::error::{GuardianError, Result};
use redb::{Database, ReadableDatabase, ReadableTable, TableDefinition};
use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    sync::{Arc, Mutex, Weak},
};
use tracing::{Span, debug, instrument};

pub const IN_MEMORY_DIRECTORY: &str = ":memory:";

const LEVEL_DOWN_TABLE: TableDefinition<&[u8], &[u8]> = TableDefinition::new("level_down");

/// A single redb-backed cache instance, keyed by its datastore path and
/// tracked by the owning `LevelDownCache` manager.
pub struct WrappedCache {
    id: String,
    db: Database,
    manager_map: Weak<Mutex<HashMap<String, Arc<WrappedCache>>>>,
    #[allow(dead_code)]
    span: Span,
    closed: Mutex<bool>,
}

impl WrappedCache {
    /// Returns the value stored under `key`, or an error if it is missing.
    #[instrument(level = "debug", skip(self, _ctx))]
    pub fn get(
        &self,
        _ctx: &mut dyn core::any::Any,
        key: &Key,
    ) -> std::result::Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync>> {
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(LEVEL_DOWN_TABLE)?;
        match table.get(key.as_bytes().as_slice())? {
            Some(v) => Ok(v.value().to_vec()),
            None => Err(format!("key not found: {}", key).into()),
        }
    }

    /// Returns whether a value exists under `key`.
    pub fn has(
        &self,
        _ctx: &mut dyn core::any::Any,
        key: &Key,
    ) -> std::result::Result<bool, Box<dyn std::error::Error + Send + Sync>> {
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(LEVEL_DOWN_TABLE)?;
        Ok(table.get(key.as_bytes().as_slice())?.is_some())
    }

    /// Returns the byte length of the value stored under `key`.
    pub fn get_size(
        &self,
        _ctx: &mut dyn core::any::Any,
        key: &Key,
    ) -> std::result::Result<usize, Box<dyn std::error::Error + Send + Sync>> {
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(LEVEL_DOWN_TABLE)?;
        let v = table
            .get(key.as_bytes().as_slice())?
            .ok_or_else(|| format!("key not found: {}", key))?;
        Ok(v.value().len())
    }

    /// Runs a query against the cache, supporting optional prefix filtering,
    /// offset, limit and ordering.
    pub fn query(
        &self,
        _ctx: &mut dyn core::any::Any,
        q: &Query,
    ) -> std::result::Result<Results, Box<dyn std::error::Error + Send + Sync>> {
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(LEVEL_DOWN_TABLE)?;

        let mut items: Results = Vec::new();
        let mut count = 0;
        let skip_count = q.offset.unwrap_or(0);
        let mut skipped = 0;

        if let Some(prefix_key) = &q.prefix {
            let prefix_bytes = prefix_key.as_bytes();
            let iter = table.range(prefix_bytes.as_slice()..)?;

            for entry_result in iter {
                let entry = entry_result?;
                let key_bytes = entry.0.value();
                if !key_bytes.starts_with(&prefix_bytes) {
                    break;
                }

                if skipped < skip_count {
                    skipped += 1;
                    continue;
                }

                let key_str = String::from_utf8(key_bytes.to_vec()).unwrap_or_default();
                items.push(ResultItem::new(Key::new(key_str), entry.1.value().to_vec()));
                count += 1;

                if let Some(n) = q.limit
                    && count >= n
                {
                    break;
                }
            }
        } else {
            let iter = table.iter()?;

            for entry_result in iter {
                let entry = entry_result?;

                if skipped < skip_count {
                    skipped += 1;
                    continue;
                }

                let key_bytes = entry.0.value();
                let key_str = String::from_utf8(key_bytes.to_vec()).unwrap_or_default();
                items.push(ResultItem::new(Key::new(key_str), entry.1.value().to_vec()));
                count += 1;

                if let Some(n) = q.limit
                    && count >= n
                {
                    break;
                }
            }
        }

        if matches!(q.order, Order::Desc) {
            items.reverse();
        }

        Ok(items)
    }

    /// Stores `value` under `key`, committing the write transaction.
    #[instrument(level = "debug", skip(self, _ctx, value))]
    pub fn put(
        &self,
        _ctx: &mut dyn core::any::Any,
        key: &Key,
        value: &[u8],
    ) -> std::result::Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let write_txn = self.db.begin_write()?;
        {
            let mut table = write_txn.open_table(LEVEL_DOWN_TABLE)?;
            table.insert(key.as_bytes().as_slice(), value)?;
        }
        write_txn.commit()?;
        Ok(())
    }

    /// Removes the value stored under `key`, committing the write transaction.
    pub fn delete(
        &self,
        _ctx: &mut dyn core::any::Any,
        key: &Key,
    ) -> std::result::Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let write_txn = self.db.begin_write()?;
        {
            let mut table = write_txn.open_table(LEVEL_DOWN_TABLE)?;
            table.remove(key.as_bytes().as_slice())?;
        }
        write_txn.commit()?;
        Ok(())
    }

    /// No-op flush: data is already persisted via redb write transactions.
    pub fn sync(
        &self,
        _ctx: &mut dyn core::any::Any,
        _key: &Key,
    ) -> std::result::Result<(), Box<dyn std::error::Error + Send + Sync>> {
        // Data is already persisted via write transactions in redb.
        Ok(())
    }

    /// Closes the cache and unregisters it from the manager map (idempotent).
    #[instrument(level = "debug", skip(self))]
    pub fn close(&self) -> std::result::Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let mut closed = self.closed.lock().unwrap();
        if *closed {
            return Ok(());
        }

        if let Some(map) = self.manager_map.upgrade() {
            let mut m = map.lock().unwrap();
            m.remove(&self.id);
        }

        // Data is already persisted via write transactions in redb.
        *closed = true;
        Ok(())
    }
}

// Send + Sync is safe because redb::Database is thread-safe
unsafe impl Send for WrappedCache {}
unsafe impl Sync for WrappedCache {}

/// Wrapper that adapts a `WrappedCache` to the `Datastore` trait.
pub struct DatastoreWrapper {
    cache: Arc<WrappedCache>,
}

impl DatastoreWrapper {
    /// Wraps a shared `WrappedCache` in the adapter.
    pub fn new(cache: Arc<WrappedCache>) -> Self {
        Self { cache }
    }
}

#[async_trait::async_trait]
impl Datastore for DatastoreWrapper {
    #[instrument(level = "debug", skip(self, key))]
    async fn has(&self, key: &[u8]) -> Result<bool> {
        let key_obj = Key::new(String::from_utf8_lossy(key));
        let mut any_ctx = ();
        self.cache
            .has(&mut any_ctx, &key_obj)
            .map_err(|e| GuardianError::Other(format!("Cache has error: {}", e)))
    }

    #[instrument(level = "debug", skip(self, key, value))]
    async fn put(&self, key: &[u8], value: &[u8]) -> Result<()> {
        let key_obj = Key::new(String::from_utf8_lossy(key));
        let mut any_ctx = ();
        self.cache
            .put(&mut any_ctx, &key_obj, value)
            .map_err(|e| GuardianError::Other(format!("Cache put error: {}", e)))
    }

    #[instrument(level = "debug", skip(self, key))]
    async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let key_obj = Key::new(String::from_utf8_lossy(key));
        let mut any_ctx = ();
        match self.cache.get(&mut any_ctx, &key_obj) {
            Ok(value) => Ok(Some(value)),
            Err(_) => Ok(None), // key not found
        }
    }

    #[instrument(level = "debug", skip(self, key))]
    async fn delete(&self, key: &[u8]) -> Result<()> {
        let key_obj = Key::new(String::from_utf8_lossy(key));
        let mut any_ctx = ();
        self.cache
            .delete(&mut any_ctx, &key_obj)
            .map_err(|e| GuardianError::Other(format!("Cache delete error: {}", e)))
    }

    #[instrument(level = "debug", skip(self, query))]
    async fn query(&self, query: &Query) -> Result<Results> {
        let mut any_ctx = ();
        self.cache
            .query(&mut any_ctx, query)
            .map_err(|e| GuardianError::Other(format!("Cache query error: {}", e)))
    }

    #[instrument(level = "debug", skip(self, prefix))]
    async fn list_keys(&self, prefix: &[u8]) -> Result<Vec<Key>> {
        // Convert the byte prefix into a prefix Query.
        let prefix_str = String::from_utf8_lossy(prefix);
        let prefix_key = Key::new(prefix_str.to_string());

        let query = Query {
            prefix: Some(prefix_key),
            limit: None,
            order: Order::Asc,
            offset: None,
        };

        let mut any_ctx = ();
        let results = self
            .cache
            .query(&mut any_ctx, &query)
            .map_err(|e| GuardianError::Other(format!("Cache list_keys error: {}", e)))?;

        // Extract only the keys from the results.
        Ok(results.into_iter().map(|item| item.key).collect())
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
}

/// Cache manager that owns a set of `WrappedCache` instances keyed by their
/// datastore path, and implements the `Cache` trait.
pub struct LevelDownCache {
    span: Span,
    caches: Arc<Mutex<HashMap<String, Arc<WrappedCache>>>>,
}

impl LevelDownCache {
    /// Creates a new, empty cache manager.
    #[instrument(level = "debug")]
    pub fn new(_opts: Option<&Options>) -> Self {
        Self {
            span: tracing::Span::current(),
            caches: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Returns a reference to the tracing span used for instrumentation.
    pub fn span(&self) -> &Span {
        &self.span
    }

    /// Loads (or creates) the `WrappedCache` for a directory and address,
    /// reusing an existing instance on a cache hit.
    #[instrument(level = "debug", skip(self, db_address))]
    pub fn load_internal(
        &self,
        directory: &str,
        db_address: &dyn Address,
    ) -> std::result::Result<Arc<WrappedCache>, Box<dyn std::error::Error + Send + Sync>> {
        let _entered = self.span.enter();
        let key_path = datastore_key(directory, db_address);

        // cache hit
        if let Some(ds) = self.caches.lock().unwrap().get(&key_path).cloned() {
            return Ok(ds);
        }

        debug!("opening cache db: path={}", key_path.as_str());

        let db = if directory == IN_MEMORY_DIRECTORY {
            Database::builder().create_with_backend(redb::backends::InMemoryBackend::new())?
        } else {
            if let Some(parent) = Path::new(&key_path).parent() {
                std::fs::create_dir_all(parent)?;
            }
            Database::create(&key_path)?
        };

        // Ensure table exists
        {
            let write_txn = db.begin_write()?;
            {
                let _ = write_txn.open_table(LEVEL_DOWN_TABLE)?;
            }
            write_txn.commit()?;
        }

        let wrapped = Arc::new(WrappedCache {
            id: key_path.clone(),
            db,
            manager_map: Arc::downgrade(&self.caches),
            span: tracing::Span::current(),
            closed: Mutex::new(false),
        });

        self.caches
            .lock()
            .unwrap()
            .insert(key_path, wrapped.clone());
        Ok(wrapped)
    }

    /// Closes every cache instance owned by this manager.
    #[instrument(level = "debug", skip(self))]
    pub fn close_internal(
        &self,
    ) -> std::result::Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let _entered = self.span.enter();
        let caches = {
            let m = self.caches.lock().unwrap();
            m.values().cloned().collect::<Vec<_>>()
        };
        for c in caches {
            let _ = c.close();
        }
        Ok(())
    }

    /// Closes and removes the cache for a directory/address, deleting its file
    /// from disk when it is not an in-memory cache.
    #[instrument(level = "debug", skip(self, db_address))]
    pub fn destroy_internal(
        &self,
        directory: &str,
        db_address: &dyn Address,
    ) -> std::result::Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let _entered = self.span.enter();
        let key_path = datastore_key(directory, db_address);

        // Close it and remove it from the map.
        if let Some(c) = self.caches.lock().unwrap().remove(&key_path) {
            let _ = c.close();
        }

        if directory != IN_MEMORY_DIRECTORY && Path::new(&key_path).exists() {
            std::fs::remove_file(&key_path)?;
        }

        Ok(())
    }
}

// Cache trait implementation for LevelDownCache.
impl Cache for LevelDownCache {
    #[instrument(level = "info", skip(self, db_address))]
    fn load(
        &self,
        directory: &str,
        db_address: &dyn Address,
    ) -> Result<Box<dyn Datastore + Send + Sync>> {
        let _entered = self.span.enter();
        let wrapped_cache = self
            .load_internal(directory, db_address)
            .map_err(|e| GuardianError::Other(format!("Failed to load cache: {}", e)))?;
        Ok(Box::new(DatastoreWrapper {
            cache: wrapped_cache,
        }))
    }

    #[instrument(level = "info", skip(self))]
    fn close(&mut self) -> Result<()> {
        let _entered = self.span.enter();
        let caches = {
            let m = self.caches.lock().unwrap();
            m.values().cloned().collect::<Vec<_>>()
        };
        for c in caches {
            let _ = c.close();
        }
        Ok(())
    }

    #[instrument(level = "info", skip(self, db_address))]
    fn destroy(&self, directory: &str, db_address: &dyn Address) -> Result<()> {
        let _entered = self.span.enter();
        self.destroy_internal(directory, db_address)
            .map_err(|e| GuardianError::Other(format!("Failed to destroy cache: {}", e)))?;
        Ok(())
    }
}

/// Builds the on-disk/in-memory key path for a datastore from its directory and
/// database address (root hash joined with the address path).
fn datastore_key(directory: &str, db_address: &dyn Address) -> String {
    let db_path = PathBuf::from(db_address.get_root().to_string()).join(db_address.get_path());
    PathBuf::from(directory)
        .join(db_path)
        .to_string_lossy()
        .into_owned()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::address::Address;
    use std::fmt;

    // Mock Address implementation for testing
    #[derive(Debug)]
    struct MockAddress {
        root: iroh_blobs::Hash,
        path: String,
    }

    impl MockAddress {
        fn new(root_str: &str, path: &str) -> Self {
            // Create a Hash from the root string for more meaningful testing.
            // For the test, use a Hash derived from root_str.
            use blake3;
            let hash_bytes: [u8; 32] = blake3::hash(root_str.as_bytes()).into();
            let hash = iroh_blobs::Hash::from(hash_bytes);
            Self {
                root: hash,
                path: path.to_string(),
            }
        }
    }

    impl Address for MockAddress {
        fn get_root(&self) -> iroh_blobs::Hash {
            self.root
        }

        fn get_path(&self) -> &str {
            &self.path
        }

        fn equals(&self, other: &dyn Address) -> bool {
            self.root == other.get_root() && self.path == other.get_path()
        }
    }

    impl fmt::Display for MockAddress {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "{}/{}", hex::encode(self.root.as_bytes()), self.path)
        }
    }

    #[tokio::test]
    async fn test_datastore_wrapper_basic_operations() {
        let cache = LevelDownCache::new(None);
        let mock_address = MockAddress::new("test_root", "test_path");

        let datastore = cache.load(IN_MEMORY_DIRECTORY, &mock_address).unwrap();

        // Test put and get
        let key = b"test_key";
        let value = b"test_value";

        datastore.put(key, value).await.unwrap();
        let retrieved = datastore.get(key).await.unwrap();
        assert_eq!(retrieved, Some(value.to_vec()));

        // Test has
        assert!(datastore.has(key).await.unwrap());
        assert!(!datastore.has(b"non_existent").await.unwrap());

        // Test delete
        datastore.delete(key).await.unwrap();
        assert!(!datastore.has(key).await.unwrap());
    }

    #[tokio::test]
    async fn test_datastore_wrapper_query() {
        let cache = LevelDownCache::new(None);
        let mock_address = MockAddress::new("test_root", "test_path");

        let datastore = cache.load(IN_MEMORY_DIRECTORY, &mock_address).unwrap();

        // Insert test data
        datastore.put(b"/users/alice", b"alice_data").await.unwrap();
        datastore.put(b"/users/bob", b"bob_data").await.unwrap();
        datastore
            .put(b"/config/database", b"db_config")
            .await
            .unwrap();

        // Test query with prefix
        let query = Query {
            prefix: Some(Key::new("/users")),
            limit: Some(10),
            order: Order::Asc,
            offset: None,
        };

        let results = datastore.query(&query).await.unwrap();
        assert_eq!(results.len(), 2);

        // Test list_keys
        let keys = datastore.list_keys(b"/users").await.unwrap();
        assert_eq!(keys.len(), 2);
    }

    #[test]
    fn test_datastore_key_generation() {
        let mock_address = MockAddress::new("root", "path/to/db");
        let key = datastore_key("/cache", &mock_address);

        // Debug: inspect what is being generated.
        println!("Generated key: {}", key);
        println!("Root Hash: {}", mock_address.get_root());
        println!("Path: {}", mock_address.get_path());

        // The exact format depends on the platform path separator
        assert!(key.contains("cache"));
        assert!(key.contains("path"));
        assert!(key.contains(&mock_address.get_root().to_string()));
    }

    #[tokio::test]
    #[ignore] // Test takes too long in CI environment
    async fn test_cache_lifecycle() {
        let mut cache = LevelDownCache::new(None);
        let mock_address = MockAddress::new("test_root", "lifecycle_test");

        // Load cache
        let datastore = cache.load(IN_MEMORY_DIRECTORY, &mock_address).unwrap();
        datastore.put(b"test", b"data").await.unwrap();

        // Destroy cache
        cache.destroy(IN_MEMORY_DIRECTORY, &mock_address).unwrap();

        // Close cache
        cache.close().unwrap();
    }
}