zclaw_openfang/crates/zclaw-growth/src/viking_adapter.rs

//! OpenViking Adapter - Interface to the OpenViking memory system
//!
//! This module provides the `VikingAdapter` which wraps the OpenViking
//! context database for storing and retrieving agent memories.

use crate::types::MemoryEntry;
use async_trait::async_trait;
use serde::{de::DeserializeOwned, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use zclaw_types::Result;

/// Search options for find operations
#[derive(Debug, Clone, Default)]
pub struct FindOptions {
    /// Scope to search within (URI prefix)
    pub scope: Option<String>,
    /// Maximum results to return
    pub limit: Option<usize>,
    /// Minimum similarity threshold
    pub min_similarity: Option<f32>,
}

/// VikingStorage trait - core storage operations (dyn-compatible)
#[async_trait]
pub trait VikingStorage: Send + Sync {
    /// Store a memory entry
    async fn store(&self, entry: &MemoryEntry) -> Result<()>;

    /// Get a memory entry by URI
    async fn get(&self, uri: &str) -> Result<Option<MemoryEntry>>;

    /// Find memories by query with options
    async fn find(&self, query: &str, options: FindOptions) -> Result<Vec<MemoryEntry>>;

    /// Find memories by URI prefix
    async fn find_by_prefix(&self, prefix: &str) -> Result<Vec<MemoryEntry>>;

    /// Delete a memory by URI
    async fn delete(&self, uri: &str) -> Result<()>;

    /// Store metadata as JSON string
    async fn store_metadata_json(&self, key: &str, json: &str) -> Result<()>;

    /// Get metadata as JSON string
    async fn get_metadata_json(&self, key: &str) -> Result<Option<String>>;
}

/// OpenViking adapter implementation
#[derive(Clone)]
pub struct VikingAdapter {
    /// Storage backend
    backend: Arc<dyn VikingStorage>,
}

impl VikingAdapter {
    /// Create a new Viking adapter with a storage backend
    pub fn new(backend: Arc<dyn VikingStorage>) -> Self {
        Self { backend }
    }

    /// Create with in-memory storage (for testing)
    pub fn in_memory() -> Self {
        Self {
            backend: Arc::new(InMemoryStorage::new()),
        }
    }

    /// Store a memory entry
    pub async fn store(&self, entry: &MemoryEntry) -> Result<()> {
        self.backend.store(entry).await
    }

    /// Get a memory entry by URI
    pub async fn get(&self, uri: &str) -> Result<Option<MemoryEntry>> {
        self.backend.get(uri).await
    }

    /// Find memories by query
    pub async fn find(&self, query: &str, options: FindOptions) -> Result<Vec<MemoryEntry>> {
        self.backend.find(query, options).await
    }

    /// Find memories by URI prefix
    pub async fn find_by_prefix(&self, prefix: &str) -> Result<Vec<MemoryEntry>> {
        self.backend.find_by_prefix(prefix).await
    }

    /// Delete a memory
    pub async fn delete(&self, uri: &str) -> Result<()> {
        self.backend.delete(uri).await
    }

    /// Store metadata (typed)
    pub async fn store_metadata<T: Serialize>(&self, key: &str, value: &T) -> Result<()> {
        let json = serde_json::to_string(value)?;
        self.backend.store_metadata_json(key, &json).await
    }

    /// Get metadata (typed)
    pub async fn get_metadata<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>> {
        match self.backend.get_metadata_json(key).await? {
            Some(json) => {
                let value: T = serde_json::from_str(&json)?;
                Ok(Some(value))
            }
            None => Ok(None),
        }
    }
}

/// In-memory storage backend (for testing and development)
pub struct InMemoryStorage {
    memories: std::sync::RwLock<HashMap<String, MemoryEntry>>,
    metadata: std::sync::RwLock<HashMap<String, String>>,
}

impl InMemoryStorage {
    /// Create a new in-memory storage
    pub fn new() -> Self {
        Self {
            memories: std::sync::RwLock::new(HashMap::new()),
            metadata: std::sync::RwLock::new(HashMap::new()),
        }
    }
}

impl Default for InMemoryStorage {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl VikingStorage for InMemoryStorage {
    async fn store(&self, entry: &MemoryEntry) -> Result<()> {
        let mut memories = self.memories.write().unwrap();
        memories.insert(entry.uri.clone(), entry.clone());
        Ok(())
    }

    async fn get(&self, uri: &str) -> Result<Option<MemoryEntry>> {
        let memories = self.memories.read().unwrap();
        Ok(memories.get(uri).cloned())
    }

    async fn find(&self, query: &str, options: FindOptions) -> Result<Vec<MemoryEntry>> {
        let memories = self.memories.read().unwrap();

        let mut results: Vec<MemoryEntry> = memories
            .values()
            .filter(|entry| {
                // Apply scope filter
                if let Some(ref scope) = options.scope {
                    if !entry.uri.starts_with(scope) {
                        return false;
                    }
                }

                // Simple text matching (in real implementation, use semantic search)
                if !query.is_empty() {
                    let query_lower = query.to_lowercase();
                    let content_lower = entry.content.to_lowercase();
                    let keywords_match = entry.keywords.iter().any(|k| k.to_lowercase().contains(&query_lower));

                    content_lower.contains(&query_lower) || keywords_match
                } else {
                    true
                }
            })
            .cloned()
            .collect();

        // Sort by importance and access count
        results.sort_by(|a, b| {
            b.importance
                .cmp(&a.importance)
                .then_with(|| b.access_count.cmp(&a.access_count))
        });

        // Apply limit
        if let Some(limit) = options.limit {
            results.truncate(limit);
        }

        Ok(results)
    }

    async fn find_by_prefix(&self, prefix: &str) -> Result<Vec<MemoryEntry>> {
        let memories = self.memories.read().unwrap();

        let results: Vec<MemoryEntry> = memories
            .values()
            .filter(|entry| entry.uri.starts_with(prefix))
            .cloned()
            .collect();

        Ok(results)
    }

    async fn delete(&self, uri: &str) -> Result<()> {
        let mut memories = self.memories.write().unwrap();
        memories.remove(uri);
        Ok(())
    }

    async fn store_metadata_json(&self, key: &str, json: &str) -> Result<()> {
        let mut metadata = self.metadata.write().unwrap();
        metadata.insert(key.to_string(), json.to_string());
        Ok(())
    }

    async fn get_metadata_json(&self, key: &str) -> Result<Option<String>> {
        let metadata = self.metadata.read().unwrap();
        Ok(metadata.get(key).cloned())
    }
}

/// OpenViking levels for storage
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VikingLevel {
    /// L0: Raw data (original content)
    L0,
    /// L1: Summarized content
    L1,
    /// L2: Keywords and metadata
    L2,
}

impl std::fmt::Display for VikingLevel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            VikingLevel::L0 => write!(f, "L0"),
            VikingLevel::L1 => write!(f, "L1"),
            VikingLevel::L2 => write!(f, "L2"),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::MemoryType;

    #[tokio::test]
    async fn test_in_memory_storage_store_and_get() {
        let storage = InMemoryStorage::new();
        let entry = MemoryEntry::new(
            "test-agent",
            MemoryType::Preference,
            "style",
            "test content".to_string(),
        );

        storage.store(&entry).await.unwrap();
        let retrieved = storage.get(&entry.uri).await.unwrap();

        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().content, "test content");
    }

    #[tokio::test]
    async fn test_in_memory_storage_find() {
        let storage = InMemoryStorage::new();

        let entry1 = MemoryEntry::new(
            "agent-1",
            MemoryType::Knowledge,
            "rust",
            "Rust programming tips".to_string(),
        );
        let entry2 = MemoryEntry::new(
            "agent-1",
            MemoryType::Knowledge,
            "python",
            "Python programming tips".to_string(),
        );

        storage.store(&entry1).await.unwrap();
        storage.store(&entry2).await.unwrap();

        let results = storage
            .find(
                "Rust",
                FindOptions {
                    scope: Some("agent://agent-1".to_string()),
                    limit: Some(10),
                    min_similarity: None,
                },
            )
            .await
            .unwrap();

        assert_eq!(results.len(), 1);
        assert!(results[0].content.contains("Rust"));
    }

    #[tokio::test]
    async fn test_in_memory_storage_delete() {
        let storage = InMemoryStorage::new();
        let entry = MemoryEntry::new(
            "test-agent",
            MemoryType::Preference,
            "style",
            "test".to_string(),
        );

        storage.store(&entry).await.unwrap();
        storage.delete(&entry.uri).await.unwrap();

        let retrieved = storage.get(&entry.uri).await.unwrap();
        assert!(retrieved.is_none());
    }

    #[tokio::test]
    async fn test_metadata_storage() {
        let storage = InMemoryStorage::new();

        #[derive(Serialize, serde::Deserialize)]
        struct TestData {
            value: String,
        }

        let data = TestData {
            value: "test".to_string(),
        };

        storage.store_metadata_json("test-key", &serde_json::to_string(&data).unwrap()).await.unwrap();
        let json = storage.get_metadata_json("test-key").await.unwrap();

        assert!(json.is_some());
        let retrieved: TestData = serde_json::from_str(&json.unwrap()).unwrap();
        assert_eq!(retrieved.value, "test");
    }

    #[tokio::test]
    async fn test_viking_adapter_typed_metadata() {
        let adapter = VikingAdapter::in_memory();

        #[derive(Serialize, serde::Deserialize)]
        struct TestData {
            value: String,
        }

        let data = TestData {
            value: "test".to_string(),
        };

        adapter.store_metadata("test-key", &data).await.unwrap();
        let retrieved: Option<TestData> = adapter.get_metadata("test-key").await.unwrap();

        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().value, "test");
    }

    #[test]
    fn test_viking_level_display() {
        assert_eq!(format!("{}", VikingLevel::L0), "L0");
        assert_eq!(format!("{}", VikingLevel::L1), "L1");
        assert_eq!(format!("{}", VikingLevel::L2), "L2");
    }
}