ruvector-sona 0.1.8

//! Simplified NAPI-RS bindings for Node.js
//! Enable with feature flag: `napi`
//!
//! This version uses a simpler API that doesn't expose TrajectoryBuilder to JS

#![cfg(feature = "napi")]

use napi_derive::napi;
use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};

use crate::{
    LearnedPattern, SonaConfig, SonaEngine as RustSonaEngine,
    TrajectoryBuilder as RustTrajectoryBuilder,
};

// Global storage for trajectory builders
fn get_trajectory_builders() -> &'static Mutex<HashMap<u32, RustTrajectoryBuilder>> {
    static BUILDERS: OnceLock<Mutex<HashMap<u32, RustTrajectoryBuilder>>> = OnceLock::new();
    BUILDERS.get_or_init(|| Mutex::new(HashMap::new()))
}

fn get_next_builder_id() -> &'static Mutex<u32> {
    static NEXT_ID: OnceLock<Mutex<u32>> = OnceLock::new();
    NEXT_ID.get_or_init(|| Mutex::new(0))
}

/// Node.js SONA Engine wrapper
#[napi]
pub struct SonaEngine {
    inner: RustSonaEngine,
}

#[napi]
impl SonaEngine {
    /// Create a new SONA engine with default configuration
    /// @param hidden_dim - Hidden dimension size (e.g., 256, 512)
    #[napi(constructor)]
    pub fn new(hidden_dim: u32) -> Self {
        Self {
            inner: RustSonaEngine::new(hidden_dim as usize),
        }
    }

    /// Create with custom configuration
    /// @param config - Custom SONA configuration object
    #[napi(factory)]
    pub fn with_config(config: JsSonaConfig) -> Self {
        let rust_config = SonaConfig {
            hidden_dim: config.hidden_dim as usize,
            embedding_dim: config.embedding_dim.unwrap_or(config.hidden_dim) as usize,
            micro_lora_rank: config.micro_lora_rank.unwrap_or(1) as usize,
            base_lora_rank: config.base_lora_rank.unwrap_or(8) as usize,
            micro_lora_lr: config.micro_lora_lr.unwrap_or(0.001) as f32,
            base_lora_lr: config.base_lora_lr.unwrap_or(0.0001) as f32,
            ewc_lambda: config.ewc_lambda.unwrap_or(1000.0) as f32,
            pattern_clusters: config.pattern_clusters.unwrap_or(50) as usize,
            trajectory_capacity: config.trajectory_capacity.unwrap_or(10000) as usize,
            background_interval_ms: config.background_interval_ms.unwrap_or(3600000) as u64,
            quality_threshold: config.quality_threshold.unwrap_or(0.5) as f32,
            enable_simd: config.enable_simd.unwrap_or(true),
        };
        Self {
            inner: RustSonaEngine::with_config(rust_config),
        }
    }

    /// Start a new trajectory recording
    /// @param query_embedding - Query embedding vector (Float64Array)
    /// @returns Trajectory ID for adding steps
    #[napi]
    pub fn begin_trajectory(&self, query_embedding: Vec<f64>) -> u32 {
        let embedding: Vec<f32> = query_embedding.iter().map(|&x| x as f32).collect();
        let builder = self.inner.begin_trajectory(embedding);

        let mut builders = get_trajectory_builders().lock().unwrap();
        let mut next_id = get_next_builder_id().lock().unwrap();
        let id = *next_id;
        *next_id += 1;
        builders.insert(id, builder);
        id
    }

    /// Add a step to trajectory
    /// @param trajectory_id - Trajectory ID from beginTrajectory
    /// @param activations - Layer activations (Float64Array)
    /// @param attention_weights - Attention weights (Float64Array)
    /// @param reward - Reward signal for this step
    #[napi]
    pub fn add_trajectory_step(
        &self,
        trajectory_id: u32,
        activations: Vec<f64>,
        attention_weights: Vec<f64>,
        reward: f64,
    ) {
        let mut builders = get_trajectory_builders().lock().unwrap();
        if let Some(builder) = builders.get_mut(&trajectory_id) {
            let act: Vec<f32> = activations.iter().map(|&x| x as f32).collect();
            let att: Vec<f32> = attention_weights.iter().map(|&x| x as f32).collect();
            builder.add_step(act, att, reward as f32);
        }
    }

    /// Set model route for trajectory
    /// @param trajectory_id - Trajectory ID
    /// @param route - Model route identifier
    #[napi]
    pub fn set_trajectory_route(&self, trajectory_id: u32, route: String) {
        let mut builders = get_trajectory_builders().lock().unwrap();
        if let Some(builder) = builders.get_mut(&trajectory_id) {
            builder.set_model_route(&route);
        }
    }

    /// Add context to trajectory
    /// @param trajectory_id - Trajectory ID
    /// @param context_id - Context identifier
    #[napi]
    pub fn add_trajectory_context(&self, trajectory_id: u32, context_id: String) {
        let mut builders = get_trajectory_builders().lock().unwrap();
        if let Some(builder) = builders.get_mut(&trajectory_id) {
            builder.add_context(&context_id);
        }
    }

    /// Complete a trajectory and submit for learning
    /// @param trajectory_id - Trajectory ID
    /// @param quality - Final quality score [0.0, 1.0]
    #[napi]
    pub fn end_trajectory(&self, trajectory_id: u32, quality: f64) {
        let mut builders = get_trajectory_builders().lock().unwrap();
        if let Some(builder) = builders.remove(&trajectory_id) {
            let trajectory = builder.build(quality as f32);
            self.inner.submit_trajectory(trajectory);
        }
    }

    /// Apply micro-LoRA transformation to input
    /// @param input - Input vector (Float64Array)
    /// @returns Transformed output vector
    #[napi]
    pub fn apply_micro_lora(&self, input: Vec<f64>) -> Vec<f64> {
        let input_f32: Vec<f32> = input.iter().map(|&x| x as f32).collect();
        let mut output = vec![0.0f32; input_f32.len()];
        self.inner.apply_micro_lora(&input_f32, &mut output);
        output.iter().map(|&x| x as f64).collect()
    }

    /// Apply base-LoRA transformation to layer output
    /// @param layer_idx - Layer index
    /// @param input - Input vector (Float64Array)
    /// @returns Transformed output vector
    #[napi]
    pub fn apply_base_lora(&self, layer_idx: u32, input: Vec<f64>) -> Vec<f64> {
        let input_f32: Vec<f32> = input.iter().map(|&x| x as f32).collect();
        let mut output = vec![0.0f32; input_f32.len()];
        self.inner
            .apply_base_lora(layer_idx as usize, &input_f32, &mut output);
        output.iter().map(|&x| x as f64).collect()
    }

    /// Run background learning cycle if due
    /// @returns Optional status message if cycle was executed
    #[napi]
    pub fn tick(&self) -> Option<String> {
        self.inner.tick()
    }

    /// Force background learning cycle immediately
    /// @returns Status message with learning results
    #[napi]
    pub fn force_learn(&self) -> String {
        self.inner.force_learn()
    }

    /// Flush instant loop updates
    #[napi]
    pub fn flush(&self) {
        self.inner.flush();
    }

    /// Find similar learned patterns to query
    /// @param query_embedding - Query embedding vector
    /// @param k - Number of patterns to return
    /// @returns Array of learned patterns
    #[napi]
    pub fn find_patterns(&self, query_embedding: Vec<f64>, k: u32) -> Vec<JsLearnedPattern> {
        let query: Vec<f32> = query_embedding.iter().map(|&x| x as f32).collect();
        self.inner
            .find_patterns(&query, k as usize)
            .into_iter()
            .map(JsLearnedPattern::from)
            .collect()
    }

    /// Get engine statistics as JSON string
    /// @returns Statistics object as JSON string
    #[napi]
    pub fn get_stats(&self) -> String {
        serde_json::to_string(&self.inner.stats())
            .unwrap_or_else(|e| format!("{{\"error\": \"{}\"}}", e))
    }

    /// Enable or disable the engine
    /// @param enabled - Whether to enable the engine
    #[napi]
    pub fn set_enabled(&mut self, enabled: bool) {
        self.inner.set_enabled(enabled);
    }

    /// Check if engine is enabled
    /// @returns Whether the engine is enabled
    #[napi]
    pub fn is_enabled(&self) -> bool {
        self.inner.is_enabled()
    }
}

/// SONA configuration for Node.js
#[napi(object)]
pub struct JsSonaConfig {
    /// Hidden dimension size
    pub hidden_dim: u32,
    /// Embedding dimension (defaults to hidden_dim)
    pub embedding_dim: Option<u32>,
    /// Micro-LoRA rank (1-2, default: 1)
    pub micro_lora_rank: Option<u32>,
    /// Base LoRA rank (default: 8)
    pub base_lora_rank: Option<u32>,
    /// Micro-LoRA learning rate (default: 0.001)
    pub micro_lora_lr: Option<f64>,
    /// Base LoRA learning rate (default: 0.0001)
    pub base_lora_lr: Option<f64>,
    /// EWC lambda regularization (default: 1000.0)
    pub ewc_lambda: Option<f64>,
    /// Number of pattern clusters (default: 50)
    pub pattern_clusters: Option<u32>,
    /// Trajectory buffer capacity (default: 10000)
    pub trajectory_capacity: Option<u32>,
    /// Background learning interval in ms (default: 3600000 = 1 hour)
    pub background_interval_ms: Option<i64>,
    /// Quality threshold for learning (default: 0.5)
    pub quality_threshold: Option<f64>,
    /// Enable SIMD optimizations (default: true)
    pub enable_simd: Option<bool>,
}

/// Learned pattern for Node.js
#[napi(object)]
pub struct JsLearnedPattern {
    /// Pattern identifier
    pub id: String,
    /// Cluster centroid embedding
    pub centroid: Vec<f64>,
    /// Number of trajectories in cluster
    pub cluster_size: u32,
    /// Total weight of trajectories
    pub total_weight: f64,
    /// Average quality of member trajectories
    pub avg_quality: f64,
    /// Creation timestamp (Unix seconds)
    pub created_at: String,
    /// Last access timestamp (Unix seconds)
    pub last_accessed: String,
    /// Total access count
    pub access_count: u32,
    /// Pattern type
    pub pattern_type: String,
}

impl From<LearnedPattern> for JsLearnedPattern {
    fn from(pattern: LearnedPattern) -> Self {
        Self {
            id: pattern.id.to_string(),
            centroid: pattern.centroid.iter().map(|&x| x as f64).collect(),
            cluster_size: pattern.cluster_size as u32,
            total_weight: pattern.total_weight as f64,
            avg_quality: pattern.avg_quality as f64,
            created_at: pattern.created_at.to_string(),
            last_accessed: pattern.last_accessed.to_string(),
            access_count: pattern.access_count,
            pattern_type: format!("{:?}", pattern.pattern_type),
        }
    }
}