mr-ability 0.8.0

//! # 记忆存储实现
//!
//! [`MemoryStore`] 基于 [`RocksDBStore`] 实现 [`MemoryStorage`] trait，
//! 提供记忆的完整生命周期管理。
//!
//! ## 存储结构
//!
//! - **memories 列族**：UUID → JSON 序列化的 `Memory`
//! - **by_tag 列族**：`tag:uuid` → UUID（标签索引）
//! - **deleted 列族**：UUID → UUID（软删除索引）
//! - **importance 列族**：UUID → 重要性分数字符串
//!
//! ## 软删除机制
//!
//! 首次删除时标记 `is_deleted=true` 并记录 `deleted_at`，
//! 第二次删除（硬删除）才真正从存储中移除。

use super::rocksdb::RocksDBStore;
use super::traits::MemoryStorage;
use anyhow::{Context, Result};
use async_trait::async_trait;
use chrono::Utc;
use mr_common::{Memory, MemoryType};
use uuid::Uuid;

/// 基于 RocksDB 的记忆存储实现。
pub struct MemoryStore {
    rocksdb: std::sync::Arc<RocksDBStore>,
}

impl MemoryStore {
    /// 创建记忆存储实例。
    pub fn new(rocksdb: std::sync::Arc<RocksDBStore>) -> Self {
        Self { rocksdb }
    }

    /// 生成记忆主键（UUID 字符串的字节表示）。
    fn memory_key(id: &Uuid) -> Vec<u8> {
        id.to_string().into_bytes()
    }

    /// 生成标签索引键（`tag:uuid` 格式）。
    fn tag_key(tag: &str, id: &Uuid) -> Vec<u8> {
        format!("{}:{}", tag, id).into_bytes()
    }

    /// JSON 序列化记忆。
    fn serialize_memory(memory: &Memory) -> Result<Vec<u8>> {
        serde_json::to_vec(memory).context("Failed to serialize memory")
    }

    /// JSON 反序列化记忆。
    fn deserialize_memory(data: &[u8]) -> Result<Memory> {
        serde_json::from_slice(data).context("Failed to deserialize memory")
    }

    /// 从迭代器中收集记忆，支持跳过已删除项。
    fn collect_memories_from_iter(
        mut iter: rocksdb::DBRawIterator<'_>,
        limit: usize,
        filter_deleted: bool,
    ) -> Vec<Memory> {
        let mut memories = Vec::new();
        let mut count = 0;

        iter.seek_to_first();
        while iter.valid() && count < limit {
            if let Some(value) = iter.value() {
                if let Ok(memory) = Self::deserialize_memory(value) {
                    if !filter_deleted || !memory.is_deleted {
                        memories.push(memory);
                        count += 1;
                    }
                }
            }
            iter.next();
        }

        memories
    }
}

#[async_trait]
impl MemoryStorage for MemoryStore {
    async fn save(&self, memory: &Memory) -> Result<()> {
        let id_key = Self::memory_key(&memory.id);
        let data = Self::serialize_memory(memory)?;

        let cf_memories = self.rocksdb.cf_memories()?;
        self.rocksdb.put_cf(cf_memories, &id_key, &data)?;

        for tag in &memory.tags {
            let tag_key = Self::tag_key(tag, &memory.id);
            let cf_by_tag = self.rocksdb.cf_by_tag()?;
            self.rocksdb.put_cf(cf_by_tag, &tag_key, &id_key)?;
        }

        let importance_key = Self::memory_key(&memory.id);
        let importance_data = memory.importance.to_string().into_bytes();
        let cf_importance = self.rocksdb.cf_importance()?;
        self.rocksdb
            .put_cf(cf_importance, &importance_key, &importance_data)?;

        Ok(())
    }

    async fn get(&self, id: &Uuid) -> Result<Option<Memory>> {
        let id_key = Self::memory_key(id);
        let cf_memories = self.rocksdb.cf_memories()?;

        match self.rocksdb.get_cf(cf_memories, &id_key)? {
            Some(bytes) => {
                let memory = Self::deserialize_memory(&bytes)?;
                Ok(Some(memory))
            }
            None => Ok(None),
        }
    }

    async fn update(&self, memory: &Memory) -> Result<()> {
        self.save(memory).await
    }

    async fn delete(&self, id: &Uuid) -> Result<bool> {
        let memory = self.get(id).await?;

        match memory {
            Some(mem) => {
                if mem.is_deleted {
                    let id_key = Self::memory_key(id);
                    let cf_memories = self.rocksdb.cf_memories()?;
                    self.rocksdb.delete_cf(cf_memories, &id_key)?;

                    let cf_deleted = self.rocksdb.cf_deleted()?;
                    self.rocksdb.delete_cf(cf_deleted, &id_key)?;

                    for tag in &mem.tags {
                        let tag_key = Self::tag_key(tag, id);
                        let cf_by_tag = self.rocksdb.cf_by_tag()?;
                        self.rocksdb.delete_cf(cf_by_tag, &tag_key)?;
                    }

                    Ok(true)
                } else {
                    let mut mem_deleted = mem;
                    mem_deleted.is_deleted = true;
                    mem_deleted.deleted_at = Some(Utc::now());
                    self.update(&mem_deleted).await?;

                    let id_key = Self::memory_key(id);
                    let cf_deleted = self.rocksdb.cf_deleted()?;
                    self.rocksdb.put_cf(cf_deleted, &id_key, &id_key)?;

                    Ok(false)
                }
            }
            None => Ok(false),
        }
    }

    async fn restore(&self, id: &Uuid) -> Result<Option<Memory>> {
        let memory = self.get(id).await?;

        match memory {
            Some(mut mem) if mem.is_deleted => {
                mem.is_deleted = false;
                mem.deleted_at = None;
                self.save(&mem).await?;

                // 清理软删除索引（与 delete 对称），保证 count() 统计一致
                let id_key = Self::memory_key(id);
                let cf_deleted = self.rocksdb.cf_deleted()?;
                self.rocksdb.delete_cf(cf_deleted, &id_key)?;

                Ok(Some(mem))
            }
            _ => Ok(None),
        }
    }

    async fn list(&self, limit: usize) -> Result<Vec<Memory>> {
        let cf_memories = self.rocksdb.cf_memories()?;
        let iter = self.rocksdb.iter_cf(cf_memories);

        Ok(Self::collect_memories_from_iter(iter, limit, true))
    }

    async fn list_by_type(&self, memory_type: MemoryType, limit: usize) -> Result<Vec<Memory>> {
        let all = self.list(limit * 10).await?;
        Ok(all
            .into_iter()
            .filter(|m| m.memory_type == memory_type)
            .take(limit)
            .collect())
    }

    async fn list_by_tag(&self, tag: &str, limit: usize) -> Result<Vec<Memory>> {
        let cf_by_tag = self.rocksdb.cf_by_tag()?;
        let mut iter = self.rocksdb.iter_cf(cf_by_tag);

        let prefix = format!("{}:", tag).into_bytes();
        let mut memories = Vec::new();
        let mut count = 0;

        iter.seek_to_first();
        while iter.valid() && count < limit {
            if let Some(key) = iter.key() {
                if key.starts_with(&prefix) {
                    if let Some(value) = iter.value() {
                        let id_str = String::from_utf8_lossy(value);
                        if let Ok(id) = Uuid::parse_str(&id_str) {
                            if let Some(memory) = self.get(&id).await? {
                                if !memory.is_deleted {
                                    memories.push(memory);
                                    count += 1;
                                }
                            }
                        }
                    }
                }
            }
            iter.next();
        }

        Ok(memories)
    }

    async fn list_by_importance(&self, min: f32, max: f32) -> Result<Vec<Memory>> {
        let cf_importance = self.rocksdb.cf_importance()?;
        let mut iter = self.rocksdb.iter_cf(cf_importance);

        let mut memories = Vec::new();

        iter.seek_to_first();
        while iter.valid() {
            if let Some(key) = iter.key() {
                if let Some(value) = iter.value() {
                    let importance_str = String::from_utf8_lossy(value);
                    if let Ok(importance) = importance_str.parse::<f32>() {
                        if importance >= min && importance <= max {
                            let id_str = String::from_utf8_lossy(key);
                            if let Ok(id) = Uuid::parse_str(&id_str) {
                                if let Some(memory) = self.get(&id).await? {
                                    if !memory.is_deleted {
                                        memories.push(memory);
                                    }
                                }
                            }
                        }
                    }
                }
            }
            iter.next();
        }

        Ok(memories)
    }

    async fn list_deleted(&self) -> Result<Vec<Memory>> {
        let cf_deleted = self.rocksdb.cf_deleted()?;
        let mut iter = self.rocksdb.iter_cf(cf_deleted);

        let mut memories = Vec::new();

        iter.seek_to_first();
        while iter.valid() {
            if let Some(key) = iter.key() {
                let id_str = String::from_utf8_lossy(key);
                if let Ok(id) = Uuid::parse_str(&id_str) {
                    if let Some(memory) = self.get(&id).await? {
                        if memory.is_deleted {
                            memories.push(memory);
                        }
                    }
                }
            }
            iter.next();
        }

        Ok(memories)
    }

    async fn count(&self) -> Result<usize> {
        // 未删除记忆数 = memories 列族 key 数 - deleted 列族 key 数（软删除索引）
        // 仅迭代 key 计数，不反序列化，避免每次调用全量 list
        let cf_memories = self.rocksdb.cf_memories()?;
        let cf_deleted = self.rocksdb.cf_deleted()?;
        let total = self.rocksdb.count_cf_keys(cf_memories)?;
        let deleted = self.rocksdb.count_cf_keys(cf_deleted)?;
        Ok(total.saturating_sub(deleted))
    }

    async fn count_deleted(&self) -> Result<usize> {
        let deleted = self.list_deleted().await?;
        Ok(deleted.len())
    }

    async fn list_by_project(&self, project_id: &Uuid) -> Result<Vec<Memory>> {
        let all = self.list(10000).await?;
        Ok(all
            .into_iter()
            .filter(|m| m.project_id == Some(*project_id))
            .collect())
    }

    async fn get_chunks_by_group(&self, chunk_group_id: &Uuid) -> Result<Vec<Memory>> {
        let all = self.list(10000).await?;
        Ok(all
            .into_iter()
            .filter(|m| m.chunk_group_id == Some(*chunk_group_id))
            .collect())
    }

    async fn list_older_than(&self, cutoff: &str, limit: usize) -> Result<Vec<Memory>> {
        let cutoff_time = chrono::DateTime::parse_from_rfc3339(cutoff)
            .context("Invalid cutoff time format")?
            .with_timezone(&Utc);

        let all = self.list(limit * 10).await?;
        Ok(all
            .into_iter()
            .filter(|m| m.created_at < cutoff_time)
            .take(limit)
            .collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[tokio::test]
    async fn test_memory_save_and_get() {
        let dir = tempdir().unwrap();
        let rocksdb = RocksDBStore::open(dir.path()).unwrap();
        let store = MemoryStore::new(std::sync::Arc::new(rocksdb));

        let memory = Memory::new("test content".to_string(), MemoryType::Knowledge)
            .with_tags(vec!["tag1".to_string()]);

        store.save(&memory).await.unwrap();

        let retrieved = store.get(&memory.id).await.unwrap();
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().content, "test content");
    }

    #[tokio::test]
    async fn test_memory_list() {
        let dir = tempdir().unwrap();
        let rocksdb = RocksDBStore::open(dir.path()).unwrap();
        let store = MemoryStore::new(std::sync::Arc::new(rocksdb));

        for i in 0..5 {
            let memory = Memory::new(format!("content {}", i), MemoryType::Conversation);
            store.save(&memory).await.unwrap();
        }

        let memories = store.list(10).await.unwrap();
        assert_eq!(memories.len(), 5);
    }

    #[tokio::test]
    async fn test_memory_list_by_tag() {
        let dir = tempdir().unwrap();
        let rocksdb = RocksDBStore::open(dir.path()).unwrap();
        let store = MemoryStore::new(std::sync::Arc::new(rocksdb));

        let m1 = Memory::new("content 1".to_string(), MemoryType::Knowledge)
            .with_tags(vec!["rust".to_string()]);
        let m2 = Memory::new("content 2".to_string(), MemoryType::Knowledge)
            .with_tags(vec!["python".to_string()]);

        store.save(&m1).await.unwrap();
        store.save(&m2).await.unwrap();

        let rust_memories = store.list_by_tag("rust", 10).await.unwrap();
        assert_eq!(rust_memories.len(), 1);
    }

    #[tokio::test]
    async fn test_memory_soft_delete() {
        let dir = tempdir().unwrap();
        let rocksdb = RocksDBStore::open(dir.path()).unwrap();
        let store = MemoryStore::new(std::sync::Arc::new(rocksdb));

        let memory = Memory::new("test".to_string(), MemoryType::Conversation);
        store.save(&memory).await.unwrap();

        let deleted = store.delete(&memory.id).await.unwrap();
        assert!(!deleted);

        let retrieved = store.get(&memory.id).await.unwrap();
        assert!(retrieved.is_some());
        assert!(retrieved.unwrap().is_deleted);
    }

    #[tokio::test]
    async fn test_memory_restore_clears_deleted_index() {
        let dir = tempdir().unwrap();
        let rocksdb = RocksDBStore::open(dir.path()).unwrap();
        let store = MemoryStore::new(std::sync::Arc::new(rocksdb));

        let memory = Memory::new("test".to_string(), MemoryType::Knowledge);
        store.save(&memory).await.unwrap();

        // 软删 → count 减少
        store.delete(&memory.id).await.unwrap();
        assert_eq!(store.count().await.unwrap(), 0);
        assert_eq!(store.count_deleted().await.unwrap(), 1);

        // 恢复 → is_deleted=false 且软删索引清理，count 恢复
        let restored = store.restore(&memory.id).await.unwrap();
        assert!(restored.is_some());
        assert!(!restored.unwrap().is_deleted);
        assert_eq!(store.count().await.unwrap(), 1);
        assert_eq!(store.count_deleted().await.unwrap(), 0);

        // 再次 restore：非删除状态返回 None
        assert!(store.restore(&memory.id).await.unwrap().is_none());
    }
}