zclaw_openfang/desktop/src-tauri/src/viking_commands.rs

//! OpenViking Memory Storage - Native Rust Implementation
//!
//! Provides native Rust memory storage using SqliteStorage with TF-IDF semantic search.
//! This is a self-contained implementation that doesn't require external Python or CLI dependencies.
//!
//! Features:
//! - SQLite persistence with FTS5 full-text search
//! - TF-IDF semantic scoring
//! - Token budget control
//! - Automatic memory indexing

use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::OnceCell;
use zclaw_growth::{
    FindOptions, MemoryEntry, MemoryType, PromptInjector, RetrievalResult, SqliteStorage,
    VikingStorage,
};

// === Types ===

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct VikingStatus {
    pub available: bool,
    pub version: Option<String>,
    pub data_dir: Option<String>,
    pub error: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct VikingResource {
    pub uri: String,
    pub name: String,
    #[serde(rename = "type")]
    pub resource_type: String,
    pub size: Option<u64>,
    pub modified_at: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct VikingFindResult {
    pub uri: String,
    pub score: f64,
    pub content: String,
    pub level: String,
    pub overview: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct VikingGrepResult {
    pub uri: String,
    pub line: u32,
    pub content: String,
    pub match_start: u32,
    pub match_end: u32,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct VikingAddResult {
    pub uri: String,
    pub status: String,
}

// === Global Storage Instance ===

/// Global storage instance
static STORAGE: OnceCell<Arc<SqliteStorage>> = OnceCell::const_new();

/// Get the storage directory path
fn get_storage_dir() -> PathBuf {
    // Use platform-specific data directory
    if let Some(data_dir) = dirs::data_dir() {
        data_dir.join("zclaw").join("memories")
    } else {
        // Fallback to current directory
        PathBuf::from("./zclaw_data/memories")
    }
}

/// Initialize the storage (should be called once at startup)
pub async fn init_storage() -> Result<(), String> {
    let storage_dir = get_storage_dir();
    let db_path = storage_dir.join("memories.db");

    tracing::info!("[VikingCommands] Initializing storage at {:?}", db_path);

    let storage = SqliteStorage::new(&db_path)
        .await
        .map_err(|e| format!("Failed to initialize storage: {}", e))?;

    let _ = STORAGE.set(Arc::new(storage));

    tracing::info!("[VikingCommands] Storage initialized successfully");
    Ok(())
}

/// Get the storage instance (public for use by other modules)
pub async fn get_storage() -> Result<Arc<SqliteStorage>, String> {
    STORAGE
        .get()
        .cloned()
        .ok_or_else(|| "Storage not initialized. Call init_storage() first.".to_string())
}

/// Get storage directory for status
fn get_data_dir_string() -> Option<String> {
    get_storage_dir().to_str().map(|s| s.to_string())
}

// === Tauri Commands ===

/// Check if memory storage is available
#[tauri::command]
pub async fn viking_status() -> Result<VikingStatus, String> {
    match get_storage().await {
        Ok(storage) => {
            // Try a simple query to verify storage is working
            let _ = storage
                .find("", FindOptions::default())
                .await
                .map_err(|e| format!("Storage health check failed: {}", e))?;

            Ok(VikingStatus {
                available: true,
                version: Some("0.1.0-native".to_string()),
                data_dir: get_data_dir_string(),
                error: None,
            })
        }
        Err(e) => Ok(VikingStatus {
            available: false,
            version: None,
            data_dir: get_data_dir_string(),
            error: Some(e),
        }),
    }
}

/// Add a memory entry
#[tauri::command]
pub async fn viking_add(uri: String, content: String) -> Result<VikingAddResult, String> {
    let storage = get_storage().await?;

    // Parse URI to extract agent_id, memory_type, and category
    // Expected format: agent://{agent_id}/{type}/{category}
    let (agent_id, memory_type, category) = parse_uri(&uri)?;

    let entry = MemoryEntry::new(&agent_id, memory_type, &category, content);

    storage
        .store(&entry)
        .await
        .map_err(|e| format!("Failed to store memory: {}", e))?;

    Ok(VikingAddResult {
        uri,
        status: "added".to_string(),
    })
}

/// Add a memory with metadata
#[tauri::command]
pub async fn viking_add_with_metadata(
    uri: String,
    content: String,
    keywords: Vec<String>,
    importance: Option<u8>,
) -> Result<VikingAddResult, String> {
    let storage = get_storage().await?;

    let (agent_id, memory_type, category) = parse_uri(&uri)?;

    let mut entry = MemoryEntry::new(&agent_id, memory_type, &category, content);
    entry.keywords = keywords;

    if let Some(imp) = importance {
        entry.importance = imp.min(10).max(1);
    }

    storage
        .store(&entry)
        .await
        .map_err(|e| format!("Failed to store memory: {}", e))?;

    Ok(VikingAddResult {
        uri,
        status: "added".to_string(),
    })
}

/// Find memories by semantic search
#[tauri::command]
pub async fn viking_find(
    query: String,
    scope: Option<String>,
    limit: Option<usize>,
) -> Result<Vec<VikingFindResult>, String> {
    let storage = get_storage().await?;

    let options = FindOptions {
        scope,
        limit,
        min_similarity: Some(0.1),
    };

    let entries = storage
        .find(&query, options)
        .await
        .map_err(|e| format!("Failed to search memories: {}", e))?;

    Ok(entries
        .into_iter()
        .enumerate()
        .map(|(i, entry)| VikingFindResult {
            uri: entry.uri,
            score: 1.0 - (i as f64 * 0.1), // Simple scoring based on rank
            content: entry.content,
            level: "L1".to_string(),
            overview: None,
        })
        .collect())
}

/// Grep memories by pattern (uses FTS5)
#[tauri::command]
pub async fn viking_grep(
    pattern: String,
    uri: Option<String>,
    _case_sensitive: Option<bool>,
    limit: Option<usize>,
) -> Result<Vec<VikingGrepResult>, String> {
    let storage = get_storage().await?;

    let scope = uri.as_ref().and_then(|u| {
        // Extract agent scope from URI
        u.strip_prefix("agent://")
            .and_then(|s| s.split('/').next())
            .map(|agent| format!("agent://{}", agent))
    });

    let options = FindOptions {
        scope,
        limit,
        min_similarity: Some(0.05), // Lower threshold for grep
    };

    let entries = storage
        .find(&pattern, options)
        .await
        .map_err(|e| format!("Failed to grep memories: {}", e))?;

    Ok(entries
        .into_iter()
        .flat_map(|entry| {
            // Find matching lines
            entry
                .content
                .lines()
                .enumerate()
                .filter(|(_, line)| {
                    line.to_lowercase()
                        .contains(&pattern.to_lowercase())
                })
                .map(|(i, line)| VikingGrepResult {
                    uri: entry.uri.clone(),
                    line: (i + 1) as u32,
                    content: line.to_string(),
                    match_start: line.find(&pattern).unwrap_or(0) as u32,
                    match_end: (line.find(&pattern).unwrap_or(0) + pattern.len()) as u32,
                })
                .collect::<Vec<_>>()
        })
        .take(limit.unwrap_or(100))
        .collect())
}

/// List memories at a path
#[tauri::command]
pub async fn viking_ls(path: String) -> Result<Vec<VikingResource>, String> {
    let storage = get_storage().await?;

    let entries = storage
        .find_by_prefix(&path)
        .await
        .map_err(|e| format!("Failed to list memories: {}", e))?;

    Ok(entries
        .into_iter()
        .map(|entry| VikingResource {
            uri: entry.uri.clone(),
            name: entry
                .uri
                .rsplit('/')
                .next()
                .unwrap_or(&entry.uri)
                .to_string(),
            resource_type: entry.memory_type.to_string(),
            size: Some(entry.content.len() as u64),
            modified_at: Some(entry.last_accessed.to_rfc3339()),
        })
        .collect())
}

/// Read memory content
#[tauri::command]
pub async fn viking_read(uri: String, _level: Option<String>) -> Result<String, String> {
    let storage = get_storage().await?;

    let entry = storage
        .get(&uri)
        .await
        .map_err(|e| format!("Failed to read memory: {}", e))?;

    match entry {
        Some(e) => Ok(e.content),
        None => Err(format!("Memory not found: {}", uri)),
    }
}

/// Remove a memory
#[tauri::command]
pub async fn viking_remove(uri: String) -> Result<(), String> {
    let storage = get_storage().await?;

    storage
        .delete(&uri)
        .await
        .map_err(|e| format!("Failed to remove memory: {}", e))?;

    Ok(())
}

/// Get memory tree
#[tauri::command]
pub async fn viking_tree(path: String, _depth: Option<usize>) -> Result<serde_json::Value, String> {
    let storage = get_storage().await?;

    let entries = storage
        .find_by_prefix(&path)
        .await
        .map_err(|e| format!("Failed to get tree: {}", e))?;

    // Build a simple tree structure
    let mut tree = serde_json::Map::new();

    for entry in entries {
        let parts: Vec<&str> = entry.uri.split('/').collect();
        let mut current = &mut tree;

        for part in &parts[..parts.len().saturating_sub(1)] {
            if !current.contains_key(*part) {
                current.insert(
                    (*part).to_string(),
                    serde_json::json!({}),
                );
            }
            current = current
                .get_mut(*part)
                .and_then(|v| v.as_object_mut())
                .unwrap();
        }

        if let Some(last) = parts.last() {
            current.insert(
                (*last).to_string(),
                serde_json::json!({
                    "type": entry.memory_type.to_string(),
                    "importance": entry.importance,
                    "access_count": entry.access_count,
                }),
            );
        }
    }

    Ok(serde_json::Value::Object(tree))
}

/// Inject memories into prompt (for agent loop integration)
#[tauri::command]
pub async fn viking_inject_prompt(
    agent_id: String,
    base_prompt: String,
    user_input: String,
    max_tokens: Option<usize>,
) -> Result<String, String> {
    let storage = get_storage().await?;

    // Retrieve relevant memories
    let options = FindOptions {
        scope: Some(format!("agent://{}", agent_id)),
        limit: Some(10),
        min_similarity: Some(0.3),
    };

    let entries = storage
        .find(&user_input, options)
        .await
        .map_err(|e| format!("Failed to retrieve memories: {}", e))?;

    // Convert to RetrievalResult
    let mut result = RetrievalResult::default();
    for entry in entries {
        match entry.memory_type {
            MemoryType::Preference => result.preferences.push(entry),
            MemoryType::Knowledge => result.knowledge.push(entry),
            MemoryType::Experience => result.experience.push(entry),
            MemoryType::Session => {} // Skip session memories
        }
    }

    // Calculate tokens
    result.total_tokens = result.calculate_tokens();

    // Apply token budget
    let budget = max_tokens.unwrap_or(500);
    if result.total_tokens > budget {
        // Truncate by priority: preferences > knowledge > experience
        while result.total_tokens > budget && !result.experience.is_empty() {
            result.experience.pop();
            result.total_tokens = result.calculate_tokens();
        }
        while result.total_tokens > budget && !result.knowledge.is_empty() {
            result.knowledge.pop();
            result.total_tokens = result.calculate_tokens();
        }
        while result.total_tokens > budget && !result.preferences.is_empty() {
            result.preferences.pop();
            result.total_tokens = result.calculate_tokens();
        }
    }

    // Inject into prompt
    let injector = PromptInjector::new();
    Ok(injector.inject_with_format(&base_prompt, &result))
}

// === Helper Functions ===

/// Parse URI to extract components
fn parse_uri(uri: &str) -> Result<(String, MemoryType, String), String> {
    // Expected format: agent://{agent_id}/{type}/{category}
    let without_prefix = uri
        .strip_prefix("agent://")
        .ok_or_else(|| format!("Invalid URI format: {}", uri))?;

    let parts: Vec<&str> = without_prefix.splitn(3, '/').collect();

    if parts.len() < 3 {
        return Err(format!("Invalid URI format, expected agent://{{agent_id}}/{{type}}/{{category}}: {}", uri));
    }

    let agent_id = parts[0].to_string();
    let memory_type = MemoryType::parse(parts[1]);
    let category = parts[2].to_string();

    Ok((agent_id, memory_type, category))
}

// === Tests ===

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

    #[tokio::test]
    async fn test_parse_uri() {
        let (agent_id, memory_type, category) =
            parse_uri("agent://test-agent/preferences/style").unwrap();

        assert_eq!(agent_id, "test-agent");
        assert_eq!(memory_type, MemoryType::Preference);
        assert_eq!(category, "style");
    }

    #[test]
    fn test_invalid_uri() {
        assert!(parse_uri("invalid-uri").is_err());
        assert!(parse_uri("agent://only-agent").is_err());
    }
}