db/heromodels/src/herodb/tst_index.rs

use crate::db::error::{DbError, DbResult};
use crate::db::model::IndexKey;
use std::path::{Path, PathBuf};
use std::collections::HashMap;
use tst::TST;

/// Manages TST-based indexes for model objects
pub struct TSTIndexManager {
    /// Base path for TST databases
    base_path: PathBuf,
    
    /// Map of model prefixes to their TST instances
    tst_instances: HashMap<String, TST>,
}

impl TSTIndexManager {
    /// Creates a new TST index manager
    pub fn new<P: AsRef<Path>>(base_path: P) -> DbResult<Self> {
        let base_path = base_path.as_ref().to_path_buf();
        
        // Create directory if it doesn't exist
        std::fs::create_dir_all(&base_path).map_err(DbError::IoError)?;
        
        Ok(Self {
            base_path,
            tst_instances: HashMap::new(),
        })
    }
    
    /// Gets or creates a TST instance for a model prefix
    pub fn get_tst(&mut self, prefix: &str) -> DbResult<&mut TST> {
        if !self.tst_instances.contains_key(prefix) {
            // Create a new TST instance for this prefix
            let tst_path = self.base_path.join(format!("{}_tst", prefix));
            let tst_path_str = tst_path.to_string_lossy().to_string();
            
            // Create the TST
            let tst = TST::new(&tst_path_str, false)
                .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?;
            
            // Insert it into the map
            self.tst_instances.insert(prefix.to_string(), tst);
        }
        
        // Return a mutable reference to the TST
        Ok(self.tst_instances.get_mut(prefix).unwrap())
    }
    
    /// Adds or updates an object in the TST index with primary key
    pub fn set(&mut self, prefix: &str, id: u32, data: Vec<u8>) -> DbResult<()> {
        // Get the TST for this prefix
        let tst = self.get_tst(prefix)?;
        
        // Create the primary key in the format prefix_id
        let key = format!("{}_{}", prefix, id);
        
        // Set the key-value pair in the TST
        tst.set(&key, data.clone())
            .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?;
        
        Ok(())
    }
    
    /// Adds or updates an object in the TST index with additional index keys
    pub fn set_with_indexes(&mut self, prefix: &str, id: u32, data: Vec<u8>, index_keys: &[IndexKey]) -> DbResult<()> {
        // First set the primary key
        self.set(prefix, id, data.clone())?;
        
        // Get the TST for this prefix
        let tst = self.get_tst(prefix)?;
        
        // Add additional index keys
        for index_key in index_keys {
            // Create the index key in the format prefix_indexname_value
            let key = format!("{}_{}_{}", prefix, index_key.name, index_key.value);
            
            // Set the key-value pair in the TST
            // For index keys, we store the ID as the value
            let id_bytes = id.to_be_bytes().to_vec();
            tst.set(&key, id_bytes)
                .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?;
        }
        
        Ok(())
    }
    
    /// Removes an object from the TST index (primary key only)
    pub fn delete(&mut self, prefix: &str, id: u32) -> DbResult<()> {
        // Get the TST for this prefix
        let tst = self.get_tst(prefix)?;
        
        // Create the key in the format prefix_id
        let key = format!("{}_{}", prefix, id);
        
        // Delete the key from the TST
        tst.delete(&key)
            .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?;
        
        Ok(())
    }
    
    /// Removes an object from the TST index including all index keys
    pub fn delete_with_indexes(&mut self, prefix: &str, id: u32, index_keys: &[IndexKey]) -> DbResult<()> {
        // First delete the primary key
        self.delete(prefix, id)?;
        
        // Get the TST for this prefix
        let tst = self.get_tst(prefix)?;
        
        // Delete additional index keys
        for index_key in index_keys {
            // Create the index key in the format prefix_indexname_value
            let key = format!("{}_{}_{}", prefix, index_key.name, index_key.value);
            
            // Delete the key from the TST
            tst.delete(&key)
                .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?;
        }
        
        Ok(())
    }
    
    /// Lists all objects with a given prefix (primary keys only)
    pub fn list(&mut self, prefix: &str) -> DbResult<Vec<(u32, Vec<u8>)>> {
        // Get the TST for this prefix
        let tst = self.get_tst(prefix)?;
        
        // Get all keys with this prefix followed by an underscore
        let search_prefix = format!("{}_", prefix);
        let keys = tst.list(&search_prefix)
            .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?;
        
        // Get all values for these keys
        let mut result = Vec::with_capacity(keys.len());
        for key in keys {
            // Check if this is a primary key (prefix_id) and not an index key (prefix_indexname_value)
            let parts: Vec<&str> = key.split('_').collect();
            if parts.len() != 2 {
                continue; // Skip index keys
            }
            
            // Extract the ID from the key (format: prefix_id)
            let id_str = parts[1];
            let id = id_str.parse::<u32>().map_err(|_| {
                DbError::GeneralError(format!("Invalid ID in key: {}", key))
            })?;
            
            // Get the value from the TST
            let data = tst.get(&key)
                .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?;
            
            result.push((id, data));
        }
        
        Ok(result)
    }
    
    /// Finds objects by a specific index key
    pub fn find_by_index(&mut self, prefix: &str, index_name: &str, index_value: &str) -> DbResult<Vec<u32>> {
        // Get the TST for this prefix
        let tst = self.get_tst(prefix)?;
        
        // Create the index key in the format prefix_indexname_value
        let key = format!("{}_{}_{}", prefix, index_name, index_value);
        
        // Try to get the value from the TST
        match tst.get(&key) {
            Ok(id_bytes) => {
                // Convert the bytes to a u32 ID
                if id_bytes.len() == 4 {
                    let mut bytes = [0u8; 4];
                    bytes.copy_from_slice(&id_bytes[0..4]);
                    let id = u32::from_be_bytes(bytes);
                    Ok(vec![id])
                } else {
                    Err(DbError::GeneralError(format!("Invalid ID bytes for key: {}", key)))
                }
            },
            Err(_) => Ok(Vec::new()), // No matches found
        }
    }
    
    /// Finds objects by a prefix of an index key
    pub fn find_by_index_prefix(&mut self, prefix: &str, index_name: &str, index_value_prefix: &str) -> DbResult<Vec<u32>> {
        // Get the TST for this prefix
        let tst = self.get_tst(prefix)?;
        
        // Create the index key prefix in the format prefix_indexname_valueprefix
        let key_prefix = format!("{}_{}_{}", prefix, index_name, index_value_prefix);
        
        // Get all keys with this prefix
        let keys = tst.list(&key_prefix)
            .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?;
        
        // Extract the IDs from the values
        let mut result = Vec::with_capacity(keys.len());
        for key in keys {
            // Get the value from the TST
            let id_bytes = tst.get(&key)
                .map_err(|e| DbError::GeneralError(format!("TST error: {:?}", e)))?;
            
            // Convert the bytes to a u32 ID
            if id_bytes.len() == 4 {
                let mut bytes = [0u8; 4];
                bytes.copy_from_slice(&id_bytes[0..4]);
                let id = u32::from_be_bytes(bytes);
                result.push(id);
            }
        }
        
        Ok(result)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;
    
    #[test]
    fn test_tst_index_manager() {
        // Create a temporary directory for the test
        let temp_dir = tempdir().unwrap();
        let path = temp_dir.path();
        
        // Create a TST index manager
        let mut manager = TSTIndexManager::new(path).unwrap();
        
        // Test setting values
        let data1 = vec![1, 2, 3];
        let data2 = vec![4, 5, 6];
        manager.set("test", 1, data1.clone()).unwrap();
        manager.set("test", 2, data2.clone()).unwrap();
        
        // Test listing values
        let items = manager.list("test").unwrap();
        assert_eq!(items.len(), 2);
        
        // Check that the values are correct
        let mut found_data1 = false;
        let mut found_data2 = false;
        for (id, data) in items {
            if id == 1 && data == data1 {
                found_data1 = true;
            } else if id == 2 && data == data2 {
                found_data2 = true;
            }
        }
        assert!(found_data1);
        assert!(found_data2);
        
        // Test deleting a value
        manager.delete("test", 1).unwrap();
        
        // Test listing again
        let items = manager.list("test").unwrap();
        assert_eq!(items.len(), 1);
        assert_eq!(items[0].0, 2);
        assert_eq!(items[0].1, data2);
    }
}