crates/bex-db/src/lib.rs

use anyhow::{Context, Result};
use bex_types::plugin_info::PluginInfo;
use redb::{Database, ReadableTable, TableDefinition};
use std::path::Path;
use std::sync::Arc;

const KV_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("kv");
const SECRETS_TABLE: TableDefinition<&str, &str> = TableDefinition::new("secrets");
const PLUGINS_TABLE: TableDefinition<&str, &str> = TableDefinition::new("plugins");
const WASM_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("wasm_blobs");
const MANIFEST_TABLE: TableDefinition<&str, &str> = TableDefinition::new("manifests");

fn kv_key(pid: &str, key: &str) -> String {
    format!("{}\x00{}", pid, key)
}

pub struct BexDb {
    db: Arc<Database>,
}

impl BexDb {
    pub fn open(data_dir: &Path) -> Result<Self> {
        std::fs::create_dir_all(data_dir)?;
        let db_path = data_dir.join("bex.redb");

        // Compact on open if DB file is > 10 MB (before Arc wrapping,
        // since compact() requires &mut Database)
        let should_compact = std::fs::metadata(&db_path)
            .map(|m| m.len() > 10 * 1024 * 1024)
            .unwrap_or(false);

        let mut db = Database::create(&db_path).context("failed to open redb")?;

        if should_compact {
            tracing::info!("Compacting database (file > 10 MB)...");
            if let Err(e) = db.compact() {
                tracing::warn!("Database compaction failed: {}", e);
            }
        }

        let db = Arc::new(db);
        let write_txn = db.begin_write()?;
        write_txn.open_table(KV_TABLE)?;
        write_txn.open_table(SECRETS_TABLE)?;
        write_txn.open_table(PLUGINS_TABLE)?;
        write_txn.open_table(WASM_TABLE)?;
        write_txn.open_table(MANIFEST_TABLE)?;
        write_txn.commit()?;
        Ok(Self { db })
    }

    pub fn kv_set(&self, plugin_id: &str, key: &str, value: &[u8]) -> Result<bool> {
        let k = kv_key(plugin_id, key);
        let write_txn = self.db.begin_write()?;
        let mut table = write_txn.open_table(KV_TABLE)?;
        table.insert(k.as_str(), value)?;
        drop(table);
        write_txn.commit()?;
        Ok(true)
    }

    pub fn kv_get(&self, plugin_id: &str, key: &str) -> Result<Option<Vec<u8>>> {
        let k = kv_key(plugin_id, key);
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(KV_TABLE)?;
        let result = table.get(k.as_str())?.map(|v| v.value().to_vec());
        Ok(result)
    }

    pub fn kv_remove(&self, plugin_id: &str, key: &str) -> Result<bool> {
        let k = kv_key(plugin_id, key);
        let write_txn = self.db.begin_write()?;
        let mut table = write_txn.open_table(KV_TABLE)?;
        let existed = table.remove(k.as_str())?.is_some();
        drop(table);
        write_txn.commit()?;
        Ok(existed)
    }

    pub fn kv_keys(&self, plugin_id: &str, prefix: &str) -> Result<Vec<String>> {
        let start = kv_key(plugin_id, prefix);
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(KV_TABLE)?;
        let range = table.range(start.as_str()..)?;
        let mut result = Vec::new();
        for item in range {
            let (k, _) = item?;
            let key_str = k.value();
            if let Some(pos) = key_str.find('\0') {
                let pid = &key_str[..pos];
                let key = &key_str[pos + 1..];
                if pid == plugin_id && key.starts_with(prefix) {
                    result.push(key.to_string());
                }
            }
        }
        Ok(result)
    }

    pub fn secret_set(&self, plugin_id: &str, key: &str, value: &str) -> Result<()> {
        let k = kv_key(plugin_id, key);
        let write_txn = self.db.begin_write()?;
        let mut table = write_txn.open_table(SECRETS_TABLE)?;
        table.insert(k.as_str(), value)?;
        drop(table);
        write_txn.commit()?;
        Ok(())
    }

    pub fn secret_get(&self, plugin_id: &str, key: &str) -> Result<Option<String>> {
        let k = kv_key(plugin_id, key);
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(SECRETS_TABLE)?;
        let result = table.get(k.as_str())?.map(|v| v.value().to_string());
        Ok(result)
    }

    pub fn secret_remove(&self, plugin_id: &str, key: &str) -> Result<bool> {
        let k = kv_key(plugin_id, key);
        let write_txn = self.db.begin_write()?;
        let mut table = write_txn.open_table(SECRETS_TABLE)?;
        let existed = table.remove(k.as_str())?.is_some();
        drop(table);
        write_txn.commit()?;
        Ok(existed)
    }

    pub fn secret_keys(&self, plugin_id: &str, prefix: &str) -> Result<Vec<String>> {
        let start = kv_key(plugin_id, prefix);
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(SECRETS_TABLE)?;
        let range = table.range(start.as_str()..)?;
        let mut result = Vec::new();
        for item in range {
            let (k, _) = item?;
            let key_str = k.value();
            if let Some(pos) = key_str.find('\0') {
                let pid = &key_str[..pos];
                let key = &key_str[pos + 1..];
                if pid == plugin_id && key.starts_with(prefix) {
                    result.push(key.to_string());
                }
            }
        }
        Ok(result)
    }

    pub fn save_plugin_info(&self, info: &PluginInfo) -> Result<()> {
        let json = serde_json::to_string(info)?;
        let write_txn = self.db.begin_write()?;
        let mut table = write_txn.open_table(PLUGINS_TABLE)?;
        table.insert(info.id.as_str(), json.as_str())?;
        drop(table);
        write_txn.commit()?;
        Ok(())
    }

    pub fn get_plugin_info(&self, id: &str) -> Result<Option<PluginInfo>> {
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(PLUGINS_TABLE)?;
        let result = match table.get(id)? {
            Some(json) => Some(serde_json::from_str::<PluginInfo>(json.value())?),
            None => None,
        };
        Ok(result)
    }

    pub fn list_plugins(&self) -> Result<Vec<PluginInfo>> {
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(PLUGINS_TABLE)?;
        let mut result = Vec::new();
        for item in table.iter()? {
            let (_, json) = item?;
            if let Ok(info) = serde_json::from_str::<PluginInfo>(json.value()) {
                result.push(info);
            }
        }
        Ok(result)
    }

    pub fn remove_plugin(&self, id: &str) -> Result<bool> {
        let write_txn = self.db.begin_write()?;
        let mut table = write_txn.open_table(PLUGINS_TABLE)?;
        let existed = table.remove(id)?.is_some();
        drop(table);
        // Also remove WASM blob and manifest
        let mut wasm_table = write_txn.open_table(WASM_TABLE)?;
        wasm_table.remove(id)?;
        drop(wasm_table);
        let mut manifest_table = write_txn.open_table(MANIFEST_TABLE)?;
        manifest_table.remove(id)?;
        drop(manifest_table);
        write_txn.commit()?;
        Ok(existed)
    }

    /// Store the WASM blob for a plugin
    pub fn save_wasm_blob(&self, id: &str, wasm: &[u8]) -> Result<()> {
        let write_txn = self.db.begin_write()?;
        let mut table = write_txn.open_table(WASM_TABLE)?;
        table.insert(id, wasm)?;
        drop(table);
        write_txn.commit()?;
        Ok(())
    }

    /// Retrieve the WASM blob for a plugin
    pub fn get_wasm_blob(&self, id: &str) -> Result<Option<Vec<u8>>> {
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(WASM_TABLE)?;
        Ok(table.get(id)?.map(|v| v.value().to_vec()))
    }

    /// Store the manifest YAML for a plugin
    pub fn save_manifest(&self, id: &str, manifest_yaml: &str) -> Result<()> {
        let write_txn = self.db.begin_write()?;
        let mut table = write_txn.open_table(MANIFEST_TABLE)?;
        table.insert(id, manifest_yaml)?;
        drop(table);
        write_txn.commit()?;
        Ok(())
    }

    /// Retrieve the manifest YAML for a plugin
    pub fn get_manifest(&self, id: &str) -> Result<Option<String>> {
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(MANIFEST_TABLE)?;
        Ok(table.get(id)?.map(|v| v.value().to_string()))
    }
}