audb_runtime/

database.rs

//! Database connection and query execution with Manifold backend
//!
//! This module provides the `Database` type for managing Manifold database connections
//! and executing queries. Each database uses Manifold's ColumnFamilyDatabase for
//! multi-tenant storage with concurrent write support.

use crate::error::{QueryError, Result};
use crate::types::QueryResult;
use manifold::column_family::ColumnFamilyDatabase;
use manifold_graph::{GraphTable, GraphTableRead};
use manifold_properties::{PropertyTable, PropertyTableRead, PropertyValue};
use manifold_vectors::{VectorTable, VectorTableRead};
use std::path::{Path, PathBuf};
use uuid::Uuid;

/// Database connection wrapper
///
/// Wraps a Manifold ColumnFamilyDatabase and provides query execution methods.
/// This is the main entry point for generated code.
///
/// ## Architecture
///
/// - Each schema becomes a column family in Manifold
/// - Properties stored via manifold-properties (typed, temporal)
/// - Relationships stored via manifold-graph (bidirectional edges)
/// - Vectors stored via manifold-vectors (zero-copy access)
///
/// ## Thread Safety
///
/// Database is `Send + Sync` and can be shared across async tasks.
pub struct Database {
    /// Path to the database
    path: PathBuf,

    /// Manifold column family database
    cf_db: ColumnFamilyDatabase,
}

impl Database {
    /// Open a database connection
    ///
    /// Creates or opens a Manifold database at the specified path.
    ///
    /// ## Example
    ///
    /// ```no_run
    /// # use audb_runtime::Database;
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let db = Database::open("./data/app.manifold").await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        let path = path.as_ref().to_path_buf();

        // Ensure parent directory exists
        if let Some(parent) = path.parent() {
            if !parent.exists() {
                std::fs::create_dir_all(parent).map_err(|e| QueryError::ConnectionError {
                    message: format!("Failed to create parent directory: {}", e),
                })?;
            }
        }

        // Open Manifold ColumnFamilyDatabase (creates file if it doesn't exist)
        let cf_db = ColumnFamilyDatabase::builder().open(&path).map_err(|e| {
            QueryError::ConnectionError {
                message: format!("Failed to open Manifold database: {}", e),
            }
        })?;

        Ok(Self { path, cf_db })
    }

    /// Close the database connection
    ///
    /// Explicitly closes the connection and flushes any pending writes.
    /// The database is automatically closed when dropped.
    pub async fn close(self) -> Result<()> {
        // Manifold handles cleanup on drop
        drop(self.cf_db);
        Ok(())
    }

    /// Get the database path
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Get a collection (column family) for a schema
    ///
    /// This is the primary method used by generated code to access schema storage.
    ///
    /// ## Example
    ///
    /// ```no_run
    /// # use audb_runtime::Database;
    /// # async fn example(db: &Database) -> Result<(), Box<dyn std::error::Error>> {
    /// let users_collection = db.collection("users")?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn collection(&self, name: &str) -> Result<Collection> {
        let cf =
            self.cf_db
                .column_family_or_create(name)
                .map_err(|e| QueryError::ConnectionError {
                    message: format!("Failed to get column family '{}': {}", name, e),
                })?;

        Ok(Collection { cf })
    }

    /// Execute a HyperQL query (stub for now)
    ///
    /// This will be implemented when HyperQL compiler is integrated.
    pub async fn execute_hyperql(&self, _query: &str) -> Result<QueryResult> {
        // TODO: Implement HyperQL query compilation and execution
        Ok(QueryResult {
            rows: Vec::new(),
            affected_rows: 0,
        })
    }

    /// Execute a SQL query (stub for now)
    pub async fn execute_sql(&self, _query: &str) -> Result<QueryResult> {
        // TODO: Implement SQL query execution
        Ok(QueryResult {
            rows: Vec::new(),
            affected_rows: 0,
        })
    }

    /// Execute a Cypher query (stub for now)
    pub async fn execute_cypher(&self, _query: &str) -> Result<QueryResult> {
        // TODO: Implement Cypher query execution
        Ok(QueryResult {
            rows: Vec::new(),
            affected_rows: 0,
        })
    }

    /// Execute a custom language query
    pub async fn execute_custom(&self, language: &str, query: &str) -> Result<QueryResult> {
        match language {
            "hyperql" => self.execute_hyperql(query).await,
            "sql" => self.execute_sql(query).await,
            "cypher" => self.execute_cypher(query).await,
            _ => Err(QueryError::UnsupportedLanguage {
                language: language.to_string(),
            }),
        }
    }
}

/// Collection wrapper for a schema
///
/// Each collection corresponds to a schema and provides CRUD operations
/// backed by Manifold's PropertyTable and GraphTable.
pub struct Collection {
    cf: manifold::column_family::ColumnFamily,
}

impl Collection {
    /// Create an entity in the collection
    ///
    /// Stores entity properties in PropertyTable with native types and temporal tracking.
    ///
    /// ## Example
    ///
    /// ```no_run
    /// # use audb_runtime::Database;
    /// # use uuid::Uuid;
    /// # use serde_json::json;
    /// # async fn example(db: &Database) -> Result<(), Box<dyn std::error::Error>> {
    /// let collection = db.collection("users")?;
    /// let id = Uuid::new_v4();
    /// let data = json!({
    ///     "name": "Alice",
    ///     "age": 42,
    ///     "active": true
    /// });
    /// collection.create_entity(id, data)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn create_entity(&self, id: Uuid, data: serde_json::Value) -> Result<()> {
        let write_txn = self
            .cf
            .begin_write()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to begin transaction: {}", e),
            })?;

        let mut props = PropertyTable::open(&write_txn, "properties").map_err(|e| {
            QueryError::ExecutionError {
                message: format!("Failed to open PropertyTable: {}", e),
            }
        })?;

        // Convert JSON to properties
        if let serde_json::Value::Object(map) = data {
            for (key, value) in map {
                let prop_value = match value {
                    serde_json::Value::Number(n) if n.is_i64() => {
                        PropertyValue::new_integer(n.as_i64().unwrap())
                    }
                    serde_json::Value::Number(n) if n.is_f64() => {
                        PropertyValue::new_float(n.as_f64().unwrap())
                    }
                    serde_json::Value::Bool(b) => PropertyValue::new_boolean(b),
                    serde_json::Value::String(s) => PropertyValue::new_string(s),
                    serde_json::Value::Null => PropertyValue::new_null(),
                    // Store complex types (arrays, objects) as JSON strings
                    _ => PropertyValue::new_string(value.to_string()),
                };

                props.set(&id, key.as_str(), prop_value).map_err(|e| {
                    QueryError::ExecutionError {
                        message: format!("Failed to set property: {}", e),
                    }
                })?;
            }
        }

        drop(props);
        write_txn.commit().map_err(|e| QueryError::ExecutionError {
            message: format!("Failed to commit transaction: {}", e),
        })?;

        Ok(())
    }

    /// Get an entity from the collection
    ///
    /// Returns None if the entity doesn't exist.
    pub fn get_entity(&self, id: Uuid) -> Result<Option<serde_json::Value>> {
        let read_txn = self
            .cf
            .begin_read()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to begin read transaction: {}", e),
            })?;

        let props = PropertyTableRead::open(&read_txn, "properties").map_err(|e| {
            QueryError::ExecutionError {
                message: format!("Failed to open PropertyTable: {}", e),
            }
        })?;

        // Get all properties for this entity
        let mut map = serde_json::Map::new();
        let properties = props.get_all(&id).map_err(|e| QueryError::ExecutionError {
            message: format!("Failed to read properties: {}", e),
        })?;

        for (key, value_guard) in properties {
            let value = value_guard.value();

            let json_value = if let Some(i) = value.as_integer() {
                serde_json::Value::Number(i.into())
            } else if let Some(f) = value.as_float() {
                serde_json::Number::from_f64(f)
                    .map(serde_json::Value::Number)
                    .unwrap_or(serde_json::Value::Null)
            } else if let Some(b) = value.as_boolean() {
                serde_json::Value::Bool(b)
            } else if let Some(s) = value.as_string() {
                // Try to parse as JSON first (for complex types stored as JSON strings)
                serde_json::from_str(s).unwrap_or_else(|_| serde_json::Value::String(s.to_string()))
            } else {
                serde_json::Value::Null
            };

            map.insert(key.to_string(), json_value);
        }

        if map.is_empty() {
            Ok(None)
        } else {
            Ok(Some(serde_json::Value::Object(map)))
        }
    }

    /// Update an entity in the collection
    ///
    /// Properties are automatically versioned with temporal tracking.
    pub fn update_entity(&self, id: Uuid, data: serde_json::Value) -> Result<()> {
        // Properties are automatically versioned, so update = create
        self.create_entity(id, data)
    }

    /// Delete an entity from the collection
    ///
    /// Removes all properties for the entity (soft delete via temporal tracking).
    pub fn delete_entity(&self, id: Uuid) -> Result<()> {
        let write_txn = self
            .cf
            .begin_write()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to begin transaction: {}", e),
            })?;

        let mut props = PropertyTable::open(&write_txn, "properties").map_err(|e| {
            QueryError::ExecutionError {
                message: format!("Failed to open PropertyTable: {}", e),
            }
        })?;

        // Read property keys first
        let read_txn = self
            .cf
            .begin_read()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to begin read transaction: {}", e),
            })?;

        let read_props = PropertyTableRead::open(&read_txn, "properties").map_err(|e| {
            QueryError::ExecutionError {
                message: format!("Failed to open PropertyTable: {}", e),
            }
        })?;

        // Collect property keys to delete
        let keys_to_delete: Vec<String> = read_props
            .get_all(&id)
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to read properties: {}", e),
            })?
            .iter()
            .map(|(key, _)| key.to_string())
            .collect();

        drop(read_props);
        drop(read_txn);

        // Convert to references for bulk delete
        let keys_refs: Vec<(Uuid, &str)> = keys_to_delete
            .iter()
            .map(|key| (id, key.as_str()))
            .collect();

        // Delete properties
        props
            .remove_bulk(&keys_refs)
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to delete properties: {}", e),
            })?;

        drop(props);
        write_txn.commit().map_err(|e| QueryError::ExecutionError {
            message: format!("Failed to commit transaction: {}", e),
        })?;

        Ok(())
    }

    /// Add an edge between two entities
    ///
    /// Creates a bidirectional edge in the GraphTable.
    pub fn add_edge(&self, source: Uuid, edge_type: &str, target: Uuid) -> Result<()> {
        let write_txn = self
            .cf
            .begin_write()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to begin transaction: {}", e),
            })?;

        let mut graph =
            GraphTable::open(&write_txn, "edges").map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to open GraphTable: {}", e),
            })?;

        graph
            .add_edge(&source, edge_type, &target, true, 1.0, None)
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to add edge: {}", e),
            })?;

        drop(graph);
        write_txn.commit().map_err(|e| QueryError::ExecutionError {
            message: format!("Failed to commit transaction: {}", e),
        })?;

        Ok(())
    }

    /// Get outgoing edges from an entity
    ///
    /// Returns all target UUIDs for edges of the specified type.
    pub fn get_outgoing_edges(&self, source: Uuid, edge_type: &str) -> Result<Vec<Uuid>> {
        let read_txn = self
            .cf
            .begin_read()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to begin read transaction: {}", e),
            })?;

        let graph =
            GraphTableRead::open(&read_txn, "edges").map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to open GraphTable: {}", e),
            })?;

        let mut targets = Vec::new();
        let edges = graph
            .outgoing_edges(&source)
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to read outgoing edges: {}", e),
            })?;

        for edge_result in edges {
            let edge = edge_result.map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to read edge: {}", e),
            })?;

            if edge.edge_type == edge_type && edge.is_active {
                targets.push(edge.target);
            }
        }

        Ok(targets)
    }

    /// Get incoming edges to an entity
    ///
    /// Returns all source UUIDs for edges of the specified type.
    pub fn get_incoming_edges(&self, target: Uuid, edge_type: &str) -> Result<Vec<Uuid>> {
        let read_txn = self
            .cf
            .begin_read()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to begin read transaction: {}", e),
            })?;

        let graph =
            GraphTableRead::open(&read_txn, "edges").map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to open GraphTable: {}", e),
            })?;

        let mut sources = Vec::new();
        let edges = graph
            .incoming_edges(&target)
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to read incoming edges: {}", e),
            })?;

        for edge_result in edges {
            let edge = edge_result.map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to read edge: {}", e),
            })?;

            if edge.edge_type == edge_type && edge.is_active {
                sources.push(edge.source);
            }
        }

        Ok(sources)
    }

    /// List all entity IDs in the collection
    ///
    /// Returns a vector of all UUIDs that have properties in this collection.
    ///
    /// ## Example
    ///
    /// ```no_run
    /// # use audb_runtime::Database;
    /// # async fn example(db: &Database) -> Result<(), Box<dyn std::error::Error>> {
    /// let collection = db.collection("users")?;
    /// let ids = collection.list_all_ids()?;
    /// println!("Found {} entities", ids.len());
    /// # Ok(())
    /// # }
    /// ```
    pub fn list_all_ids(&self) -> Result<Vec<Uuid>> {
        let read_txn = self
            .cf
            .begin_read()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to begin read transaction: {}", e),
            })?;

        let props = PropertyTableRead::open(&read_txn, "properties").map_err(|e| {
            QueryError::ExecutionError {
                message: format!("Failed to open PropertyTable: {}", e),
            }
        })?;

        let mut entity_ids = std::collections::HashSet::new();

        // Scan all properties and collect unique entity IDs
        for result in props.iter().map_err(|e| QueryError::ExecutionError {
            message: format!("Failed to iterate properties: {}", e),
        })? {
            let ((entity_id, _key), _value) = result.map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to read property: {}", e),
            })?;
            entity_ids.insert(entity_id);
        }

        Ok(entity_ids.into_iter().collect())
    }

    /// List all entities in the collection
    ///
    /// Returns all entities with their JSON data.
    ///
    /// ## Example
    ///
    /// ```no_run
    /// # use audb_runtime::Database;
    /// # async fn example(db: &Database) -> Result<(), Box<dyn std::error::Error>> {
    /// let collection = db.collection("users")?;
    /// let entities = collection.list_all_entities()?;
    /// for (id, data) in entities {
    ///     println!("Entity {}: {:?}", id, data);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn list_all_entities(&self) -> Result<Vec<(Uuid, serde_json::Value)>> {
        let ids = self.list_all_ids()?;
        let mut entities = Vec::new();

        for id in ids {
            if let Some(data) = self.get_entity(id)? {
                entities.push((id, data));
            }
        }

        Ok(entities)
    }

    /// Get a VectorTable for storing and retrieving embeddings
    ///
    /// The dimension is a compile-time const generic, typically provided
    /// by the code generator based on Tessera model metadata.
    ///
    /// ## Example
    ///
    /// ```no_run
    /// # use audb_runtime::Database;
    /// # use uuid::Uuid;
    /// # async fn example(db: &Database) -> Result<(), Box<dyn std::error::Error>> {
    /// let collection = db.collection("documents")?;
    /// let vectors = collection.vectors::<384>("embeddings")?;
    ///
    /// let id = Uuid::new_v4();
    /// let embedding = vec![0.1f32; 384];
    /// vectors.insert(&id, &embedding)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn vectors<const DIM: usize>(&self, table_name: &str) -> Result<VectorTableWrapper<DIM>> {
        Ok(VectorTableWrapper {
            cf: self.cf.clone(),
            table_name: table_name.to_string(),
        })
    }
}

/// Wrapper for VectorTable operations
///
/// Provides write and read access to vector embeddings with const generic dimensions.
pub struct VectorTableWrapper<const DIM: usize> {
    cf: manifold::column_family::ColumnFamily,
    table_name: String,
}

impl<const DIM: usize> VectorTableWrapper<DIM> {
    /// Insert a vector for an entity
    pub fn insert(&self, id: &Uuid, vector: &[f32]) -> Result<()> {
        if vector.len() != DIM {
            return Err(QueryError::ExecutionError {
                message: format!(
                    "Vector dimension mismatch: expected {}, got {}",
                    DIM,
                    vector.len()
                ),
            });
        }

        let write_txn = self
            .cf
            .begin_write()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to begin transaction: {}", e),
            })?;

        let mut vectors = VectorTable::<DIM>::open(&write_txn, &self.table_name).map_err(|e| {
            QueryError::ExecutionError {
                message: format!("Failed to open VectorTable: {}", e),
            }
        })?;

        // Convert to fixed-size array
        let mut arr = [0.0f32; DIM];
        arr.copy_from_slice(vector);

        vectors
            .insert(id, &arr)
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to insert vector: {}", e),
            })?;

        drop(vectors);
        write_txn.commit().map_err(|e| QueryError::ExecutionError {
            message: format!("Failed to commit transaction: {}", e),
        })?;

        Ok(())
    }

    /// Get a vector for an entity
    pub fn get(&self, id: &Uuid) -> Result<Option<Vec<f32>>> {
        let read_txn = self
            .cf
            .begin_read()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to begin read transaction: {}", e),
            })?;

        let vectors = VectorTableRead::<DIM>::open(&read_txn, &self.table_name).map_err(|e| {
            QueryError::ExecutionError {
                message: format!("Failed to open VectorTable: {}", e),
            }
        })?;

        let result = vectors.get(id).map_err(|e| QueryError::ExecutionError {
            message: format!("Failed to read vector: {}", e),
        })?;

        Ok(result.map(|guard| guard.value().to_vec()))
    }

    /// Search for similar vectors using brute-force cosine similarity
    ///
    /// This implementation scans all vectors and computes cosine similarity.
    /// For production use with large datasets, consider integrating HNSW.
    ///
    /// # Arguments
    ///
    /// * `query` - Query vector to search for
    /// * `limit` - Maximum number of results to return
    ///
    /// # Returns
    ///
    /// Returns UUIDs sorted by similarity (highest first)
    pub fn search_similar(&self, query: &[f32], limit: usize) -> Result<Vec<Uuid>> {
        if query.len() != DIM {
            return Err(QueryError::ExecutionError {
                message: format!(
                    "Query vector dimension mismatch: expected {}, got {}",
                    DIM,
                    query.len()
                ),
            });
        }

        let read_txn = self
            .cf
            .begin_read()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to begin read transaction: {}", e),
            })?;

        let vectors = VectorTableRead::<DIM>::open(&read_txn, &self.table_name).map_err(|e| {
            QueryError::ExecutionError {
                message: format!("Failed to open VectorTable: {}", e),
            }
        })?;

        // Compute cosine similarity for all vectors
        let mut similarities: Vec<(Uuid, f32)> = Vec::new();

        // Iterate over all vectors in the table
        // Note: This is a brute-force scan - O(n) where n is number of vectors
        let iter = vectors
            .all_vectors()
            .map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to iterate vectors: {}", e),
            })?;

        for result in iter {
            let (id, vector_guard) = result.map_err(|e| QueryError::ExecutionError {
                message: format!("Failed to read vector entry: {}", e),
            })?;

            let vector = vector_guard.value();
            let similarity = cosine_similarity(query, vector);
            similarities.push((id, similarity));
        }

        // Sort by similarity (descending)
        similarities.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        // Take top K results
        let results: Vec<Uuid> = similarities
            .into_iter()
            .take(limit)
            .map(|(id, _)| id)
            .collect();

        Ok(results)
    }
}

/// Compute cosine similarity between two vectors
///
/// Returns a value between -1.0 and 1.0, where 1.0 means identical direction.
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    debug_assert_eq!(a.len(), b.len(), "Vectors must have same length");

    let mut dot_product = 0.0;
    let mut norm_a = 0.0;
    let mut norm_b = 0.0;

    for i in 0..a.len() {
        dot_product += a[i] * b[i];
        norm_a += a[i] * a[i];
        norm_b += b[i] * b[i];
    }

    let magnitude = (norm_a * norm_b).sqrt();
    if magnitude == 0.0 {
        0.0
    } else {
        dot_product / magnitude
    }
}

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

    #[tokio::test]
    async fn test_database_open() {
        let temp_dir = std::env::temp_dir().join("audb_test_manifold_open.manifold");
        let wal_path = temp_dir.with_extension("wal");
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);

        let db = Database::open(&temp_dir).await;
        assert!(db.is_ok());

        let db = db.unwrap();
        assert_eq!(db.path(), temp_dir.as_path());

        // Cleanup
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);
    }

    #[tokio::test]
    async fn test_create_and_get_entity() {
        let temp_dir = std::env::temp_dir().join("audb_test_manifold_crud.manifold");
        let wal_path = temp_dir.with_extension("wal");
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);

        let db = Database::open(&temp_dir).await.unwrap();
        let collection = db.collection("users").unwrap();

        let id = Uuid::new_v4();
        let data = json!({
            "name": "Alice",
            "age": 42,
            "active": true
        });

        // Create entity
        collection.create_entity(id, data.clone()).unwrap();

        // Get entity
        let retrieved = collection.get_entity(id).unwrap();
        assert!(retrieved.is_some());

        let retrieved = retrieved.unwrap();
        assert_eq!(retrieved["name"], "Alice");
        assert_eq!(retrieved["age"], 42);
        assert_eq!(retrieved["active"], true);

        // Cleanup
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);
    }

    #[tokio::test]
    async fn test_update_entity() {
        let temp_dir = std::env::temp_dir().join("audb_test_manifold_update.manifold");
        let wal_path = temp_dir.with_extension("wal");
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);

        let db = Database::open(&temp_dir).await.unwrap();
        let collection = db.collection("users").unwrap();

        let id = Uuid::new_v4();
        let data = json!({ "name": "Alice", "age": 42 });

        collection.create_entity(id, data).unwrap();

        // Update
        let updated = json!({ "name": "Alice", "age": 43 });
        collection.update_entity(id, updated).unwrap();

        // Verify
        let retrieved = collection.get_entity(id).unwrap().unwrap();
        assert_eq!(retrieved["age"], 43);

        // Cleanup
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);
    }

    #[tokio::test]
    async fn test_delete_entity() {
        let temp_dir = std::env::temp_dir().join("audb_test_manifold_delete.manifold");
        let wal_path = temp_dir.with_extension("wal");
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);

        let db = Database::open(&temp_dir).await.unwrap();
        let collection = db.collection("users").unwrap();

        let id = Uuid::new_v4();
        let data = json!({ "name": "Bob" });

        collection.create_entity(id, data).unwrap();
        assert!(collection.get_entity(id).unwrap().is_some());

        collection.delete_entity(id).unwrap();
        assert!(collection.get_entity(id).unwrap().is_none());

        // Cleanup
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);
    }

    #[tokio::test]
    async fn test_edges() {
        let temp_dir = std::env::temp_dir().join("audb_test_manifold_edges.manifold");
        let wal_path = temp_dir.with_extension("wal");
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);

        let db = Database::open(&temp_dir).await.unwrap();
        let collection = db.collection("test").unwrap();

        let user = Uuid::new_v4();
        let post1 = Uuid::new_v4();
        let post2 = Uuid::new_v4();

        // Add edges
        collection.add_edge(post1, "authored_by", user).unwrap();
        collection.add_edge(post2, "authored_by", user).unwrap();

        // Get outgoing edges from posts
        let authors1 = collection.get_outgoing_edges(post1, "authored_by").unwrap();
        assert_eq!(authors1, vec![user]);

        // Get incoming edges to user
        let authored_posts = collection.get_incoming_edges(user, "authored_by").unwrap();
        assert_eq!(authored_posts.len(), 2);
        assert!(authored_posts.contains(&post1));
        assert!(authored_posts.contains(&post2));

        // Cleanup
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);
    }

    #[tokio::test]
    async fn test_vectors() {
        let temp_dir = std::env::temp_dir().join("audb_test_manifold_vectors.manifold");
        let wal_path = temp_dir.with_extension("wal");
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);

        let db = Database::open(&temp_dir).await.unwrap();
        let collection = db.collection("documents").unwrap();

        let id = Uuid::new_v4();
        let embedding = vec![0.1f32, 0.2, 0.3, 0.4]; // 4D vector for testing

        // Get vector table with dimension 4
        let vectors = collection.vectors::<4>("embeddings").unwrap();

        // Insert vector
        vectors.insert(&id, &embedding).unwrap();

        // Get vector
        let retrieved = vectors.get(&id).unwrap();
        assert!(retrieved.is_some());
        let retrieved_vec = retrieved.unwrap();
        assert_eq!(retrieved_vec.len(), 4);
        assert_eq!(retrieved_vec, embedding);

        // Test dimension mismatch error
        let wrong_dim = vec![0.1f32; 8];
        let result = vectors.insert(&id, &wrong_dim);
        assert!(result.is_err());

        // Cleanup
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);
    }

    #[tokio::test]
    async fn test_vector_similarity_search() {
        let temp_dir = std::env::temp_dir().join("audb_test_manifold_similarity.manifold");
        let wal_path = temp_dir.with_extension("wal");
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);

        let db = Database::open(&temp_dir).await.unwrap();
        let collection = db.collection("documents").unwrap();

        let vectors = collection.vectors::<4>("embeddings").unwrap();

        // Insert test vectors
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();
        let id3 = Uuid::new_v4();

        // Similar to query
        let vec1 = vec![1.0f32, 0.0, 0.0, 0.0];
        // Somewhat similar
        let vec2 = vec![0.7f32, 0.7, 0.0, 0.0];
        // Different
        let vec3 = vec![0.0f32, 0.0, 1.0, 0.0];

        vectors.insert(&id1, &vec1).unwrap();
        vectors.insert(&id2, &vec2).unwrap();
        vectors.insert(&id3, &vec3).unwrap();

        // Search with query similar to vec1
        let query = vec![0.9f32, 0.1, 0.0, 0.0];
        let results = vectors.search_similar(&query, 2).unwrap();

        // Should return id1 first (most similar)
        assert_eq!(results.len(), 2);
        assert_eq!(results[0], id1);

        // Cleanup
        let _ = std::fs::remove_file(&temp_dir);
        let _ = std::fs::remove_file(&wal_path);
    }

    #[test]
    fn test_cosine_similarity() {
        // Identical vectors
        let a = vec![1.0f32, 0.0, 0.0];
        let b = vec![1.0f32, 0.0, 0.0];
        let sim = cosine_similarity(&a, &b);
        assert!((sim - 1.0).abs() < 0.001);

        // Orthogonal vectors
        let a = vec![1.0f32, 0.0, 0.0];
        let b = vec![0.0f32, 1.0, 0.0];
        let sim = cosine_similarity(&a, &b);
        assert!(sim.abs() < 0.001);

        // Opposite vectors
        let a = vec![1.0f32, 0.0, 0.0];
        let b = vec![-1.0f32, 0.0, 0.0];
        let sim = cosine_similarity(&a, &b);
        assert!((sim + 1.0).abs() < 0.001);

        // Similar vectors
        let a = vec![1.0f32, 0.5, 0.0];
        let b = vec![0.9f32, 0.4, 0.0];
        let sim = cosine_similarity(&a, &b);
        assert!(sim > 0.99); // Very similar
    }
}