engram-retired/crates/engram-jni/src/lib.rs

/// JNI bindings for engram-core.
///
/// These functions expose the Engram API to Kotlin/JVM callers via Java Native Interface.
/// The convention is:
///
///   Java_<package>_<class>_<method>
///   → Java_ai_neuron_engram_EngramDb_<method>
///
/// The `EngramDb` handle is stored as a Java `long` (native pointer). The Kotlin
/// wrapper class casts it to/from `Long` and keeps it private.
///
/// # Memory model
/// - `open` allocates an `EngramHandle` on the Rust heap and returns its address as `jlong`.
/// - `close` takes that `jlong`, reconstructs the Box, and drops it.
/// - All other methods borrow the handle via `&*ptr`.
///
/// # JSON wire format
/// Nodes and results are passed as JSON strings to avoid bespoke JNI object marshalling.
/// The Kotlin layer converts between the data classes and JSON.
use engram_core::{ActivatedNode, EngramDb, MemoryTier, Node, NodeType, ScoredNode};
use jni::objects::{JClass, JString};
use jni::sys::{jfloatArray, jint, jlong, jstring};
use jni::JNIEnv;
use std::path::Path;
use uuid::Uuid;

// ── Handle ────────────────────────────────────────────────────────────────────

struct EngramHandle {
    db: EngramDb,
}

// ── Helper macros ─────────────────────────────────────────────────────────────

macro_rules! handle_ref {
    ($handle:expr) => {
        unsafe { &*($handle as *const EngramHandle) }
    };
}

// ── JNI methods: lifecycle ────────────────────────────────────────────────────

/// Open an engram database and return a native handle as jlong.
///
/// Kotlin: `external fun open(path: String): Long`
#[no_mangle]
pub extern "system" fn Java_ai_neuron_engram_EngramDb_open(
    mut env: JNIEnv,
    _class: JClass,
    path: JString,
) -> jlong {
    let path_str: String = match env.get_string(&path) {
        Ok(s) => s.into(),
        Err(_) => return 0,
    };
    match EngramDb::open(Path::new(&path_str)) {
        Ok(db) => {
            let handle = Box::new(EngramHandle { db });
            Box::into_raw(handle) as jlong
        }
        Err(e) => {
            let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
            0
        }
    }
}

/// Close and free a database handle.
///
/// Kotlin: `external fun close(handle: Long)`
#[no_mangle]
pub extern "system" fn Java_ai_neuron_engram_EngramDb_close(
    _env: JNIEnv,
    _class: JClass,
    handle: jlong,
) {
    if handle != 0 {
        unsafe {
            drop(Box::from_raw(handle as *mut EngramHandle));
        }
    }
}

// ── JNI methods: statistics ───────────────────────────────────────────────────

/// Kotlin: `external fun nodeCount(handle: Long): Long`
#[no_mangle]
pub extern "system" fn Java_ai_neuron_engram_EngramDb_nodeCount(
    mut env: JNIEnv,
    _class: JClass,
    handle: jlong,
) -> jlong {
    let h = handle_ref!(handle);
    match h.db.node_count() {
        Ok(n) => n as jlong,
        Err(e) => {
            let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
            -1
        }
    }
}

/// Kotlin: `external fun edgeCount(handle: Long): Long`
#[no_mangle]
pub extern "system" fn Java_ai_neuron_engram_EngramDb_edgeCount(
    mut env: JNIEnv,
    _class: JClass,
    handle: jlong,
) -> jlong {
    let h = handle_ref!(handle);
    match h.db.edge_count() {
        Ok(n) => n as jlong,
        Err(e) => {
            let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
            -1
        }
    }
}

// ── JNI methods: nodes ────────────────────────────────────────────────────────

/// Store a node from JSON and return the assigned UUID string.
///
/// Kotlin: `external fun putNode(handle: Long, nodeJson: String): String`
#[no_mangle]
pub extern "system" fn Java_ai_neuron_engram_EngramDb_putNode(
    mut env: JNIEnv,
    _class: JClass,
    handle: jlong,
    node_json: JString,
) -> jstring {
    let json: String = match env.get_string(&node_json) {
        Ok(s) => s.into(),
        Err(_) => return std::ptr::null_mut(),
    };

    let node = match node_from_json(&json) {
        Some(n) => n,
        None => {
            let _ = env.throw_new("java/lang/IllegalArgumentException", "Invalid node JSON");
            return std::ptr::null_mut();
        }
    };

    let h = handle_ref!(handle);
    match h.db.put_node(node) {
        Ok(id) => {
            let id_str = id.to_string();
            env.new_string(&id_str)
                .map(|s| s.into_raw())
                .unwrap_or(std::ptr::null_mut())
        }
        Err(e) => {
            let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
            std::ptr::null_mut()
        }
    }
}

/// Retrieve a node by UUID, returned as JSON, or null if not found.
///
/// Kotlin: `external fun getNode(handle: Long, id: String): String?`
#[no_mangle]
pub extern "system" fn Java_ai_neuron_engram_EngramDb_getNode(
    mut env: JNIEnv,
    _class: JClass,
    handle: jlong,
    id: JString,
) -> jstring {
    let id_str: String = match env.get_string(&id) {
        Ok(s) => s.into(),
        Err(_) => return std::ptr::null_mut(),
    };
    let uuid = match id_str.parse::<Uuid>() {
        Ok(u) => u,
        Err(_) => return std::ptr::null_mut(),
    };

    let h = handle_ref!(handle);
    match h.db.get_node(uuid) {
        Ok(Some(node)) => {
            let json = node_to_json(&node);
            env.new_string(&json)
                .map(|s| s.into_raw())
                .unwrap_or(std::ptr::null_mut())
        }
        Ok(None) => std::ptr::null_mut(),
        Err(e) => {
            let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
            std::ptr::null_mut()
        }
    }
}

// ── JNI methods: search ───────────────────────────────────────────────────────

/// Search for similar nodes by embedding vector.
/// Returns a JSON array of scored nodes.
///
/// Kotlin: `external fun searchEmbedding(handle: Long, embedding: FloatArray, limit: Int): String`
#[no_mangle]
pub extern "system" fn Java_ai_neuron_engram_EngramDb_searchEmbedding(
    mut env: JNIEnv,
    _class: JClass,
    handle: jlong,
    embedding: jfloatArray,
    limit: jint,
) -> jstring {
    let emb = match float_array_from_jni(&mut env, embedding) {
        Some(v) => v,
        None => return std::ptr::null_mut(),
    };

    let h = handle_ref!(handle);
    match h.db.search_embedding(&emb, limit as usize) {
        Ok(results) => {
            let json = scored_nodes_to_json(&results);
            env.new_string(&json)
                .map(|s| s.into_raw())
                .unwrap_or(std::ptr::null_mut())
        }
        Err(e) => {
            let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
            std::ptr::null_mut()
        }
    }
}

/// Run spreading activation.
/// `seeds_json` is a JSON array of UUID strings.
/// Returns a JSON array of activated nodes.
///
/// Kotlin: `external fun activate(handle: Long, seedsJson: String, queryEmbedding: FloatArray, maxDepth: Int, limit: Int): String`
#[no_mangle]
pub extern "system" fn Java_ai_neuron_engram_EngramDb_activate(
    mut env: JNIEnv,
    _class: JClass,
    handle: jlong,
    seeds_json: JString,
    query_embedding: jfloatArray,
    max_depth: jint,
    limit: jint,
) -> jstring {
    let seeds_str: String = match env.get_string(&seeds_json) {
        Ok(s) => s.into(),
        Err(_) => return std::ptr::null_mut(),
    };

    let seeds: Vec<Uuid> = parse_uuid_array(&seeds_str);
    let query_emb = match float_array_from_jni(&mut env, query_embedding) {
        Some(v) => v,
        None => return std::ptr::null_mut(),
    };

    let h = handle_ref!(handle);
    match h.db.activate(&seeds, &query_emb, max_depth as u8, limit as usize) {
        Ok(results) => {
            let json = activated_nodes_to_json(&results);
            env.new_string(&json)
                .map(|s| s.into_raw())
                .unwrap_or(std::ptr::null_mut())
        }
        Err(e) => {
            let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
            std::ptr::null_mut()
        }
    }
}

// ── JNI methods: salience ─────────────────────────────────────────────────────

/// Touch a node (increment activation count and update salience).
///
/// Kotlin: `external fun touch(handle: Long, id: String)`
#[no_mangle]
pub extern "system" fn Java_ai_neuron_engram_EngramDb_touch(
    mut env: JNIEnv,
    _class: JClass,
    handle: jlong,
    id: JString,
) {
    let id_str: String = match env.get_string(&id) {
        Ok(s) => s.into(),
        Err(_) => return,
    };
    let uuid = match id_str.parse::<Uuid>() {
        Ok(u) => u,
        Err(_) => return,
    };
    let h = handle_ref!(handle);
    if let Err(e) = h.db.touch(uuid) {
        let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
    }
}

/// Apply salience decay. Returns the number of nodes updated.
///
/// Kotlin: `external fun decay(handle: Long, factor: Float): Int`
#[no_mangle]
pub extern "system" fn Java_ai_neuron_engram_EngramDb_decay(
    mut env: JNIEnv,
    _class: JClass,
    handle: jlong,
    factor: f32,
) -> jint {
    let h = handle_ref!(handle);
    match h.db.decay(factor) {
        Ok(n) => n as jint,
        Err(e) => {
            let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
            -1
        }
    }
}

// ── JNI helpers ───────────────────────────────────────────────────────────────

fn float_array_from_jni(env: &mut JNIEnv, arr: jfloatArray) -> Option<Vec<f32>> {
    if arr.is_null() {
        return None;
    }
    let arr_obj = unsafe { jni::objects::JFloatArray::from_raw(arr) };
    let len = env.get_array_length(&arr_obj).ok()? as usize;
    let mut buf = vec![0f32; len];
    env.get_float_array_region(&arr_obj, 0, &mut buf).ok()?;
    Some(buf)
}

fn parse_uuid_array(json: &str) -> Vec<Uuid> {
    // Minimal parser: `["uuid1","uuid2",...]`
    json.trim_matches(|c| c == '[' || c == ']')
        .split(',')
        .filter_map(|s| {
            let s = s.trim().trim_matches('"');
            s.parse::<Uuid>().ok()
        })
        .collect()
}

// ── JSON helpers ──────────────────────────────────────────────────────────────

fn node_from_json(json: &str) -> Option<Node> {
    let content = extract_string_field(json, "content").unwrap_or_default();
    let node_type_str =
        extract_string_field(json, "node_type").unwrap_or_else(|| "Memory".into());
    let tier_str = extract_string_field(json, "tier").unwrap_or_else(|| "Episodic".into());
    let importance: f32 = extract_f32_field(json, "importance").unwrap_or(0.5);
    let embedding = extract_f32_array(json, "embedding").unwrap_or_default();

    let node_type = match node_type_str.as_str() {
        "Concept" => NodeType::Concept,
        "Event" => NodeType::Event,
        "Entity" => NodeType::Entity,
        "Process" => NodeType::Process,
        "InternalState" => NodeType::InternalState,
        _ => NodeType::Memory,
    };
    let tier = match tier_str.as_str() {
        "Working" => MemoryTier::Working,
        "Semantic" => MemoryTier::Semantic,
        "Procedural" => MemoryTier::Procedural,
        _ => MemoryTier::Episodic,
    };
    Some(Node::new(node_type, embedding, content.into_bytes(), tier, importance))
}

fn node_to_json(node: &Node) -> String {
    let content = String::from_utf8_lossy(&node.content)
        .replace('\\', "\\\\")
        .replace('"', "\\\"");
    let emb_str = node
        .embedding
        .iter()
        .map(|f| format!("{:.6}", f))
        .collect::<Vec<_>>()
        .join(",");
    format!(
        r#"{{"id":"{}","node_type":"{:?}","tier":"{:?}","content":"{}","salience":{:.6},"importance":{:.6},"activation_count":{},"embedding":[{}]}}"#,
        node.id, node.node_type, node.tier, content, node.salience, node.importance,
        node.activation_count, emb_str,
    )
}

fn scored_nodes_to_json(nodes: &[ScoredNode]) -> String {
    let items: Vec<String> = nodes
        .iter()
        .map(|s| format!(r#"{{"node":{},"score":{:.6}}}"#, node_to_json(&s.node), s.score))
        .collect();
    format!("[{}]", items.join(","))
}

fn activated_nodes_to_json(nodes: &[ActivatedNode]) -> String {
    let items: Vec<String> = nodes
        .iter()
        .map(|a| {
            format!(
                r#"{{"node":{},"activation_strength":{:.6},"hops":{}}}"#,
                node_to_json(&a.node), a.activation_strength, a.hops,
            )
        })
        .collect();
    format!("[{}]", items.join(","))
}

fn extract_string_field(json: &str, key: &str) -> Option<String> {
    let needle = format!("\"{}\":", key);
    let start = json.find(&needle)? + needle.len();
    let rest = json[start..].trim_start();
    if !rest.starts_with('"') {
        return None;
    }
    let inner = &rest[1..];
    let end = inner.find('"')?;
    Some(inner[..end].to_string())
}

fn extract_f32_field(json: &str, key: &str) -> Option<f32> {
    let needle = format!("\"{}\":", key);
    let start = json.find(&needle)? + needle.len();
    let rest = json[start..].trim_start();
    let end = rest
        .find(|c: char| c == ',' || c == '}')
        .unwrap_or(rest.len());
    rest[..end].trim().parse::<f32>().ok()
}

fn extract_f32_array(json: &str, key: &str) -> Option<Vec<f32>> {
    let needle = format!("\"{}\":", key);
    let start = json.find(&needle)? + needle.len();
    let rest = json[start..].trim_start();
    if !rest.starts_with('[') {
        return None;
    }
    let end = rest.find(']')?;
    let inner = &rest[1..end];
    Some(
        inner
            .split(',')
            .filter_map(|s| s.trim().parse::<f32>().ok())
            .collect(),
    )
}

// ── Tests ─────────────────────────────────────────────────────────────────────

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

    #[test]
    fn node_json_roundtrip() {
        let node = Node::new(
            NodeType::Memory,
            vec![0.1, 0.2, 0.3],
            b"test content".to_vec(),
            MemoryTier::Episodic,
            0.8,
        );
        let json = node_to_json(&node);
        assert!(json.contains("test content"));
        assert!(json.contains("Memory"));
        assert!(json.contains("Episodic"));
    }

    #[test]
    fn parse_uuid_array_valid() {
        let uuids = parse_uuid_array(r#"["550e8400-e29b-41d4-a716-446655440000"]"#);
        assert_eq!(uuids.len(), 1);
    }

    #[test]
    fn parse_uuid_array_empty() {
        let uuids = parse_uuid_array("[]");
        assert_eq!(uuids.len(), 0);
    }

    #[test]
    fn extract_string_field_works() {
        let json = r#"{"content":"hello","type":"Memory"}"#;
        assert_eq!(extract_string_field(json, "content"), Some("hello".into()));
        assert_eq!(extract_string_field(json, "type"), Some("Memory".into()));
    }

    #[test]
    fn extract_f32_array_works() {
        let json = r#"{"embedding":[0.1,0.2,0.3]}"#;
        let arr = extract_f32_array(json, "embedding").unwrap();
        assert_eq!(arr.len(), 3);
    }

    #[test]
    fn activated_nodes_json_is_array() {
        let json = activated_nodes_to_json(&[]);
        assert_eq!(json, "[]");
    }
}