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,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct EmbeddingConfigResult {
    pub provider: String,
    pub configured: bool,
}

/// Industry keyword config received from the frontend (JSON string).
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct IndustryConfigPayload {
    pub id: String,
    pub name: String,
    pub keywords: Vec<String>,
    pub system_prompt: String,
}

// === Global Storage Instance ===

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

/// Get the storage directory path
pub 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, initializing on first access if needed
pub async fn get_storage() -> Result<Arc<SqliteStorage>, String> {
    if let Some(storage) = STORAGE.get() {
        return Ok(storage.clone());
    }

    // Attempt lazy initialization
    tracing::info!("[VikingCommands] Storage not yet initialized, attempting lazy init...");
    init_storage().await?;

    STORAGE
        .get()
        .cloned()
        .ok_or_else(|| "Storage initialization failed. Check logs for details.".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
// @connected
#[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
// @connected
#[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
// @connected
#[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
// @connected
#[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)| {
            // Use overview (L1) when available, full content otherwise (L2)
            let (content, level, overview) = if let Some(ref ov) = entry.overview {
                if !ov.is_empty() {
                    (ov.clone(), "L1".to_string(), None)
                } else {
                    (entry.content.clone(), "L2".to_string(), None)
                }
            } else {
                (entry.content.clone(), "L2".to_string(), None)
            };
            VikingFindResult {
                uri: entry.uri,
                score: 1.0 - (i as f64 * 0.1), // Simple scoring based on rank
                content,
                level,
                overview,
            }
        })
        .collect())
}

/// Grep memories by pattern (uses FTS5)
// @connected
#[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
// @connected
#[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
// @connected
#[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) => {
            // Support level-based content retrieval
            let content = match level.as_deref() {
                Some("L0") | Some("l0") => {
                    // L0: abstract_summary (keywords)
                    e.abstract_summary
                        .filter(|s| !s.is_empty())
                        .unwrap_or_else(|| {
                            // Fallback: first 50 chars of overview
                            e.overview
                                .as_ref()
                                .map(|ov| ov.chars().take(50).collect())
                                .unwrap_or_else(|| e.content.chars().take(50).collect())
                        })
                }
                Some("L1") | Some("l1") => {
                    // L1: overview (1-2 sentence summary)
                    e.overview
                        .filter(|s| !s.is_empty())
                        .unwrap_or_else(|| truncate_text(&e.content, 200))
                }
                _ => {
                    // L2 or default: full content
                    e.content
                }
            };
            Ok(content)
        }
        None => Err(format!("Memory not found: {}", uri)),
    }
}

/// Remove a memory
// @connected
#[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
// @connected
#[tauri::command]
pub async fn viking_tree(path: String, depth: Option<usize>) -> Result<serde_json::Value, String> {
    let max_depth = depth.unwrap_or(5);
    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 level = parts.len().saturating_sub(1);
        if level > max_depth {
            continue;
        }
        let mut current = &mut tree;

        for part in &parts[..level] {
            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)
// @connected
#[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 ===

/// Truncate text to approximately max_chars characters
fn truncate_text(text: &str, max_chars: usize) -> String {
    if text.chars().count() <= max_chars {
        text.to_string()
    } else {
        let truncated: String = text.chars().take(max_chars).collect();
        format!("{}...", truncated)
    }
}

/// 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))
}

/// Configure embedding for semantic memory search
/// Configures both SqliteStorage (VikingPanel) and PersistentMemoryStore (chat flow)
// @connected
#[tauri::command]
pub async fn viking_configure_embedding(
    provider: String,
    api_key: String,
    model: Option<String>,
    endpoint: Option<String>,
) -> Result<EmbeddingConfigResult, String> {
    let storage = get_storage().await?;

    // 1. Configure SqliteStorage (VikingPanel / VikingCommands)
    let config_viking = crate::llm::EmbeddingConfig {
        provider: provider.clone(),
        api_key: api_key.clone(),
        endpoint: endpoint.clone(),
        model: model.clone(),
    };

    let client_viking = crate::llm::EmbeddingClient::new(config_viking);
    let adapter = crate::embedding_adapter::TauriEmbeddingAdapter::new(client_viking);

    storage
        .configure_embedding(std::sync::Arc::new(adapter))
        .await
        .map_err(|e| format!("Failed to configure embedding: {}", e))?;

    // 2. Configure PersistentMemoryStore (chat flow)
    let config_memory = crate::llm::EmbeddingConfig {
        provider: provider.clone(),
        api_key,
        endpoint,
        model,
    };
    let client_memory = std::sync::Arc::new(crate::llm::EmbeddingClient::new(config_memory));

    let embed_fn: crate::memory::EmbedFn = {
        let client_arc = client_memory.clone();
        std::sync::Arc::new(move |text: &str| {
            let client = client_arc.clone();
            let text = text.to_string();
            Box::pin(async move {
                let response = client.embed(&text).await?;
                Ok(response.embedding)
            })
        })
    };

    crate::memory::configure_embedding_client(embed_fn);

    tracing::info!("[VikingCommands] Embedding configured with provider: {} (both storage systems)", provider);

    Ok(EmbeddingConfigResult {
        provider,
        configured: true,
    })
}

/// Configure summary driver for L0/L1 auto-generation
// @connected
#[tauri::command]
pub async fn viking_configure_summary_driver(
    endpoint: String,
    api_key: String,
    model: Option<String>,
) -> Result<bool, String> {
    let driver = crate::summarizer_adapter::TauriSummaryDriver::new(endpoint, api_key, model);
    crate::summarizer_adapter::configure_summary_driver(driver);

    tracing::info!("[VikingCommands] Summary driver configured");
    Ok(true)
}

/// Store a memory and optionally generate L0/L1 summaries in the background
// @connected
#[tauri::command]
pub async fn viking_store_with_summaries(
    uri: String,
    content: String,
) -> Result<VikingAddResult, String> {
    let storage = get_storage().await?;
    let (agent_id, memory_type, category) = parse_uri(&uri)?;

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

    // Store the entry immediately (L2 full content)
    storage
        .store(&entry)
        .await
        .map_err(|e| format!("Failed to store memory: {}", e))?;

    // Background: generate L0/L1 summaries if driver is configured
    if crate::summarizer_adapter::is_summary_driver_configured() {
        let entry_uri = entry.uri.clone();
        let storage_clone = storage.clone();

        tokio::spawn(async move {
            if let Some(driver) = crate::summarizer_adapter::get_summary_driver() {
                let (overview, abstract_summary) =
                    zclaw_growth::summarizer::generate_summaries(driver.as_ref(), &entry).await;

                if overview.is_some() || abstract_summary.is_some() {
                    // Update the entry with summaries
                    let updated = MemoryEntry {
                        overview,
                        abstract_summary,
                        ..entry
                    };

                    if let Err(e) = storage_clone.store(&updated).await {
                        tracing::debug!(
                            "[VikingCommands] Failed to update summaries for {}: {}",
                            entry_uri,
                            e
                        );
                    } else {
                        tracing::debug!(
                            "[VikingCommands] Updated L0/L1 summaries for {}",
                            entry_uri
                        );
                    }
                }
            }
        });
    }

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

// === Tests ===

// ---------------------------------------------------------------------------
// Industry Keywords Loader
// ---------------------------------------------------------------------------

/// Load industry keywords into the ButlerRouter middleware.
///
/// Called from the frontend after fetching industry configs from SaaS.
/// Updates the shared `industry_keywords` Arc on the Kernel, which the
/// ButlerRouterMiddleware reads automatically (same Arc instance).
#[tauri::command]
pub async fn viking_load_industry_keywords(
    kernel_state: tauri::State<'_, crate::kernel_commands::KernelState>,
    configs: String,
) -> Result<(), String> {
    let raw: Vec<IndustryConfigPayload> = serde_json::from_str(&configs)
        .map_err(|e| format!("Failed to parse industry configs: {}", e))?;

    let industry_configs: Vec<zclaw_runtime::IndustryKeywordConfig> = raw
        .into_iter()
        .map(|c| zclaw_runtime::IndustryKeywordConfig {
            id: c.id,
            name: c.name,
            keywords: c.keywords,
            system_prompt: c.system_prompt,
        })
        .collect();

    tracing::info!(
        "[viking_commands] Loading {} industry keyword configs into Kernel",
        industry_configs.len()
    );

    // Update through the Kernel's shared Arc (connected to ButlerRouterMiddleware)
    // Clone the Arc first, then drop the KernelState Mutex before awaiting RwLock.
    // This prevents blocking all other kernel commands during the write.
    let shared = {
        let kernel_guard = kernel_state.lock().await;
        kernel_guard.as_ref().map(|k| k.industry_keywords())
    };
    if let Some(shared) = shared {
        let mut guard = shared.write().await;
        *guard = industry_configs;
        tracing::info!("[viking_commands] Industry keywords synced to ButlerRouter middleware");
    } else {
        tracing::warn!("[viking_commands] Kernel not initialized, industry keywords not loaded");
    }

    Ok(())
}

// ---------------------------------------------------------------------------
// 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());
    }
}