rsclaw 0.0.1-alpha.1

use anyhow::{Context, Result};
use redb::{Database, ReadableTable, TableDefinition, WriteTransaction};
use std::path::{Path, PathBuf};
use std::sync::Arc;

/// Table definitions for redb.
const SESSIONS_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("sessions");
const AGENTS_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("agents");
const MESSAGES_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("messages");
const METADATA_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("metadata");

/// Main database for persistent storage.
pub struct Store {
    db: Arc<Database>,
    data_dir: PathBuf,
}

impl Store {
    /// Create or open a store at the given path.
    pub fn new(data_dir: &Path) -> Result<Self> {
        std::fs::create_dir_all(data_dir).context("Failed to create data directory")?;

        let db_path = data_dir.join("rsclaw.redb");
        let db = Database::create(&db_path).context("Failed to open database")?;

        let store = Self {
            db: Arc::new(db),
            data_dir: data_dir.to_path_buf(),
        };

        store.init_tables()?;
        Ok(store)
    }

    /// Initialize required tables.
    fn init_tables(&self) -> Result<()> {
        let write_txn = self
            .db
            .begin_write()
            .context("Failed to begin write transaction")?;

        write_txn
            .open_table(SESSIONS_TABLE)
            .context("Failed to open sessions table")?;
        write_txn
            .open_table(AGENTS_TABLE)
            .context("Failed to open agents table")?;
        write_txn
            .open_table(MESSAGES_TABLE)
            .context("Failed to open messages table")?;
        write_txn
            .open_table(METADATA_TABLE)
            .context("Failed to open metadata table")?;

        write_txn.commit().context("Failed to commit transaction")?;

        Ok(())
    }

    /// Get data directory path.
    pub fn data_dir(&self) -> &Path {
        &self.data_dir
    }

    /// Get the database reference.
    pub fn db(&self) -> &Database {
        &self.db
    }
}

/// Key-value store operations.
impl Store {
    /// Put a value in a table.
    pub fn put(&self, table: TableDefinition<&str, &[u8]>, key: &str, value: &[u8]) -> Result<()> {
        let write_txn = self.db.begin_write()?;
        {
            let mut table = write_txn.open_table(table)?;
            table.insert(key, value)?;
        }
        write_txn.commit()?;
        Ok(())
    }

    /// Get a value from a table.
    pub fn get(&self, table: TableDefinition<&str, &[u8]>, key: &str) -> Result<Option<Vec<u8>>> {
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(table)?;

        match table.get(key)? {
            Some(value) => Ok(Some(value.value().to_vec())),
            None => Ok(None),
        }
    }

    /// Delete a value from a table.
    pub fn delete(&self, table: TableDefinition<&str, &[u8]>, key: &str) -> Result<bool> {
        let write_txn = self.db.begin_write()?;
        let existed = {
            let mut table = write_txn.open_table(table)?;
            table.remove(key)?.is_some()
        };
        write_txn.commit()?;
        Ok(existed)
    }

    /// List all keys in a table.
    pub fn list_keys(&self, table: TableDefinition<&str, &[u8]>) -> Result<Vec<String>> {
        let read_txn = self.db.begin_read()?;
        let table = read_txn.open_table(table)?;

        let mut keys = Vec::new();
        for item in table.iter()? {
            let (key, _) = item?;
            keys.push(key.value().to_string());
        }
        Ok(keys)
    }
}

/// Session operations.
impl Store {
    /// Save a session.
    pub fn save_session(&self, session_id: &str, data: &[u8]) -> Result<()> {
        self.put(SESSIONS_TABLE, session_id, data)
    }

    /// Get a session.
    pub fn get_session(&self, session_id: &str) -> Result<Option<Vec<u8>>> {
        self.get(SESSIONS_TABLE, session_id)
    }

    /// Delete a session.
    pub fn delete_session(&self, session_id: &str) -> Result<bool> {
        self.delete(SESSIONS_TABLE, session_id)
    }

    /// List all session IDs.
    pub fn list_sessions(&self) -> Result<Vec<String>> {
        self.list_keys(SESSIONS_TABLE)
    }
}

/// Agent operations.
impl Store {
    /// Save an agent state.
    pub fn save_agent(&self, agent_name: &str, data: &[u8]) -> Result<()> {
        self.put(AGENTS_TABLE, agent_name, data)
    }

    /// Get an agent state.
    pub fn get_agent(&self, agent_name: &str) -> Result<Option<Vec<u8>>> {
        self.get(AGENTS_TABLE, agent_name)
    }

    /// Delete an agent state.
    pub fn delete_agent(&self, agent_name: &str) -> Result<bool> {
        self.delete(AGENTS_TABLE, agent_name)
    }

    /// List all agent names.
    pub fn list_agents(&self) -> Result<Vec<String>> {
        self.list_keys(AGENTS_TABLE)
    }
}

/// Message operations.
impl Store {
    /// Save a message.
    pub fn save_message(&self, message_id: &str, data: &[u8]) -> Result<()> {
        self.put(MESSAGES_TABLE, message_id, data)
    }

    /// Get a message.
    pub fn get_message(&self, message_id: &str) -> Result<Option<Vec<u8>>> {
        self.get(MESSAGES_TABLE, message_id)
    }

    /// Delete a message.
    pub fn delete_message(&self, message_id: &str) -> Result<bool> {
        self.delete(MESSAGES_TABLE, message_id)
    }

    /// List all message IDs.
    pub fn list_messages(&self) -> Result<Vec<String>> {
        self.list_keys(MESSAGES_TABLE)
    }
}

/// Metadata operations.
impl Store {
    /// Set metadata.
    pub fn set_metadata(&self, key: &str, value: &[u8]) -> Result<()> {
        self.put(METADATA_TABLE, key, value)
    }

    /// Get metadata.
    pub fn get_metadata(&self, key: &str) -> Result<Option<Vec<u8>>> {
        self.get(METADATA_TABLE, key)
    }
}

/// Create default store location.
pub fn default_store() -> Result<Store> {
    let home = dirs::home_dir().context("Cannot determine home directory")?;
    let data_dir = home.join(".rsclaw").join("data");
    Store::new(&data_dir)
}