tenzro-storage 0.1.0

//! Key-value store abstraction for Tenzro Network
//!
//! This module provides a unified key-value store interface with
//! implementations for RocksDB and in-memory storage.

use crate::config::StorageConfig;
use crate::error::{Result, StorageError};
use parking_lot::RwLock;
use rocksdb::{ColumnFamily, ColumnFamilyDescriptor, Options, WriteBatch, WriteOptions, DB};
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;

/// Column family names used in the storage layer
pub const CF_BLOCKS: &str = "blocks";
pub const CF_STATE: &str = "state";
pub const CF_ACCOUNTS: &str = "accounts";
pub const CF_TRANSACTIONS: &str = "transactions";
pub const CF_METADATA: &str = "metadata";
pub const CF_SNAPSHOTS: &str = "snapshots";
pub const CF_IDENTITIES: &str = "identities";
pub const CF_DELEGATIONS: &str = "delegations";
pub const CF_CREDENTIALS: &str = "credentials";
pub const CF_CHANNELS: &str = "channels";
pub const CF_AGENTS: &str = "agents";
pub const CF_MODELS: &str = "models";
pub const CF_PROVIDERS: &str = "providers";
pub const CF_TASKS: &str = "tasks";
pub const CF_AGENT_TEMPLATES: &str = "agent_templates";
pub const CF_SKILLS: &str = "skills";
pub const CF_TOOLS: &str = "tools";
/// Operator-curated knowledge resource registry: vector DBs, RAG
/// indices, document corpora, indexed historical datasets, real-time
/// data feeds. Persists `KnowledgeRecord` rows. Mirrors `CF_TOOLS`
/// pattern but for queryable data resources.
pub const CF_KNOWLEDGE: &str = "knowledge";
/// Operator-curated workflow template catalog: reusable multi-step
/// saga specs that tenants can instantiate via `tenzro_instantiateWorkflow`.
/// Persists `WorkflowTemplate` rows. Distinct from running-workflow
/// state (which lives in tenzro-workflow's own storage).
pub const CF_WORKFLOW_TEMPLATES: &str = "workflow_templates";
pub const CF_TOKENS: &str = "tokens";
pub const CF_SETTLEMENTS: &str = "settlements";
pub const CF_MODEL_SERVICES: &str = "model_services";
pub const CF_NFTS: &str = "nfts";
pub const CF_EVENTS: &str = "events";
pub const CF_WEBHOOKS: &str = "webhooks";
pub const CF_COMPLIANCE: &str = "compliance";
/// Tenzro Train: in-flight and completed training runs (`TrainingRun` records).
pub const CF_TRAINING_RUNS: &str = "training_runs";
/// Tenzro Train: sealed training receipts at run finalization.
pub const CF_TRAINING_RECEIPTS: &str = "training_receipts";
/// Tenzro Auth: append-only audit log for token issuance, validation, revocation,
/// signing decisions, and HITL approval decisions. Keys: `audit:<ulid>`,
/// `audit_did:<did>:<ulid>`, `audit_jti:<jti>` → ULID of the issuance event.
pub const CF_AUDIT: &str = "audit";
/// Tenzro Auth: HITL approval state. Keys: `approval:<approval_id>` → JSON
/// `ApprovalRecord`, `approval_pending:<approver_did>:<approval_id>` →
/// empty value (secondary index for the approver's queue).
pub const CF_APPROVALS: &str = "approvals";
/// Tenzro API keys: per-client bearer credentials for gating proxied
/// services where the node mediates a credential the client does not
/// hold directly (e.g. Canton devnet JWT). Keys: `apikey:<sha256_hex>` →
/// JSON `ApiKeyRecord` (subject DID, scopes, created_at, revoked_at).
pub const CF_API_KEYS: &str = "api_keys";
/// Tenzro Bridge MPC: TEE-sealed DKLS23 threshold-ECDSA keyshares.
/// Keys: `mpc/keyshare/<group_id_hex>/<epoch_le_u64_hex>` → JSON
/// `KeyshareEnvelope` (sealed `Party<C>` ciphertext + public coordinates);
/// `mpc/keyshare/by_group/<group_id_hex>` → JSON `Vec<u64>` (epoch index).
pub const CF_MPC_KEYSHARES: &str = "mpc_keyshares";
/// Canton multi-tenant analytics. Per-API-key transaction counters,
/// last-seen timestamp, per-method call counts, and party usage —
/// scoped to each tenant so an RPC operator can answer "how many
/// DAML txes has the tenzro-labs team submitted this month".
/// Keys: `canton_analytics:<key_id>` → JSON `CantonKeyAnalytics`
/// (calls_total, calls_by_method, errors_total, last_called_at,
/// first_seen_at, canton_user_id, party_id_hint).
pub const CF_CANTON_ANALYTICS: &str = "canton_analytics";
/// Bridge multi-tenant analytics. Per-API-key call counters for
/// `chainlink`-scoped methods (Chainlink Data Feed reads via the
/// operator's paid Ethereum mainnet RPC) so the RPC operator can
/// attribute upstream costs to each tenant. Keys:
/// `bridge_analytics:<key_id>` → JSON `BridgeKeyAnalytics`
/// (calls_total, calls_by_method, errors_total, last_called_at,
/// first_seen_at, cu_consumed_total).
pub const CF_BRIDGE_ANALYTICS: &str = "bridge_analytics";
/// ERC-7579 validator modules state. Per-smart-account configuration
/// for SocialRecovery, SessionKey, SpendingLimit validators, plus
/// the installed-modules index. Custody enforcement at signing time
/// depends on this being durable across restarts.
/// Keys:
///   `erc7579/social/<account20>` → JSON `SocialRecoveryConfig`
///   `erc7579/session/<account20>` → JSON `SessionKeyConfig`
///   `erc7579/session_state/<account20>` → JSON `SessionKeyState`
///   `erc7579/spending/<account20>` → JSON `SpendingLimitConfig`
///   `erc7579/spending_state/<account20>` → JSON `SpendingLimitState`
///   `erc7579/installed/<account20>/<validator20>` → JSON
///       `ValidatorModuleConfig` (installed-module index for the
///       AND-combined gate).
pub const CF_VALIDATOR_MODULES: &str = "validator_modules";

/// Key-value store trait
pub trait KvStore: Send + Sync {
    /// Gets a value by key from a column family
    fn get(&self, cf: &str, key: &[u8]) -> Result<Option<Vec<u8>>>;

    /// Puts a key-value pair into a column family
    fn put(&self, cf: &str, key: &[u8], value: &[u8]) -> Result<()>;

    /// Deletes a key from a column family
    fn delete(&self, cf: &str, key: &[u8]) -> Result<()>;

    /// Checks if a key exists in a column family
    fn contains(&self, cf: &str, key: &[u8]) -> Result<bool> {
        Ok(self.get(cf, key)?.is_some())
    }

    /// Batch write operation (buffered, not guaranteed durable until WAL flush)
    fn write_batch(&self, operations: Vec<WriteOp>) -> Result<()>;

    /// Batch write operation with fsync for durability.
    /// Use this for finalized blocks and other critical data that must survive power loss.
    fn write_batch_sync(&self, operations: Vec<WriteOp>) -> Result<()>;

    /// Gets all keys with a given prefix
    fn get_keys_with_prefix(&self, cf: &str, prefix: &[u8]) -> Result<Vec<Vec<u8>>>;

    /// Scans all key-value pairs with a given prefix
    fn scan_prefix(&self, cf: &str, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
        let keys = self.get_keys_with_prefix(cf, prefix)?;
        let mut results = Vec::with_capacity(keys.len());
        for key in keys {
            if let Some(value) = self.get(cf, &key)? {
                results.push((key, value));
            }
        }
        Ok(results)
    }

    /// Streams every key-value pair with the given prefix to `f` without
    /// materializing the full result set. Use this instead of
    /// [`KvStore::scan_prefix`] whenever the prefix can match an unbounded
    /// number of rows (e.g. a whole-column-family walk) — `scan_prefix`
    /// builds one `Vec` holding every row, which for CF_BLOCKS on a
    /// long-running chain is a multi-gigabyte allocation. An error returned
    /// from `f` aborts the scan and propagates.
    fn scan_prefix_for_each(
        &self,
        cf: &str,
        prefix: &[u8],
        f: &mut dyn FnMut(&[u8], &[u8]) -> Result<()>,
    ) -> Result<()> {
        for (key, value) in self.scan_prefix(cf, prefix)? {
            f(&key, &value)?;
        }
        Ok(())
    }
}

/// Write operation for batch writes
#[derive(Debug, Clone)]
pub enum WriteOp {
    /// Put operation
    Put {
        cf: String,
        key: Vec<u8>,
        value: Vec<u8>,
    },
    /// Delete operation
    Delete { cf: String, key: Vec<u8> },
}

/// RocksDB-based key-value store
pub struct RocksDbStore {
    db: Arc<DB>,
}

impl RocksDbStore {
    /// Creates the column family descriptors for all storage column families
    fn column_family_descriptors() -> Vec<ColumnFamilyDescriptor> {
        vec![
            ColumnFamilyDescriptor::new(CF_BLOCKS, Options::default()),
            ColumnFamilyDescriptor::new(CF_STATE, Options::default()),
            ColumnFamilyDescriptor::new(CF_ACCOUNTS, Options::default()),
            ColumnFamilyDescriptor::new(CF_TRANSACTIONS, Options::default()),
            ColumnFamilyDescriptor::new(CF_METADATA, Options::default()),
            ColumnFamilyDescriptor::new(CF_SNAPSHOTS, Options::default()),
            ColumnFamilyDescriptor::new(CF_IDENTITIES, Options::default()),
            ColumnFamilyDescriptor::new(CF_DELEGATIONS, Options::default()),
            ColumnFamilyDescriptor::new(CF_CREDENTIALS, Options::default()),
            ColumnFamilyDescriptor::new(CF_CHANNELS, Options::default()),
            ColumnFamilyDescriptor::new(CF_AGENTS, Options::default()),
            ColumnFamilyDescriptor::new(CF_MODELS, Options::default()),
            ColumnFamilyDescriptor::new(CF_PROVIDERS, Options::default()),
            ColumnFamilyDescriptor::new(CF_TASKS, Options::default()),
            ColumnFamilyDescriptor::new(CF_AGENT_TEMPLATES, Options::default()),
            ColumnFamilyDescriptor::new(CF_SKILLS, Options::default()),
            ColumnFamilyDescriptor::new(CF_TOOLS, Options::default()),
            ColumnFamilyDescriptor::new(CF_KNOWLEDGE, Options::default()),
            ColumnFamilyDescriptor::new(CF_WORKFLOW_TEMPLATES, Options::default()),
            ColumnFamilyDescriptor::new(CF_TOKENS, Options::default()),
            ColumnFamilyDescriptor::new(CF_SETTLEMENTS, Options::default()),
            ColumnFamilyDescriptor::new(CF_MODEL_SERVICES, Options::default()),
            ColumnFamilyDescriptor::new(CF_NFTS, Options::default()),
            ColumnFamilyDescriptor::new(CF_EVENTS, Options::default()),
            ColumnFamilyDescriptor::new(CF_WEBHOOKS, Options::default()),
            ColumnFamilyDescriptor::new(CF_COMPLIANCE, Options::default()),
            ColumnFamilyDescriptor::new(CF_TRAINING_RUNS, Options::default()),
            ColumnFamilyDescriptor::new(CF_TRAINING_RECEIPTS, Options::default()),
            ColumnFamilyDescriptor::new(CF_AUDIT, Options::default()),
            ColumnFamilyDescriptor::new(CF_APPROVALS, Options::default()),
            ColumnFamilyDescriptor::new(CF_API_KEYS, Options::default()),
            ColumnFamilyDescriptor::new(CF_MPC_KEYSHARES, Options::default()),
            ColumnFamilyDescriptor::new(CF_CANTON_ANALYTICS, Options::default()),
            ColumnFamilyDescriptor::new(CF_BRIDGE_ANALYTICS, Options::default()),
            ColumnFamilyDescriptor::new(CF_VALIDATOR_MODULES, Options::default()),
        ]
    }

    /// Opens a RocksDB store with the given configuration
    pub fn open(config: &StorageConfig) -> Result<Self> {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);
        opts.set_max_open_files(config.max_open_files);
        opts.set_write_buffer_size(config.write_buffer_size);
        opts.set_max_write_buffer_number(config.max_write_buffer_number);
        opts.set_target_file_size_base(config.target_file_size_base);

        if config.compression {
            opts.set_compression_type(rocksdb::DBCompressionType::Lz4);
        }

        if config.enable_statistics {
            opts.enable_statistics();
        }

        // Attempt to open the database, handling WAL corruption gracefully
        let db = match DB::open_cf_descriptors(&opts, &config.db_path, Self::column_family_descriptors()) {
            Ok(db) => db,
            Err(e) => {
                let error_str = e.to_string();
                // Check if the error indicates WAL corruption
                if error_str.contains("Corruption") || error_str.contains("corruption") ||
                   error_str.contains("log") || error_str.contains("WAL") {
                    tracing::warn!(
                        "Database corruption detected at {:?}, attempting repair: {}",
                        config.db_path,
                        error_str
                    );

                    // Attempt to repair the database
                    if let Err(repair_err) = Self::repair_database(&config.db_path) {
                        tracing::error!("Database repair failed: {}", repair_err);
                        return Err(StorageError::DatabaseError(format!(
                            "Failed to open database and repair failed: {} (repair error: {})",
                            error_str,
                            repair_err
                        )));
                    }

                    tracing::info!("Database repair completed, reopening...");

                    // Try opening again after repair with fresh descriptors
                    DB::open_cf_descriptors(&opts, &config.db_path, Self::column_family_descriptors())
                        .map_err(|e| StorageError::DatabaseError(format!(
                            "Failed to open database after repair: {}",
                            e
                        )))?
                } else {
                    // Not a corruption error, propagate it
                    return Err(e.into());
                }
            }
        };

        Ok(Self { db: Arc::new(db) })
    }

    /// Attempts to repair a corrupted RocksDB database.
    ///
    /// WARNING: This may result in data loss. Only finalized blocks should survive.
    /// Recent uncommitted data in the WAL may be lost.
    pub fn repair_database<P: AsRef<Path>>(path: P) -> Result<()> {
        tracing::warn!(
            "Repairing database at {:?}. This may result in loss of recent uncommitted data.",
            path.as_ref()
        );

        DB::repair(&Options::default(), path.as_ref())
            .map_err(|e| StorageError::DatabaseError(format!("Repair failed: {}", e)))?;

        tracing::info!("Database repair completed successfully");
        Ok(())
    }

    /// Opens a RocksDB store at the given path with default options
    pub fn open_default<P: AsRef<Path>>(path: P) -> Result<Self> {
        let config = StorageConfig::new(path.as_ref().to_path_buf());
        Self::open(&config)
    }

    /// Gets a column family handle
    fn cf_handle(&self, name: &str) -> Result<&ColumnFamily> {
        self.db
            .cf_handle(name)
            .ok_or_else(|| StorageError::ColumnFamilyNotFound(name.to_string()))
    }

    /// Gets a reference to the underlying RocksDB instance
    /// This is useful for advanced operations like creating checkpoints
    pub fn db(&self) -> &DB {
        &self.db
    }
}

impl KvStore for RocksDbStore {
    fn get(&self, cf: &str, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let cf_handle = self.cf_handle(cf)?;
        Ok(self.db.get_cf(cf_handle, key)?)
    }

    fn put(&self, cf: &str, key: &[u8], value: &[u8]) -> Result<()> {
        let cf_handle = self.cf_handle(cf)?;
        Ok(self.db.put_cf(cf_handle, key, value)?)
    }

    fn delete(&self, cf: &str, key: &[u8]) -> Result<()> {
        let cf_handle = self.cf_handle(cf)?;
        Ok(self.db.delete_cf(cf_handle, key)?)
    }

    fn write_batch(&self, operations: Vec<WriteOp>) -> Result<()> {
        let mut batch = WriteBatch::default();

        for op in operations {
            match op {
                WriteOp::Put { cf, key, value } => {
                    let cf_handle = self.cf_handle(&cf)?;
                    batch.put_cf(cf_handle, key, value);
                }
                WriteOp::Delete { cf, key } => {
                    let cf_handle = self.cf_handle(&cf)?;
                    batch.delete_cf(cf_handle, key);
                }
            }
        }

        Ok(self.db.write(batch)?)
    }

    fn write_batch_sync(&self, operations: Vec<WriteOp>) -> Result<()> {
        let mut batch = WriteBatch::default();

        for op in operations {
            match op {
                WriteOp::Put { cf, key, value } => {
                    let cf_handle = self.cf_handle(&cf)?;
                    batch.put_cf(cf_handle, key, value);
                }
                WriteOp::Delete { cf, key } => {
                    let cf_handle = self.cf_handle(&cf)?;
                    batch.delete_cf(cf_handle, key);
                }
            }
        }

        // Force fsync of the WAL so finalized blocks survive power loss.
        // This is the whole point of `write_batch_sync` — without `set_sync(true)`,
        // RocksDB only writes to the OS page cache and the data is lost on a
        // hard reboot. The previous `set_sync(false)` made this method
        // indistinguishable from `write_batch` (HIGH #78 in the production audit).
        let mut write_opts = WriteOptions::default();
        write_opts.set_sync(true);
        Ok(self.db.write_opt(batch, &write_opts)?)
    }

    fn get_keys_with_prefix(&self, cf: &str, prefix: &[u8]) -> Result<Vec<Vec<u8>>> {
        let cf_handle = self.cf_handle(cf)?;
        let mut keys = Vec::new();
        let iter = self.db.prefix_iterator_cf(cf_handle, prefix);

        for item in iter {
            let (key, _) = item?;
            if key.starts_with(prefix) {
                keys.push(key.to_vec());
            } else {
                break;
            }
        }

        Ok(keys)
    }

    /// Single-pass prefix scan that pulls keys *and* values in one iterator
    /// traversal. The default trait impl issues N follow-up `get_cf` calls;
    /// for hot paths like `BlockStorage::blocks_by_height_range` (catch-up
    /// sync) that is O(N) extra lookups per range. RocksDB already returns
    /// the value alongside the key inside `prefix_iterator_cf`, so we just
    /// take both.
    fn scan_prefix(&self, cf: &str, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
        let cf_handle = self.cf_handle(cf)?;
        let mut results = Vec::new();
        let iter = self.db.prefix_iterator_cf(cf_handle, prefix);

        for item in iter {
            let (key, value) = item?;
            if key.starts_with(prefix) {
                results.push((key.to_vec(), value.to_vec()));
            } else {
                break;
            }
        }

        Ok(results)
    }

    /// Streaming prefix scan — borrows each (key, value) directly from the
    /// RocksDB iterator and never accumulates rows, so peak memory is one
    /// row regardless of how many rows the prefix matches. This is the
    /// path the snapshot producer uses to walk entire column families.
    fn scan_prefix_for_each(
        &self,
        cf: &str,
        prefix: &[u8],
        f: &mut dyn FnMut(&[u8], &[u8]) -> Result<()>,
    ) -> Result<()> {
        let cf_handle = self.cf_handle(cf)?;
        let iter = self.db.prefix_iterator_cf(cf_handle, prefix);

        for item in iter {
            let (key, value) = item?;
            if !key.starts_with(prefix) {
                break;
            }
            f(&key, &value)?;
        }

        Ok(())
    }
}

/// In-memory key-value store for testing
pub struct MemoryStore {
    data: Arc<RwLock<HashMap<String, HashMap<Vec<u8>, Vec<u8>>>>>,
}

impl MemoryStore {
    /// Creates a new in-memory store
    pub fn new() -> Self {
        let mut data = HashMap::new();
        // Initialize column families
        data.insert(CF_BLOCKS.to_string(), HashMap::new());
        data.insert(CF_STATE.to_string(), HashMap::new());
        data.insert(CF_ACCOUNTS.to_string(), HashMap::new());
        data.insert(CF_TRANSACTIONS.to_string(), HashMap::new());
        data.insert(CF_METADATA.to_string(), HashMap::new());
        data.insert(CF_SNAPSHOTS.to_string(), HashMap::new());
        data.insert(CF_IDENTITIES.to_string(), HashMap::new());
        data.insert(CF_DELEGATIONS.to_string(), HashMap::new());
        data.insert(CF_CREDENTIALS.to_string(), HashMap::new());
        data.insert(CF_CHANNELS.to_string(), HashMap::new());
        data.insert(CF_AGENTS.to_string(), HashMap::new());
        data.insert(CF_MODELS.to_string(), HashMap::new());
        data.insert(CF_PROVIDERS.to_string(), HashMap::new());
        data.insert(CF_TASKS.to_string(), HashMap::new());
        data.insert(CF_AGENT_TEMPLATES.to_string(), HashMap::new());
        data.insert(CF_SKILLS.to_string(), HashMap::new());
        data.insert(CF_TOOLS.to_string(), HashMap::new());
        data.insert(CF_TOKENS.to_string(), HashMap::new());
        data.insert(CF_SETTLEMENTS.to_string(), HashMap::new());
        data.insert(CF_MODEL_SERVICES.to_string(), HashMap::new());
        data.insert(CF_NFTS.to_string(), HashMap::new());
        data.insert(CF_EVENTS.to_string(), HashMap::new());
        data.insert(CF_WEBHOOKS.to_string(), HashMap::new());
        data.insert(CF_COMPLIANCE.to_string(), HashMap::new());
        data.insert(CF_TRAINING_RUNS.to_string(), HashMap::new());
        data.insert(CF_TRAINING_RECEIPTS.to_string(), HashMap::new());
        data.insert(CF_AUDIT.to_string(), HashMap::new());
        data.insert(CF_APPROVALS.to_string(), HashMap::new());

        Self {
            data: Arc::new(RwLock::new(data)),
        }
    }
}

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

impl KvStore for MemoryStore {
    fn get(&self, cf: &str, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let data = self.data.read();
        Ok(data
            .get(cf)
            .and_then(|cf_data| cf_data.get(key))
            .cloned())
    }

    fn put(&self, cf: &str, key: &[u8], value: &[u8]) -> Result<()> {
        let mut data = self.data.write();
        data.entry(cf.to_string())
            .or_default()
            .insert(key.to_vec(), value.to_vec());
        Ok(())
    }

    fn delete(&self, cf: &str, key: &[u8]) -> Result<()> {
        let mut data = self.data.write();
        if let Some(cf_data) = data.get_mut(cf) {
            cf_data.remove(key);
        }
        Ok(())
    }

    fn write_batch(&self, operations: Vec<WriteOp>) -> Result<()> {
        let mut data = self.data.write();
        for op in operations {
            match op {
                WriteOp::Put { cf, key, value } => {
                    data.entry(cf)
                        .or_default()
                        .insert(key, value);
                }
                WriteOp::Delete { cf, key } => {
                    if let Some(cf_data) = data.get_mut(&cf) {
                        cf_data.remove(&key);
                    }
                }
            }
        }
        Ok(())
    }

    fn write_batch_sync(&self, operations: Vec<WriteOp>) -> Result<()> {
        // In-memory store has no persistence; sync is a no-op
        self.write_batch(operations)
    }

    fn get_keys_with_prefix(&self, cf: &str, prefix: &[u8]) -> Result<Vec<Vec<u8>>> {
        let data = self.data.read();
        let keys = data
            .get(cf)
            .map(|cf_data| {
                cf_data
                    .keys()
                    .filter(|k| k.starts_with(prefix))
                    .cloned()
                    .collect()
            })
            .unwrap_or_default();
        Ok(keys)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_memory_store() {
        let store = MemoryStore::new();

        // Test put and get
        store.put(CF_STATE, b"key1", b"value1").unwrap();
        let value = store.get(CF_STATE, b"key1").unwrap();
        assert_eq!(value, Some(b"value1".to_vec()));

        // Test delete
        store.delete(CF_STATE, b"key1").unwrap();
        let value = store.get(CF_STATE, b"key1").unwrap();
        assert_eq!(value, None);

        // Test batch write
        let ops = vec![
            WriteOp::Put {
                cf: CF_STATE.to_string(),
                key: b"key2".to_vec(),
                value: b"value2".to_vec(),
            },
            WriteOp::Put {
                cf: CF_STATE.to_string(),
                key: b"key3".to_vec(),
                value: b"value3".to_vec(),
            },
        ];
        store.write_batch(ops).unwrap();

        let value2 = store.get(CF_STATE, b"key2").unwrap();
        let value3 = store.get(CF_STATE, b"key3").unwrap();
        assert_eq!(value2, Some(b"value2".to_vec()));
        assert_eq!(value3, Some(b"value3".to_vec()));
    }

    #[test]
    fn test_prefix_search() {
        let store = MemoryStore::new();

        store.put(CF_STATE, b"prefix_key1", b"value1").unwrap();
        store.put(CF_STATE, b"prefix_key2", b"value2").unwrap();
        store.put(CF_STATE, b"other_key", b"value3").unwrap();

        let keys = store.get_keys_with_prefix(CF_STATE, b"prefix_").unwrap();
        assert_eq!(keys.len(), 2);
    }

    /// Returns a unique temporary directory path for a RocksDB test instance.
    /// We avoid pulling in a `tempfile` dev-dependency by composing the path
    /// from the system temp dir, the process id, and a per-call counter.
    fn unique_temp_db_path(label: &str) -> std::path::PathBuf {
        use std::sync::atomic::{AtomicU64, Ordering};
        static COUNTER: AtomicU64 = AtomicU64::new(0);
        let id = COUNTER.fetch_add(1, Ordering::SeqCst);
        let pid = std::process::id();
        let nanos = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or(0);
        std::env::temp_dir().join(format!("tenzro-storage-test-{}-{}-{}-{}", label, pid, id, nanos))
    }

    #[test]
    fn test_rocksdb_write_batch_sync_durability() {
        // Verifies that `write_batch_sync` actually persists data and that
        // we can reopen the DB and read what we wrote. This is the regression
        // test for HIGH #78 — the previous implementation called
        // `write_opts.set_sync(false)`, which made the API silently
        // non-durable.
        let path = unique_temp_db_path("sync");

        {
            let store = RocksDbStore::open_default(&path).expect("open db");
            store
                .write_batch_sync(vec![WriteOp::Put {
                    cf: CF_STATE.to_string(),
                    key: b"durable-key".to_vec(),
                    value: b"durable-value".to_vec(),
                }])
                .expect("sync write");
            // Drop closes the DB cleanly.
        }

        {
            let store = RocksDbStore::open_default(&path).expect("reopen db");
            let value = store.get(CF_STATE, b"durable-key").expect("read");
            assert_eq!(value, Some(b"durable-value".to_vec()));
        }

        // Best-effort cleanup; ignore failures.
        let _ = std::fs::remove_dir_all(&path);
    }

    #[test]
    fn test_rocksdb_repair_recovers_from_truncated_wal() {
        // Verifies WAL recovery (HIGH #78). We open a DB, write some data,
        // close it, corrupt one of the SST or log files by truncating it,
        // then call `repair_database` and confirm we can reopen the DB.
        // RocksDB repair is best-effort: finalized writes generally survive,
        // unflushed WAL entries may be lost. We therefore only assert that
        // (a) repair succeeds and (b) the DB reopens.
        let path = unique_temp_db_path("repair");

        {
            let store = RocksDbStore::open_default(&path).expect("open db");
            for i in 0..32u64 {
                store
                    .write_batch_sync(vec![WriteOp::Put {
                        cf: CF_STATE.to_string(),
                        key: format!("key-{}", i).into_bytes(),
                        value: format!("value-{}", i).into_bytes(),
                    }])
                    .expect("sync write");
            }
        }

        // Truncate every .log file in the DB directory to simulate WAL corruption.
        if let Ok(entries) = std::fs::read_dir(&path) {
            for entry in entries.flatten() {
                let p = entry.path();
                if p.extension().and_then(|s| s.to_str()) == Some("log")
                    && let Ok(meta) = std::fs::metadata(&p)
                {
                    let new_len = meta.len() / 2;
                    if let Ok(file) = std::fs::OpenOptions::new().write(true).open(&p) {
                        let _ = file.set_len(new_len);
                    }
                }
            }
        }

        // Repair should not panic; it may log warnings/errors but should return Ok.
        let repair_result = RocksDbStore::repair_database(&path);
        assert!(
            repair_result.is_ok(),
            "repair_database should succeed on truncated WAL: {:?}",
            repair_result
        );

        // After repair we must be able to reopen the database.
        let store = RocksDbStore::open_default(&path).expect("reopen after repair");

        // We don't assert which specific keys survive — RocksDB repair is
        // not transactional w.r.t. unflushed writes. We just confirm the DB
        // is openable and queryable, which is what matters for crash recovery.
        let _ = store.get(CF_STATE, b"key-0");

        let _ = std::fs::remove_dir_all(&path);
    }

    #[test]
    fn test_rocksdb_open_recovers_from_wal_corruption() {
        // Verifies the auto-repair-on-open path in `RocksDbStore::open()`.
        // We write data, truncate the WAL, then re-open. The open path
        // should detect corruption, run repair, and reopen successfully.
        let path = unique_temp_db_path("auto-repair");

        {
            let store = RocksDbStore::open_default(&path).expect("open db");
            for i in 0..16u64 {
                store
                    .write_batch_sync(vec![WriteOp::Put {
                        cf: CF_BLOCKS.to_string(),
                        key: format!("block-{}", i).into_bytes(),
                        value: format!("payload-{}", i).into_bytes(),
                    }])
                    .expect("sync write");
            }
        }

        if let Ok(entries) = std::fs::read_dir(&path) {
            for entry in entries.flatten() {
                let p = entry.path();
                if p.extension().and_then(|s| s.to_str()) == Some("log")
                    && let Ok(meta) = std::fs::metadata(&p)
                {
                    let new_len = meta.len() / 3;
                    if let Ok(file) = std::fs::OpenOptions::new().write(true).open(&p) {
                        let _ = file.set_len(new_len);
                    }
                }
            }
        }

        // RocksDB does not always classify a truncated WAL as a "Corruption"
        // error — it may simply silently truncate the recovered tail. So
        // this test only asserts that the open path does not panic on a
        // damaged WAL. The auto-repair branch is exercised for actual
        // corruption errors at runtime in production.
        let _ = RocksDbStore::open_default(&path);

        let _ = std::fs::remove_dir_all(&path);
    }
}