//! Python bindings for AiFastDb core database.
//!
//! Provides `AiFastDb` (wraps `CollectionStore`) and `PyCollection` (wraps `Collection`)
//! with high-level CRUD, search, and collection management APIs.

use std::path::PathBuf;
use std::sync::Arc;

use pyo3::prelude::*;
use pyo3::types::PyDict;
use pyo3::exceptions::{PyValueError, PyOSError, PyRuntimeError};

use aifastdb_core::storage::collection::{CollectionStore, Collection};
use aifastdb_core::types::collection::CollectionConfig;
use aifastdb_core::types::memory::{MemoryNode, ContentType};
use aifastdb_core::query::hybrid::{HybridQueryEngine, HybridQuery};

use crate::types::{json_value_to_py, py_to_json_value};

// ---- Helper: MemoryNode → Python dict ----

fn memory_node_to_dict(py: Python<'_>, node: &MemoryNode) -> PyResult<PyObject> {
    let dict = PyDict::new_bound(py);

    dict.set_item("id", &node.id)?;
    dict.set_item("content", &node.content)?;
    dict.set_item("content_type", format!("{:?}", node.content_type))?;
    dict.set_item("created_at", node.created_at)?;
    dict.set_item("last_accessed_at", node.last_accessed_at)?;
    dict.set_item("access_count", node.access_count)?;
    dict.set_item("importance", node.importance)?;
    dict.set_item("decay_rate", node.decay_rate)?;
    dict.set_item("depth", node.depth)?;

    if let Some(ref parent_id) = node.parent_id {
        dict.set_item("parent_id", parent_id)?;
    }

    if !node.tags.is_empty() {
        dict.set_item("tags", &node.tags)?;
    }

    if !node.children.is_empty() {
        dict.set_item("children", &node.children)?;
    }

    if let Some(ref embedding) = node.embedding {
        dict.set_item("embedding", embedding.clone())?;
    }

    if !node.metadata.is_empty() {
        let meta_dict = PyDict::new_bound(py);
        for (key, val) in &node.metadata {
            meta_dict.set_item(key, json_value_to_py(py, val)?)?;
        }
        dict.set_item("metadata", meta_dict)?;
    }

    if let Some(ref owner_id) = node.owner_id {
        dict.set_item("owner_id", owner_id)?;
    }

    Ok(dict.into_py(py))
}

// ---- Helper: parse ContentType from string ----

fn parse_content_type(s: &str) -> ContentType {
    match s.to_lowercase().as_str() {
        "text" => ContentType::Text,
        "image" => ContentType::Image,
        "code" => ContentType::Code,
        "conversation" => ContentType::Conversation,
        "document" => ContentType::Document,
        _ => ContentType::Other,
    }
}

// ---- AiFastDb class ----

/// AiFastDb — High-level Python interface to the AiFastDb database.
///
/// Usage:
///     with AiFastDb("/path/to/db") as db:
///         db.remember("id1", "hello world", tags=["greeting"])
///         result = db.get("id1")
///         db.forget("id1")
#[pyclass(name = "AiFastDb")]
pub struct PyAiFastDb {
    store: Arc<CollectionStore>,
    path: PathBuf,
    default_dimension: usize,
    #[cfg(feature = "lance-storage")]
    lance_manager: Option<aifastdb_core::LanceCollectionManager>,
}

#[pymethods]
impl PyAiFastDb {
    /// Create or open an AiFastDb database.
    ///
    /// Args:
    ///     path: Path to the database directory.
    ///     dimension: Default vector dimension for collections (default: 128).
    ///     storage_backend: "legacy" (default JSONL+mmap) or "lance" (Lance columnar).
    #[new]
    #[pyo3(signature = (path, *, dimension=128, storage_backend=None))]
    fn new(path: String, dimension: usize, storage_backend: Option<String>) -> PyResult<Self> {
        let path_buf = PathBuf::from(&path);
        let backend = storage_backend.as_deref().unwrap_or("legacy");

        #[cfg(feature = "lance-storage")]
        let lance_manager = if backend == "lance" {
            let mgr = aifastdb_core::LanceCollectionManager::open_or_create(&path_buf, dimension)
                .map_err(|e| PyOSError::new_err(format!("Failed to open Lance manager: {}", e)))?;
            Some(mgr)
        } else {
            None
        };

        #[cfg(feature = "lance-storage")]
        let store = if let Some(ref mgr) = lance_manager {
            mgr.store_arc()
        } else {
            Arc::new(CollectionStore::open_or_create(&path_buf)
                .map_err(|e| PyOSError::new_err(format!("Failed to open database: {}", e)))?)
        };

        #[cfg(not(feature = "lance-storage"))]
        let store = {
            if backend == "lance" {
                return Err(PyValueError::new_err(
                    "Lance storage requires the 'lance-storage' feature. Rebuild with --features lance-storage"
                ));
            }
            Arc::new(CollectionStore::open_or_create(&path_buf)
                .map_err(|e| PyOSError::new_err(format!("Failed to open database: {}", e)))?)
        };

        Ok(Self {
            store,
            path: path_buf,
            default_dimension: dimension,
            #[cfg(feature = "lance-storage")]
            lance_manager,
        })
    }

    /// Context manager entry.
    fn __enter__(slf: Py<Self>) -> Py<Self> {
        slf
    }

    /// Context manager exit — flushes all data to disk.
    #[pyo3(signature = (_exc_type=None, _exc_val=None, _exc_tb=None))]
    fn __exit__(
        &self,
        _exc_type: Option<&Bound<'_, PyAny>>,
        _exc_val: Option<&Bound<'_, PyAny>>,
        _exc_tb: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<bool> {
        self.store.flush()
            .map_err(|e| PyOSError::new_err(format!("Flush failed: {}", e)))?;
        Ok(false) // Don't suppress exceptions
    }

    /// Get the database path.
    #[getter]
    fn path(&self) -> String {
        self.path.to_string_lossy().to_string()
    }

    /// Get the default vector dimension.
    #[getter]
    fn dimension(&self) -> usize {
        self.default_dimension
    }

    // ---- Collection Management (T24.2) ----

    /// List all collection names.
    fn list_collections(&self) -> Vec<String> {
        self.store.list_collections()
    }

    /// Check if a collection exists.
    fn has_collection(&self, name: &str) -> bool {
        self.store.has_collection(name)
    }

    /// Delete a collection.
    fn delete_collection(&self, name: &str) -> PyResult<bool> {
        self.store.delete_collection(name)
            .map_err(|e| PyOSError::new_err(format!("Delete collection failed: {}", e)))
    }

    /// Get or create a collection and return a PyCollection wrapper.
    ///
    /// In Lance mode, collections are automatically backed by Lance datasets.
    #[pyo3(signature = (name, *, dimension=None, description=None))]
    fn collection(
        &self,
        name: String,
        dimension: Option<usize>,
        description: Option<String>,
    ) -> PyResult<PyCollection> {
        let dim = dimension.unwrap_or(self.default_dimension);
        let mut config = CollectionConfig::new(&name, dim);
        if let Some(desc) = description {
            config = config.with_description(desc);
        }

        #[cfg(feature = "lance-storage")]
        if let Some(ref mgr) = self.lance_manager {
            let collection = mgr.get_or_create_collection(config)
                .map_err(|e| PyOSError::new_err(format!("Lance collection error: {}", e)))?;
            return Ok(PyCollection { inner: collection, name });
        }

        let collection = self.store.get_or_create_collection(config)
            .map_err(|e| PyOSError::new_err(format!("Collection error: {}", e)))?;
        Ok(PyCollection { inner: collection, name })
    }

    // ---- Basic CRUD on default collection (T24.3) ----

    /// Store a memory node in the default collection.
    ///
    /// Args:
    ///     id: Unique identifier for the record.
    ///     content: Text content to store.
    ///     content_type: One of "text", "image", "code", "conversation", "document", "other".
    ///     tags: Optional list of tags.
    ///     metadata: Optional dict of metadata key-value pairs.
    ///     embedding: Optional vector embedding.
    ///     importance: Importance score (default: 1.0).
    ///     parent_id: Optional parent record ID.
    #[pyo3(signature = (id, content, *, content_type="text", tags=None, metadata=None, embedding=None, importance=1.0, parent_id=None))]
    fn remember(
        &self,
        py: Python<'_>,
        id: String,
        content: String,
        content_type: &str,
        tags: Option<Vec<String>>,
        metadata: Option<&Bound<'_, PyDict>>,
        embedding: Option<Vec<f32>>,
        importance: f64,
        parent_id: Option<String>,
    ) -> PyResult<()> {
        let collection = self._get_default_collection()?;
        Self::_remember_impl(py, &collection, id, content, content_type, tags, metadata, embedding, importance, parent_id)
    }

    /// Retrieve a memory node by ID from the default collection.
    ///
    /// Returns a dict or None if not found.
    fn get(&self, py: Python<'_>, id: &str) -> PyResult<Option<PyObject>> {
        let collection = self._get_default_collection()?;
        Self::_get_impl(py, &collection, id)
    }

    /// Delete a memory node by ID from the default collection.
    ///
    /// Returns True if the node was found and deleted.
    fn forget(&self, id: &str) -> PyResult<bool> {
        let collection = self._get_default_collection()?;
        Self::_forget_impl(&collection, id)
    }

    /// Record an access to a memory node (updates access_count and last_accessed_at).
    ///
    /// Returns True if the node was found.
    fn touch(&self, id: &str) -> PyResult<bool> {
        let collection = self._get_default_collection()?;
        Self::_touch_impl(&collection, id)
    }

    // ---- Batch Operations (T24.4) ----

    /// Store multiple memory nodes at once.
    ///
    /// Args:
    ///     items: List of dicts, each with at least "id" and "content" keys.
    ///            Optional keys: "content_type", "tags", "metadata", "embedding", "importance", "parent_id".
    fn remember_many(&self, py: Python<'_>, items: Vec<Bound<'_, PyDict>>) -> PyResult<usize> {
        let collection = self._get_default_collection()?;
        let mut count = 0;

        for item in &items {
            let id: String = item.get_item("id")?
                .ok_or_else(|| PyValueError::new_err("Each item must have an 'id' key"))?
                .extract()?;
            let content: String = item.get_item("content")?
                .ok_or_else(|| PyValueError::new_err("Each item must have a 'content' key"))?
                .extract()?;

            let content_type: String = item.get_item("content_type")?
                .map(|v| v.extract::<String>())
                .transpose()?
                .unwrap_or_else(|| "text".to_string());
            let tags: Option<Vec<String>> = item.get_item("tags")?
                .map(|v| v.extract::<Vec<String>>())
                .transpose()?;
            let embedding: Option<Vec<f32>> = item.get_item("embedding")?
                .map(|v| v.extract::<Vec<f32>>())
                .transpose()?;
            let importance: f64 = item.get_item("importance")?
                .map(|v| v.extract::<f64>())
                .transpose()?
                .unwrap_or(1.0);
            let parent_id: Option<String> = item.get_item("parent_id")?
                .map(|v| v.extract::<String>())
                .transpose()?;
            let metadata_obj = item.get_item("metadata")?;
            let metadata: Option<Bound<'_, PyDict>> = match metadata_obj {
                Some(ref v) => Some(v.downcast::<PyDict>()
                    .map_err(|_| PyValueError::new_err("'metadata' must be a dict"))?
                    .clone()),
                None => None,
            };

            Self::_remember_impl(
                py, &collection, id, content, &content_type,
                tags, metadata.as_ref(), embedding, importance, parent_id,
            )?;
            count += 1;
        }

        Ok(count)
    }

    /// Retrieve multiple memory nodes by IDs.
    ///
    /// Returns a dict mapping ID → node dict. Missing IDs are omitted.
    fn get_many(&self, py: Python<'_>, ids: Vec<String>) -> PyResult<PyObject> {
        let collection = self._get_default_collection()?;
        let result = PyDict::new_bound(py);

        for id in &ids {
            if let Some(node) = collection.docs.get(id).ok().flatten() {
                result.set_item(id, memory_node_to_dict(py, &node)?)?;
            }
        }

        Ok(result.into_py(py))
    }

    // ---- Search Operations (T24.5) ----

    /// Search for similar memories using hybrid vector + text search.
    ///
    /// Args:
    ///     query: Text query string.
    ///     embedding: Optional query vector for vector search.
    ///     limit: Maximum results to return (default: 10).
    ///     min_importance: Optional minimum importance filter.
    ///     tags: Optional tag filter (records must have at least one matching tag).
    ///
    /// Returns a list of dicts with "node" (the record dict) and "score" (float).
    #[pyo3(signature = (*, query=None, embedding=None, limit=10, min_importance=None))]
    fn recall(
        &self,
        py: Python<'_>,
        query: Option<String>,
        embedding: Option<Vec<f32>>,
        limit: usize,
        min_importance: Option<f64>,
    ) -> PyResult<Vec<PyObject>> {
        let collection = self._get_default_collection()?;
        Self::_recall_impl(py, &collection, query, embedding, limit, min_importance, false, 60.0)
    }

    /// Search using Reciprocal Rank Fusion (RRF) for better multi-signal results.
    ///
    /// Args:
    ///     Same as recall(), plus:
    ///     k: RRF parameter (default: 60.0). Lower = more weight to top results.
    #[pyo3(signature = (*, query=None, embedding=None, limit=10, min_importance=None, k=60.0))]
    fn recall_rrf(
        &self,
        py: Python<'_>,
        query: Option<String>,
        embedding: Option<Vec<f32>>,
        limit: usize,
        min_importance: Option<f64>,
        k: f64,
    ) -> PyResult<Vec<PyObject>> {
        let collection = self._get_default_collection()?;
        Self::_recall_impl(py, &collection, query, embedding, limit, min_importance, true, k)
    }

    /// Get all records in the default collection.
    ///
    /// Returns a list of dicts.
    fn recall_all(&self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
        let collection = self._get_default_collection()?;
        let nodes = collection.docs.all().unwrap_or_default();
        let mut results = Vec::with_capacity(nodes.len());
        for node in &nodes {
            results.push(memory_node_to_dict(py, node)?);
        }
        Ok(results)
    }

    // ---- Utility ----

    /// Flush all data to disk.
    fn flush(&self) -> PyResult<()> {
        self.store.flush()
            .map_err(|e| PyOSError::new_err(format!("Flush failed: {}", e)))
    }

    /// Get the document count in the default collection.
    fn document_count(&self) -> PyResult<usize> {
        let collection = self._get_default_collection()?;
        Ok(collection.document_count())
    }

    // ---- Lance Storage API ----

    /// Whether the database is using Lance storage backend.
    #[getter]
    fn is_lance_mode(&self) -> bool {
        #[cfg(feature = "lance-storage")]
        { self.lance_manager.is_some() }
        #[cfg(not(feature = "lance-storage"))]
        { false }
    }

    /// List all versions of a Lance collection.
    ///
    /// Returns a list of dicts with "version" (int) and "timestamp_ms" (int).
    fn lance_list_versions(&self, py: Python<'_>, collection_name: &str) -> PyResult<Vec<PyObject>> {
        #[cfg(feature = "lance-storage")]
        {
            let mgr = self.lance_manager.as_ref()
                .ok_or_else(|| PyValueError::new_err("Lance mode is not enabled"))?;
            let versions = mgr.list_versions(collection_name)
                .map_err(|e| PyRuntimeError::new_err(format!("list_versions failed: {}", e)))?;
            let mut out = Vec::with_capacity(versions.len());
            for v in &versions {
                let dict = PyDict::new_bound(py);
                dict.set_item("version", v.version)?;
                dict.set_item("timestamp_ms", v.timestamp_ms)?;
                out.push(dict.into_py(py));
            }
            Ok(out)
        }
        #[cfg(not(feature = "lance-storage"))]
        {
            let _ = (py, collection_name);
            Err(PyValueError::new_err("Lance storage feature is not enabled at compile time"))
        }
    }

    /// Get the current version number of a Lance collection.
    fn lance_current_version(&self, collection_name: &str) -> PyResult<u64> {
        #[cfg(feature = "lance-storage")]
        {
            let mgr = self.lance_manager.as_ref()
                .ok_or_else(|| PyValueError::new_err("Lance mode is not enabled"))?;
            mgr.current_version(collection_name)
                .map_err(|e| PyRuntimeError::new_err(format!("current_version failed: {}", e)))
        }
        #[cfg(not(feature = "lance-storage"))]
        {
            let _ = collection_name;
            Err(PyValueError::new_err("Lance storage feature is not enabled at compile time"))
        }
    }

    /// Checkout a specific version of a Lance collection (time-travel).
    fn lance_checkout_version(&self, collection_name: &str, version: u64) -> PyResult<()> {
        #[cfg(feature = "lance-storage")]
        {
            let mgr = self.lance_manager.as_ref()
                .ok_or_else(|| PyValueError::new_err("Lance mode is not enabled"))?;
            mgr.checkout_version(collection_name, version)
                .map_err(|e| PyRuntimeError::new_err(format!("checkout_version failed: {}", e)))
        }
        #[cfg(not(feature = "lance-storage"))]
        {
            let _ = (collection_name, version);
            Err(PyValueError::new_err("Lance storage feature is not enabled at compile time"))
        }
    }

    /// Compact a Lance collection's dataset.
    ///
    /// Returns the number of compacted fragments.
    fn lance_compact(&self, collection_name: &str) -> PyResult<usize> {
        #[cfg(feature = "lance-storage")]
        {
            let mgr = self.lance_manager.as_ref()
                .ok_or_else(|| PyValueError::new_err("Lance mode is not enabled"))?;
            mgr.compact(collection_name)
                .map_err(|e| PyRuntimeError::new_err(format!("compact failed: {}", e)))
        }
        #[cfg(not(feature = "lance-storage"))]
        {
            let _ = collection_name;
            Err(PyValueError::new_err("Lance storage feature is not enabled at compile time"))
        }
    }

    /// Clean up old versions of a Lance collection.
    ///
    /// Args:
    ///     older_than_secs: Remove versions older than this many seconds.
    fn lance_cleanup_old_versions(&self, collection_name: &str, older_than_secs: u64) -> PyResult<()> {
        #[cfg(feature = "lance-storage")]
        {
            let mgr = self.lance_manager.as_ref()
                .ok_or_else(|| PyValueError::new_err("Lance mode is not enabled"))?;
            let dur = std::time::Duration::from_secs(older_than_secs);
            mgr.cleanup_old_versions(collection_name, dur)
                .map_err(|e| PyRuntimeError::new_err(format!("cleanup_old_versions failed: {}", e)))
        }
        #[cfg(not(feature = "lance-storage"))]
        {
            let _ = (collection_name, older_than_secs);
            Err(PyValueError::new_err("Lance storage feature is not enabled at compile time"))
        }
    }

    /// Get the Arrow schema of a Lance collection.
    ///
    /// Returns a list of dicts with "name" (str) and "data_type" (str).
    fn lance_get_schema(&self, py: Python<'_>, collection_name: &str) -> PyResult<Vec<PyObject>> {
        #[cfg(feature = "lance-storage")]
        {
            let mgr = self.lance_manager.as_ref()
                .ok_or_else(|| PyValueError::new_err("Lance mode is not enabled"))?;
            let fields = mgr.get_schema(collection_name)
                .map_err(|e| PyRuntimeError::new_err(format!("get_schema failed: {}", e)))?;
            let mut out = Vec::with_capacity(fields.len());
            for f in &fields {
                let dict = PyDict::new_bound(py);
                dict.set_item("name", f.name().to_string())?;
                dict.set_item("data_type", format!("{:?}", f.data_type()))?;
                out.push(dict.into_py(py));
            }
            Ok(out)
        }
        #[cfg(not(feature = "lance-storage"))]
        {
            let _ = (py, collection_name);
            Err(PyValueError::new_err("Lance storage feature is not enabled at compile time"))
        }
    }

    /// Drop columns from a Lance collection.
    fn lance_drop_columns(&self, collection_name: &str, columns: Vec<String>) -> PyResult<()> {
        #[cfg(feature = "lance-storage")]
        {
            let mgr = self.lance_manager.as_ref()
                .ok_or_else(|| PyValueError::new_err("Lance mode is not enabled"))?;
            let col_refs: Vec<&str> = columns.iter().map(|s| s.as_str()).collect();
            mgr.drop_columns(collection_name, &col_refs)
                .map_err(|e| PyRuntimeError::new_err(format!("drop_columns failed: {}", e)))
        }
        #[cfg(not(feature = "lance-storage"))]
        {
            let _ = (collection_name, columns);
            Err(PyValueError::new_err("Lance storage feature is not enabled at compile time"))
        }
    }

    /// Rename a column in a Lance collection.
    fn lance_rename_column(&self, collection_name: &str, old_name: &str, new_name: &str) -> PyResult<()> {
        #[cfg(feature = "lance-storage")]
        {
            let mgr = self.lance_manager.as_ref()
                .ok_or_else(|| PyValueError::new_err("Lance mode is not enabled"))?;
            mgr.rename_column(collection_name, old_name, new_name)
                .map_err(|e| PyRuntimeError::new_err(format!("rename_column failed: {}", e)))
        }
        #[cfg(not(feature = "lance-storage"))]
        {
            let _ = (collection_name, old_name, new_name);
            Err(PyValueError::new_err("Lance storage feature is not enabled at compile time"))
        }
    }

    fn __repr__(&self) -> String {
        format!(
            "AiFastDb(path='{}', collections={})",
            self.path.display(),
            self.store.list_collections().len()
        )
    }

    fn __len__(&self) -> PyResult<usize> {
        self.document_count()
    }
}

// Private helper methods (not exposed to Python)
impl PyAiFastDb {
    fn _get_default_collection(&self) -> PyResult<Arc<Collection>> {
        self.store.get_default_collection(self.default_dimension)
            .map_err(|e| PyOSError::new_err(format!("Default collection error: {}", e)))
    }

    fn _remember_impl(
        py: Python<'_>,
        collection: &Arc<Collection>,
        id: String,
        content: String,
        content_type: &str,
        tags: Option<Vec<String>>,
        metadata: Option<&Bound<'_, PyDict>>,
        embedding: Option<Vec<f32>>,
        importance: f64,
        parent_id: Option<String>,
    ) -> PyResult<()> {
        let mut node = MemoryNode::new(content)
            .with_id(id)
            .with_content_type(parse_content_type(content_type))
            .with_importance(importance);

        if let Some(tags) = tags {
            node = node.with_tags(tags);
        }
        if let Some(emb) = embedding {
            node = node.with_embedding(emb);
        }
        if let Some(pid) = parent_id {
            node = node.with_parent(pid);
        }
        if let Some(meta) = metadata {
            for (key, val) in meta.iter() {
                let key_str: String = key.extract()?;
                let json_val = py_to_json_value(&val)?;
                node = node.with_metadata(key_str, json_val);
            }
        }

        // Store in doc store
        collection.docs.put(node.clone())
            .map_err(|e| PyOSError::new_err(format!("Write failed: {}", e)))?;

        // If has embedding, add to HNSW index
        if let Some(ref emb) = node.embedding {
            collection.hnsw.insert(&node.id, emb.to_vec())
                .map_err(|e| PyRuntimeError::new_err(format!("HNSW insert failed: {}", e)))?;
        }

        Ok(())
    }

    fn _get_impl(py: Python<'_>, collection: &Arc<Collection>, id: &str) -> PyResult<Option<PyObject>> {
        match collection.docs.get(id).ok().flatten() {
            Some(node) => Ok(Some(memory_node_to_dict(py, &node)?)),
            None => Ok(None),
        }
    }

    fn _forget_impl(collection: &Arc<Collection>, id: &str) -> PyResult<bool> {
        Ok(collection.docs.delete(id).ok().flatten().is_some())
    }

    fn _touch_impl(collection: &Arc<Collection>, id: &str) -> PyResult<bool> {
        Ok(collection.docs.record_access(id).is_ok())
    }

    fn _recall_impl(
        py: Python<'_>,
        collection: &Arc<Collection>,
        query_text: Option<String>,
        embedding: Option<Vec<f32>>,
        limit: usize,
        min_importance: Option<f64>,
        use_rrf: bool,
        k: f64,
    ) -> PyResult<Vec<PyObject>> {
        if query_text.is_none() && embedding.is_none() {
            return Err(PyValueError::new_err(
                "At least one of 'query' (text) or 'embedding' (vector) must be provided"
            ));
        }

        let engine = HybridQueryEngine::with_backend(
            Arc::clone(&collection.hnsw),
            Arc::clone(&collection.docs),
        );

        let mut hq = HybridQuery::empty().limit(limit);

        if let Some(ref text) = query_text {
            hq = hq.with_text(text.clone(), 1.0);
        }
        if let Some(ref emb) = embedding {
            hq = hq.with_vector(emb.clone(), 1.0);
        }
        if let Some(min_imp) = min_importance {
            hq = hq.min_importance(min_imp);
        }

        let results = if use_rrf {
            engine.query_rrf(&hq, k)
        } else {
            engine.query(&hq)
        };

        let mut py_results = Vec::with_capacity(results.len());
        for result in &results {
            let dict = PyDict::new_bound(py);
            dict.set_item("node", memory_node_to_dict(py, &result.node)?)?;
            dict.set_item("score", result.score)?;

            let breakdown = PyDict::new_bound(py);
            breakdown.set_item("vector_score", result.score_breakdown.vector_score)?;
            breakdown.set_item("text_score", result.score_breakdown.text_score)?;
            breakdown.set_item("importance_score", result.score_breakdown.importance_score)?;
            breakdown.set_item("recency_score", result.score_breakdown.recency_score)?;
            dict.set_item("score_breakdown", breakdown)?;

            py_results.push(dict.into_py(py));
        }

        Ok(py_results)
    }
}

// ---- PyCollection class (T24.8) ----

/// A single collection in the AiFastDb database.
///
/// Obtained via `db.collection("name")`.
#[pyclass(name = "Collection")]
pub struct PyCollection {
    inner: Arc<Collection>,
    name: String,
}

#[pymethods]
impl PyCollection {
    /// Collection name.
    #[getter]
    fn name(&self) -> &str {
        &self.name
    }

    /// Number of documents in this collection.
    fn document_count(&self) -> usize {
        self.inner.document_count()
    }

    /// Number of vectors in this collection.
    fn vector_count(&self) -> usize {
        self.inner.vector_count()
    }

    /// Store a memory node.
    #[pyo3(signature = (id, content, *, content_type="text", tags=None, metadata=None, embedding=None, importance=1.0, parent_id=None))]
    fn remember(
        &self,
        py: Python<'_>,
        id: String,
        content: String,
        content_type: &str,
        tags: Option<Vec<String>>,
        metadata: Option<&Bound<'_, PyDict>>,
        embedding: Option<Vec<f32>>,
        importance: f64,
        parent_id: Option<String>,
    ) -> PyResult<()> {
        PyAiFastDb::_remember_impl(py, &self.inner, id, content, content_type, tags, metadata, embedding, importance, parent_id)
    }

    /// Retrieve a memory node by ID.
    fn get(&self, py: Python<'_>, id: &str) -> PyResult<Option<PyObject>> {
        PyAiFastDb::_get_impl(py, &self.inner, id)
    }

    /// Delete a memory node by ID.
    fn forget(&self, id: &str) -> PyResult<bool> {
        PyAiFastDb::_forget_impl(&self.inner, id)
    }

    /// Record an access.
    fn touch(&self, id: &str) -> PyResult<bool> {
        PyAiFastDb::_touch_impl(&self.inner, id)
    }

    /// Search for similar memories.
    #[pyo3(signature = (*, query=None, embedding=None, limit=10, min_importance=None))]
    fn recall(
        &self,
        py: Python<'_>,
        query: Option<String>,
        embedding: Option<Vec<f32>>,
        limit: usize,
        min_importance: Option<f64>,
    ) -> PyResult<Vec<PyObject>> {
        PyAiFastDb::_recall_impl(py, &self.inner, query, embedding, limit, min_importance, false, 60.0)
    }

    /// Search using RRF.
    #[pyo3(signature = (*, query=None, embedding=None, limit=10, min_importance=None, k=60.0))]
    fn recall_rrf(
        &self,
        py: Python<'_>,
        query: Option<String>,
        embedding: Option<Vec<f32>>,
        limit: usize,
        min_importance: Option<f64>,
        k: f64,
    ) -> PyResult<Vec<PyObject>> {
        PyAiFastDb::_recall_impl(py, &self.inner, query, embedding, limit, min_importance, true, k)
    }

    /// Get all records.
    fn recall_all(&self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
        let nodes = self.inner.docs.all().unwrap_or_default();
        let mut results = Vec::with_capacity(nodes.len());
        for node in &nodes {
            results.push(memory_node_to_dict(py, &node)?);
        }
        Ok(results)
    }

    /// Flush this collection to disk.
    fn flush(&self) -> PyResult<()> {
        self.inner.flush()
            .map_err(|e| PyOSError::new_err(format!("Flush failed: {}", e)))
    }

    fn __repr__(&self) -> String {
        format!(
            "Collection(name='{}', documents={}, vectors={})",
            self.name,
            self.inner.document_count(),
            self.inner.vector_count()
        )
    }

    fn __len__(&self) -> usize {
        self.inner.document_count()
    }
}

/// Register database classes into the Python module.
pub fn register_database(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyAiFastDb>()?;
    m.add_class::<PyCollection>()?;
    Ok(())
}
