consciousness/src/agent/tools/memory.rs

use std::sync::Arc;
// tools/memory.rs — Native memory graph operations
//
// If running in the daemon process (STORE_HANDLE set), accesses
// the store directly. Otherwise forwards to the daemon via socket.

use anyhow::{Context, Result};
use std::sync::OnceLock;

use crate::store::Store;

// ── Store handle ───────────────────────────────────────────────

/// Global store handle. Set by daemon at startup.
/// If None, tools forward to daemon socket.
static STORE_HANDLE: OnceLock<Arc<crate::Mutex<Store>>> = OnceLock::new();

// Thread-local store for rpc_local fallback path.
thread_local! {
    static LOCAL_STORE: std::cell::RefCell<Option<Arc<crate::Mutex<Store>>>> =
        const { std::cell::RefCell::new(None) };
}

/// Set the global store handle. Call once at daemon startup.
pub fn set_store(store: Arc<crate::Mutex<Store>>) {
    STORE_HANDLE.set(store).ok();
}

/// Check if we're running in daemon mode (have direct store access).
pub fn is_daemon() -> bool {
    STORE_HANDLE.get().is_some() || LOCAL_STORE.with(|s| s.borrow().is_some())
}

// ── Helpers ────────────────────────────────────────────────────

fn get_str<'a>(args: &'a serde_json::Value, name: &'a str) -> Result<&'a str> {
    args.get(name).and_then(|v| v.as_str()).context(format!("{} is required", name))
}

fn get_f64(args: &serde_json::Value, name: &str) -> Result<f64> {
    args.get(name).and_then(|v| v.as_f64()).context(format!("{} is required", name))
}

async fn cached_store() -> Result<Arc<crate::Mutex<Store>>> {
    // Check thread-local first (rpc_local fallback path)
    if let Some(store) = LOCAL_STORE.with(|s| s.borrow().clone()) {
        return Ok(store);
    }
    // Use global handle if set (daemon mode)
    if let Some(store) = STORE_HANDLE.get() {
        return Ok(store.clone());
    }
    // Fallback to loading (for backwards compat during transition)
    Store::cached().await.map_err(|e| anyhow::anyhow!("{}", e))
}

/// Run a tool with a temporarily-opened store (for rpc_local fallback).
pub async fn run_with_local_store(tool_name: &str, args: serde_json::Value) -> Result<String> {
    let store = Store::cached().await.map_err(|e| anyhow::anyhow!("{}", e))?;

    LOCAL_STORE.with(|s| *s.borrow_mut() = Some(store));
    let result = dispatch(tool_name, &None, args).await;
    LOCAL_STORE.with(|s| *s.borrow_mut() = None);

    result
}

/// Get provenance from args._provenance, or "manual".
fn get_provenance(args: &serde_json::Value) -> String {
    args.get("_provenance")
        .and_then(|v| v.as_str())
        .unwrap_or("manual")
        .to_string()
}

// ── Macro for generating tool wrappers ─────────────────────────
//
// memory_tool!(name, mut, arg1: [str], arg2: [Option<bool>])
//   - mut/ref for store mutability
//   - generates jsonargs_* (internal, JSON args) and public typed API

macro_rules! memory_tool {
    // ── Helper rules (must come first) ─────────────────────────────

    // Extract from JSON
    (@extract $args:ident, $name:ident, str) => {
        get_str($args, stringify!($name))?
    };
    (@extract $args:ident, $name:ident, f32) => {
        get_f64($args, stringify!($name))? as f32
    };
    (@extract $args:ident, $name:ident, Vec<String>) => {
        $args.get(stringify!($name))
            .and_then(|v| v.as_array())
            .map(|arr| arr.iter().filter_map(|v| v.as_str().map(String::from)).collect::<Vec<_>>())
            .unwrap_or_default()
    };
    (@extract $args:ident, $name:ident, Option<&str>) => {
        $args.get(stringify!($name)).and_then(|v| v.as_str())
    };
    (@extract $args:ident, $name:ident, Option<bool>) => {
        $args.get(stringify!($name)).and_then(|v| v.as_bool())
    };
    (@extract $args:ident, $name:ident, Option<u64>) => {
        $args.get(stringify!($name)).and_then(|v| v.as_u64())
    };
    (@extract $args:ident, $name:ident, Option<i64>) => {
        $args.get(stringify!($name)).and_then(|v| v.as_i64())
    };
    (@extract $args:ident, $name:ident, Option<usize>) => {
        $args.get(stringify!($name)).and_then(|v| v.as_u64()).map(|v| v as usize)
    };
    (@extract $args:ident, $name:ident, Option<u32>) => {
        $args.get(stringify!($name)).and_then(|v| v.as_u64()).map(|v| v as u32)
    };
    (@extract $args:ident, $name:ident, Option<f64>) => {
        $args.get(stringify!($name)).and_then(|v| v.as_f64())
    };

    // Parameter types for function signatures
    (@param_type str) => { &str };
    (@param_type f32) => { f32 };
    (@param_type Vec<String>) => { Vec<String> };
    (@param_type Option<&str>) => { Option<&str> };
    (@param_type Option<bool>) => { Option<bool> };
    (@param_type Option<u64>) => { Option<u64> };
    (@param_type Option<i64>) => { Option<i64> };
    (@param_type Option<usize>) => { Option<usize> };
    (@param_type Option<u32>) => { Option<u32> };
    (@param_type Option<f64>) => { Option<f64> };

    // Serialize to JSON for RPC
    (@insert_json $map:ident, $name:ident, str) => {
        $map.insert(stringify!($name).into(), serde_json::json!($name));
    };
    (@insert_json $map:ident, $name:ident, f32) => {
        $map.insert(stringify!($name).into(), serde_json::json!($name));
    };
    (@insert_json $map:ident, $name:ident, Vec<String>) => {
        $map.insert(stringify!($name).into(), serde_json::json!($name));
    };
    (@insert_json $map:ident, $name:ident, Option<&str>) => {
        if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
    };
    (@insert_json $map:ident, $name:ident, Option<bool>) => {
        if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
    };
    (@insert_json $map:ident, $name:ident, Option<u64>) => {
        if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
    };
    (@insert_json $map:ident, $name:ident, Option<i64>) => {
        if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
    };
    (@insert_json $map:ident, $name:ident, Option<usize>) => {
        if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
    };
    (@insert_json $map:ident, $name:ident, Option<u32>) => {
        if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
    };
    (@insert_json $map:ident, $name:ident, Option<f64>) => {
        if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
    };

    // ── Main rules ─────────────────────────────────────────────────

    // Mutable store variant
    ($name:ident, mut $(, $($arg:ident : [$($typ:tt)+]),* $(,)?)?) => {
        paste::paste! {
            async fn [<jsonargs_ $name>](args: &serde_json::Value) -> Result<String> {
                $($(let $arg = memory_tool!(@extract args, $arg, $($typ)+);)*)?
                let prov = get_provenance(args);
                let arc = cached_store().await?;
                let mut store = arc.lock().await;
                crate::hippocampus::$name(&mut store, &prov $($(, $arg)*)?)
            }

            pub async fn $name(agent: Option<&crate::agent::Agent> $($(, $arg: memory_tool!(@param_type $($typ)+))*)?) -> Result<String> {
                if !is_daemon() {
                    #[allow(unused_mut)]
                    let mut map = serde_json::Map::new();
                    $($(memory_tool!(@insert_json map, $arg, $($typ)+);)*)?
                    return crate::mcp_server::memory_rpc(concat!("memory_", stringify!($name)), serde_json::Value::Object(map));
                }
                let prov = match agent {
                    Some(a) => a.state.lock().await.provenance.clone(),
                    None => "manual".to_string(),
                };
                let arc = cached_store().await?;
                let mut store = arc.lock().await;
                crate::hippocampus::$name(&mut store, &prov $($(, $arg)*)?)
            }
        }
    };

    // Immutable store variant
    ($name:ident, ref $(, $($arg:ident : [$($typ:tt)+]),* $(,)?)?) => {
        paste::paste! {
            async fn [<jsonargs_ $name>](args: &serde_json::Value) -> Result<String> {
                $($(let $arg = memory_tool!(@extract args, $arg, $($typ)+);)*)?
                let prov = get_provenance(args);
                let arc = cached_store().await?;
                let store = arc.lock().await;
                crate::hippocampus::$name(&store, &prov $($(, $arg)*)?)
            }

            pub async fn $name(agent: Option<&crate::agent::Agent> $($(, $arg: memory_tool!(@param_type $($typ)+))*)?) -> Result<String> {
                if !is_daemon() {
                    #[allow(unused_mut)]
                    let mut map = serde_json::Map::new();
                    $($(memory_tool!(@insert_json map, $arg, $($typ)+);)*)?
                    return crate::mcp_server::memory_rpc(concat!("memory_", stringify!($name)), serde_json::Value::Object(map));
                }
                let prov = match agent {
                    Some(a) => a.state.lock().await.provenance.clone(),
                    None => "manual".to_string(),
                };
                let arc = cached_store().await?;
                let store = arc.lock().await;
                crate::hippocampus::$name(&store, &prov $($(, $arg)*)?)
            }
        }
    };
}

// ── Memory tools ───────────────────────────────────────────────

memory_tool!(render, ref, key: [str], raw: [Option<bool>]);
memory_tool!(write, mut, key: [str], content: [str]);
memory_tool!(search, ref, keys: [Vec<String>], max_hops: [Option<u32>], edge_decay: [Option<f64>], min_activation: [Option<f64>], limit: [Option<usize>]);
memory_tool!(links, ref, key: [str]);
memory_tool!(link_set, mut, source: [str], target: [str], strength: [f32]);
memory_tool!(link_add, mut, source: [str], target: [str]);
memory_tool!(delete, mut, key: [str]);
memory_tool!(history, ref, key: [str], full: [Option<bool>]);
memory_tool!(weight_set, mut, key: [str], weight: [f32]);
memory_tool!(rename, mut, old_key: [str], new_key: [str]);
memory_tool!(supersede, mut, old_key: [str], new_key: [str], reason: [Option<&str>]);
memory_tool!(query, ref, query: [str], format: [Option<&str>]);

// ── Journal tools ──────────────────────────────────────────────

memory_tool!(journal_tail, ref, count: [Option<u64>], level: [Option<u64>], format: [Option<&str>], after: [Option<&str>]);
memory_tool!(journal_new, mut, name: [str], title: [str], body: [str], level: [Option<i64>]);
memory_tool!(journal_update, mut, body: [str], level: [Option<i64>]);

// ── Graph tools ───────────────────────────────────────────────

memory_tool!(graph_topology, ref);
memory_tool!(graph_health, ref);
memory_tool!(graph_communities, ref, top_n: [Option<usize>], min_size: [Option<usize>]);
memory_tool!(graph_normalize_strengths, mut, apply: [Option<bool>]);
memory_tool!(graph_link_impact, ref, source: [str], target: [str]);
memory_tool!(graph_hubs, ref, count: [Option<usize>]);
memory_tool!(graph_trace, ref, key: [str]);

/// Single entry point for all memory/journal tool calls.
/// If not daemon, forwards to daemon with provenance attached.
async fn dispatch(
    tool_name: &str,
    agent: &Option<std::sync::Arc<crate::agent::Agent>>,
    args: serde_json::Value,
) -> Result<String> {
    let mut args = args;
    if let Some(a) = agent {
        let prov = a.state.lock().await.provenance.clone();
        args.as_object_mut().map(|o| o.insert("_provenance".into(), prov.into()));
    }

    if !is_daemon() {
        // Forward to daemon
        let name = tool_name.to_string();
        return tokio::task::spawn_blocking(move || {
            crate::mcp_server::memory_rpc(&name, args)
        }).await.map_err(|e| anyhow::anyhow!("spawn_blocking: {}", e))?;
    }

    // Daemon path - dispatch to implementation
    match tool_name {
        "memory_render" => jsonargs_render(&args).await,
        "memory_write" => jsonargs_write(&args).await,
        "memory_search" => jsonargs_search(&args).await,
        "memory_links" => jsonargs_links(&args).await,
        "memory_link_set" => jsonargs_link_set(&args).await,
        "memory_link_add" => jsonargs_link_add(&args).await,
        "memory_delete" => jsonargs_delete(&args).await,
        "memory_history" => jsonargs_history(&args).await,
        "memory_weight_set" => jsonargs_weight_set(&args).await,
        "memory_rename" => jsonargs_rename(&args).await,
        "memory_supersede" => jsonargs_supersede(&args).await,
        "memory_query" => jsonargs_query(&args).await,
        "graph_topology" => jsonargs_graph_topology(&args).await,
        "graph_health" => jsonargs_graph_health(&args).await,
        "graph_communities" => jsonargs_graph_communities(&args).await,
        "graph_normalize_strengths" => jsonargs_graph_normalize_strengths(&args).await,
        "graph_trace" => jsonargs_graph_trace(&args).await,
        "graph_link_impact" => jsonargs_graph_link_impact(&args).await,
        "graph_hubs" => jsonargs_graph_hubs(&args).await,
        "journal_tail" => jsonargs_journal_tail(&args).await,
        "journal_new" => jsonargs_journal_new(&args).await,
        "journal_update" => jsonargs_journal_update(&args).await,
        _ => anyhow::bail!("unknown tool: {}", tool_name),
    }
}

// ── Definitions ────────────────────────────────────────────────

pub fn memory_tools() -> [super::Tool; 15] {
    use super::Tool;
    [
        Tool { name: "memory_render", description: "Read a memory node's content and links.",
               parameters_json: r#"{"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_render", &a, v).await })) },
        Tool { name: "memory_write", description: "Create or update a memory node.",
               parameters_json: r#"{"type":"object","properties":{"key":{"type":"string","description":"Node key"},"content":{"type":"string","description":"Full content (markdown)"}},"required":["key","content"]}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_write", &a, v).await })) },
        Tool { name: "memory_search", description: "Search the memory graph via spreading activation. Give 2-4 seed node keys.",
               parameters_json: r#"{"type":"object","properties":{"keys":{"type":"array","items":{"type":"string"},"description":"Seed node keys to activate from"},"max_hops":{"type":"integer","description":"Max graph hops (default 3)"},"edge_decay":{"type":"number","description":"Decay per hop (default 0.3)"},"min_activation":{"type":"number","description":"Cutoff threshold (default 0.01)"},"limit":{"type":"integer","description":"Max results (default 20)"}},"required":["keys"]}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_search", &a, v).await })) },
        Tool { name: "memory_links", description: "Show a node's neighbors with link strengths.",
               parameters_json: r#"{"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_links", &a, v).await })) },
        Tool { name: "memory_link_set", description: "Set link strength between two nodes.",
               parameters_json: r#"{"type":"object","properties":{"source":{"type":"string"},"target":{"type":"string"},"strength":{"type":"number","description":"0.01 to 1.0"}},"required":["source","target","strength"]}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_link_set", &a, v).await })) },
        Tool { name: "memory_link_add", description: "Add a new link between two nodes.",
               parameters_json: r#"{"type":"object","properties":{"source":{"type":"string"},"target":{"type":"string"}},"required":["source","target"]}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_link_add", &a, v).await })) },
        Tool { name: "memory_delete", description: "Delete a memory node.",
               parameters_json: r#"{"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_delete", &a, v).await })) },
        Tool { name: "memory_history", description: "Show version history for a node.",
               parameters_json: r#"{"type":"object","properties":{"key":{"type":"string","description":"Node key"},"full":{"type":"boolean","description":"Show full content for each version"}},"required":["key"]}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_history", &a, v).await })) },
        Tool { name: "memory_weight_set", description: "Set a node's weight directly (0.01 to 1.0).",
               parameters_json: r#"{"type":"object","properties":{"key":{"type":"string"},"weight":{"type":"number","description":"0.01 to 1.0"}},"required":["key","weight"]}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_weight_set", &a, v).await })) },
        Tool { name: "memory_rename", description: "Rename a node key in place.",
               parameters_json: r#"{"type":"object","properties":{"old_key":{"type":"string"},"new_key":{"type":"string"}},"required":["old_key","new_key"]}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_rename", &a, v).await })) },
        Tool { name: "memory_supersede", description: "Mark a node as superseded by another (sets weight to 0.01).",
               parameters_json: r#"{"type":"object","properties":{"old_key":{"type":"string"},"new_key":{"type":"string"},"reason":{"type":"string"}},"required":["old_key","new_key"]}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_supersede", &a, v).await })) },
        Tool { name: "memory_query",
               description: "Run a structured query against the memory graph.",
               parameters_json: r#"{
                   "type": "object",
                   "properties": {
                       "query":  {"type": "string", "description": "Query expression"},
                       "format": {"type": "string", "description": "compact (default) or full (with content and graph metrics)", "default": "compact"}
                   },
                   "required": ["query"]
               }"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("memory_query", &a, v).await })) },
        Tool { name: "graph_topology", description: "Show graph topology stats (nodes, edges, clustering, hubs).",
               parameters_json: r#"{"type":"object","properties":{}}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("graph_topology", &a, v).await })) },
        Tool { name: "graph_health", description: "Show graph health report with maintenance recommendations.",
               parameters_json: r#"{"type":"object","properties":{}}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("graph_health", &a, v).await })) },
        Tool { name: "graph_hubs", description: "Show top hub nodes by degree, spread apart for diverse link targets.",
               parameters_json: r#"{"type":"object","properties":{"count":{"type":"integer","description":"Number of hubs to return (default 20)"}}}"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("graph_hubs", &a, v).await })) },
    ]
}

pub fn journal_tools() -> [super::Tool; 3] {
    use super::Tool;
    [
        Tool { name: "journal_tail",
               description: "Read the last N entries at a given level.",
               parameters_json: r#"{
                   "type": "object",
                   "properties": {
                       "count":  {"type": "integer", "description": "Number of entries", "default": 1},
                       "level":  {"type": "integer", "description": "0=journal, 1=daily, 2=weekly, 3=monthly", "default": 0},
                       "format": {"type": "string", "description": "compact or full (with content)", "default": "full"},
                       "after":  {"type": "string", "description": "Only entries after this date (YYYY-MM-DD)"}
                   }
               }"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("journal_tail", &a, v).await })) },
        Tool { name: "journal_new", description: "Start a new journal/digest entry.",
               parameters_json: r#"{
                   "type": "object",
                   "properties": {
                       "name":  {"type": "string", "description": "Short node name (becomes the key)"},
                       "title": {"type": "string", "description": "Descriptive title"},
                       "body":  {"type": "string", "description": "Entry body"},
                       "level": {"type": "integer", "description": "0=journal, 1=daily, 2=weekly, 3=monthly", "default": 0}
                   },
                   "required": ["name", "title", "body"]
               }"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("journal_new", &a, v).await })) },
        Tool { name: "journal_update", description: "Append text to the most recent entry at a level.",
               parameters_json: r#"{
                   "type": "object",
                   "properties": {
                       "body":  {"type": "string", "description": "Text to append"},
                       "level": {"type": "integer", "description": "0=journal, 1=daily, 2=weekly, 3=monthly", "default": 0}
                   },
                   "required": ["body"]
               }"#,
               handler: Arc::new(|a, v| Box::pin(async move { dispatch("journal_update", &a, v).await })) },
    ]
}