lib.rs

pub mod engine;
pub mod executor;
pub mod units;

#[cfg(target_arch = "wasm32")]
getrandom::register_custom_getrandom!(sdk::js_interop::getrandom_custom);

#[cfg(test)]
pub mod benchmarks;

use engine::ComputeEngine;
use log::info;
use sdk::{Epoch, Reactor, IDX_SYSTEM_EPOCH};
use units::{
    AudioUnit, BoidUnit, CryptoUnit, DataUnit, GpuUnit, ImageUnit, MathUnit, PhysicsEngine,
    VideoUnit,
};

// --- PERSISTENT SAB CACHE ---
use std::sync::OnceLock;

// Use OnceLock for thread-safe one-time initialization without spin-waiting
static GLOBAL_SAB: OnceLock<sdk::sab::SafeSAB> = OnceLock::new();

const MAX_FFI_ROUTE_PART_BYTES: usize = 64;
const MAX_FFI_PARAMS_BYTES: usize = 1024 * 1024;
const MAX_FFI_RESULT_BYTES: usize = 64 * 1024 * 1024;

pub(crate) fn set_cached_sab(sab: sdk::sab::SafeSAB) {
    let _ = GLOBAL_SAB.set(sab);
}

pub(crate) fn get_cached_sab() -> Option<sdk::sab::SafeSAB> {
    GLOBAL_SAB.get().cloned()
}

// Use OnceLock for engine to avoid lock overhead on every access
static COMPUTE_ENGINE: OnceLock<ComputeEngine> = OnceLock::new();

fn get_engine() -> &'static ComputeEngine {
    COMPUTE_ENGINE.get_or_init(|| initialize_engine())
}

fn initialize_engine() -> ComputeEngine {
    use std::sync::Arc;
    let mut engine = ComputeEngine::new();

    // Register Unit Proxies (Arc for thread-safety)
    engine.register(Arc::new(ImageUnit::new()));
    engine.register(Arc::new(CryptoUnit::new()));
    engine.register(Arc::new(DataUnit::new()));
    engine.register(Arc::new(AudioUnit::new()));
    engine.register(Arc::new(GpuUnit::new()));
    engine.register(Arc::new(PhysicsEngine::new()));
    engine.register(Arc::new(BoidUnit::new()));
    engine.register(Arc::new(VideoUnit::new()));
    // NOTE: ApiProxy is NOT registered here - it's handled separately
    // due to browser API constraints (HTTP/WebSocket use non-Send types)
    engine.register(Arc::new(MathUnit::new()));
    // NOTE: StorageUnit is NOT registered here - it's handled separately via
    // StorageSupervisor due to browser API constraints (IndexedDB/OPFS are non-Send)

    engine
}

/// Standardized Memory Allocator for WebAssembly
#[no_mangle]
pub extern "C" fn compute_alloc(size: usize) -> *mut u8 {
    let mut buf = Vec::with_capacity(size);
    let ptr = buf.as_mut_ptr();
    std::mem::forget(buf);
    ptr
}

/// Standardized Memory Deallocator for WebAssembly
#[no_mangle]
pub extern "C" fn compute_free(ptr: *mut u8, size: usize) {
    if !ptr.is_null() {
        unsafe {
            let _ = Vec::from_raw_parts(ptr, 0, size);
        }
    }
}

/// Standardized Initialization with SharedArrayBuffer
#[no_mangle]
pub extern "C" fn compute_init_with_sab() -> i32 {
    let global = sdk::js_interop::get_global();
    let sab_key = sdk::js_interop::create_string("__INOS_SAB__");
    let sab_val = sdk::js_interop::reflect_get(&global, &sab_key);
    let offset_key = sdk::js_interop::create_string("__INOS_SAB_OFFSET__");
    let offset_val = sdk::js_interop::reflect_get(&global, &offset_key);
    let size_key = sdk::js_interop::create_string("__INOS_SAB_SIZE__");
    let size_val = sdk::js_interop::reflect_get(&global, &size_key);
    let id_key = sdk::js_interop::create_string("__INOS_MODULE_ID__");
    let id_val = sdk::js_interop::reflect_get(&global, &id_key);

    if let (Ok(val), Ok(off), Ok(sz), Ok(id)) = (sab_val, offset_val, size_val, id_val) {
        if !val.is_undefined() && !val.is_null() {
            let offset = sdk::js_interop::as_f64(&off).unwrap_or(0.0) as u32;
            let size = sdk::js_interop::as_f64(&sz).unwrap_or(0.0) as u32;
            let module_id = sdk::js_interop::as_f64(&id).unwrap_or(0.0) as u32;

            let global_sab = sdk::sab::SafeSAB::new(&val);
            sdk::set_module_id(module_id);
            sdk::identity::init_identity_from_js();
            sdk::init_logging();
            sdk::sab::set_global_barrier_view(global_sab.barrier_view().clone());
            sdk::init_context();

            info!("Compute module initialized (ID: {}) with synchronized SAB bridge (Offset: 0x{:x}, Size: {}MB)", 
                module_id, offset, size / 1024 / 1024);

            set_cached_sab(global_sab.clone());
            register_compute_capabilities(&global_sab);

            return 1;
        }
    }
    0
}

// --- GENERIC UNIT DISPATCHER ---
// This allows JS to call ANY registered unit method via a single entry point

/// Generic compute execution dispatcher
/// Allows JavaScript to call any registered unit via:
///   compute_execute("math", "matrix_multiply", input_ptr, input_len, params_ptr, params_len)
///
/// Returns: pointer to result buffer (first 4 bytes = length, rest = data), or 0 on error
///
/// Example usage from JS:
///   const result = compute_execute("math", "matrix_identity", 0, 0, paramsPtr, paramsLen);
#[no_mangle]
pub extern "C" fn compute_execute(
    service_ptr: *const u8,
    service_len: usize,
    action_ptr: *const u8,
    action_len: usize,
    input_ptr: *const u8,
    input_len: usize,
    params_ptr: *const u8,
    params_len: usize,
) -> *mut u8 {
    // 0. Context Validation
    if !sdk::is_context_valid() {
        sdk::js_interop::console_log(
            "[compute_execute] FAILED: Context is invalid (Zombie Module)",
            2,
        );
        return std::ptr::null_mut();
    }

    // 1. Marshall service name
    let service = unsafe {
        if service_ptr.is_null() || service_len == 0 {
            sdk::js_interop::console_log(
                "[compute_execute] FAILED: service_ptr is null or len=0",
                1,
            );
            return std::ptr::null_mut();
        }
        if service_len > MAX_FFI_ROUTE_PART_BYTES {
            sdk::js_interop::console_log("[compute_execute] FAILED: service name too long", 1);
            return std::ptr::null_mut();
        }
        let bytes = std::slice::from_raw_parts(service_ptr, service_len);
        match std::str::from_utf8(bytes) {
            Ok(s) => s,
            Err(_) => {
                let bytes_hex: Vec<String> =
                    bytes.iter().take(8).map(|b| format!("{:02x}", b)).collect();
                let msg = format!("[compute_execute] FAILED: service name not valid UTF-8. Ptr: {:p}, Len: {}, Bytes(8): {:?}", 
                    service_ptr, service_len, bytes_hex);
                sdk::js_interop::console_log(&msg, 1);
                return std::ptr::null_mut();
            }
        }
    };

    // 2. Marshall action name
    let action = unsafe {
        if action_ptr.is_null() || action_len == 0 {
            sdk::js_interop::console_log(
                "[compute_execute] FAILED: action_ptr is null or len=0",
                1,
            );
            return std::ptr::null_mut();
        }
        if action_len > MAX_FFI_ROUTE_PART_BYTES {
            sdk::js_interop::console_log("[compute_execute] FAILED: action name too long", 1);
            return std::ptr::null_mut();
        }
        match std::str::from_utf8(std::slice::from_raw_parts(action_ptr, action_len)) {
            Ok(s) => s,
            Err(_) => {
                sdk::js_interop::console_log(
                    "[compute_execute] FAILED: action name is not valid UTF-8",
                    1,
                );
                return std::ptr::null_mut();
            }
        }
    };

    let input = unsafe {
        if input_ptr.is_null() || input_len == 0 {
            &[]
        } else {
            std::slice::from_raw_parts(input_ptr, input_len)
        }
    };

    let params = unsafe {
        if params_ptr.is_null() || params_len == 0 {
            b"{}"
        } else {
            if params_len > MAX_FFI_PARAMS_BYTES {
                sdk::js_interop::console_log("[compute_execute] FAILED: params too large", 1);
                return std::ptr::null_mut();
            }
            std::slice::from_raw_parts(params_ptr, params_len)
        }
    };

    // Run the async execute in a synchronous context by polling it once
    // We CANNOT use block_on because it uses Atomics.wait which is forbidden on main thread

    // Initialize engine if needed (thread-safe spinlock)
    let engine = get_engine();
    let result = match poll_sync(engine.execute(service, action, input, params)) {
        Ok(res) => res,
        Err(e) => {
            let msg = format!("[compute_execute] Sync Execution Failed: {}", e);
            sdk::js_interop::console_log(&msg, 1);
            return std::ptr::null_mut();
        }
    };

    // Debug logging removed - was running 120+ times/sec at 60 FPS

    match result {
        Ok(output) => wrap_result_buffer(output),
        Err(e) => {
            // Log error and return null
            let msg = format!("[compute_execute] Logic Error: {}", e);
            sdk::js_interop::console_log(&msg, 1);
            std::ptr::null_mut()
        }
    }
}

/// Zero-Copy Protocol Dispatcher
/// Standardizes communication with Cap'n Proto JobRequest capsule
/// Standard usage: compute_dispatch(request_ptr, request_len)
#[no_mangle]
pub extern "C" fn compute_dispatch(request_ptr: *const u8, request_len: usize) -> *mut u8 {
    if !sdk::is_context_valid() || request_ptr.is_null() || request_len == 0 {
        return std::ptr::null_mut();
    }
    if request_len > MAX_FFI_PARAMS_BYTES {
        return std::ptr::null_mut();
    }

    let request_bytes = unsafe { std::slice::from_raw_parts(request_ptr, request_len) };
    let mut reader = std::io::Cursor::new(request_bytes);

    // Read message from slice (zero-copy from WASM heap perspective)
    let message_reader =
        match capnp::serialize::read_message(&mut reader, capnp::message::ReaderOptions::new()) {
            Ok(r) => r,
            Err(_) => return std::ptr::null_mut(),
        };

    let job =
        match message_reader.get_root::<sdk::protocols::compute::compute::job_request::Reader>() {
            Ok(j) => j,
            Err(_) => return std::ptr::null_mut(),
        };

    // Extract fields using zero-copy lenses
    let service = job
        .get_library() // Field remains 'library' in schema for compatibility
        .unwrap_or(capnp::text::Reader::from(""))
        .to_str()
        .unwrap_or("");
    let action = job
        .get_method() // Field remains 'method' in schema for compatibility
        .unwrap_or(capnp::text::Reader::from(""))
        .to_str()
        .unwrap_or("");
    let input = job.get_input().unwrap_or(&[]);

    let params_reader = job.get_params().unwrap();
    let params = match params_reader.which().unwrap() {
        sdk::protocols::compute::compute::job_params::Which::Binary(data) => data.unwrap_or(&[]),
        _ => &[], // Structured params handled inside specialized units if needed
    };

    let engine = get_engine();
    let result = match poll_sync(engine.execute(service, action, input, params)) {
        Ok(res) => res,
        Err(_) => return std::ptr::null_mut(),
    };

    match result {
        Ok(output) => wrap_result_buffer(output),
        Err(_) => std::ptr::null_mut(),
    }
}

fn wrap_result_buffer(output: Vec<u8>) -> *mut u8 {
    if output.len() > MAX_FFI_RESULT_BYTES {
        sdk::js_interop::console_log(
            "[compute_execute] FAILED: result too large for FFI frame",
            1,
        );
        return std::ptr::null_mut();
    }

    let output_len = output.len();
    let total_len = match 4usize.checked_add(output_len) {
        Some(len) => len,
        None => return std::ptr::null_mut(),
    };
    let mut buffer = Vec::with_capacity(total_len);
    buffer.extend_from_slice(&(output_len as u32).to_le_bytes());
    buffer.extend_from_slice(&output);

    let ptr = buffer.as_mut_ptr();
    std::mem::forget(buffer);
    ptr
}

/// Helper to poll a future once synchronously
/// Panics or errors if the future yields (is not ready immediately)
fn poll_sync<T>(future: impl std::future::Future<Output = T>) -> Result<T, String> {
    use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};

    unsafe fn clone(_: *const ()) -> RawWaker {
        RawWaker::new(std::ptr::null(), &VTABLE)
    }
    unsafe fn wake(_: *const ()) {}
    unsafe fn wake_by_ref(_: *const ()) {}
    unsafe fn drop(_: *const ()) {}

    static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop);

    // Safety: we are single-threaded on WASM main thread usually, or this is just a dummy waker
    let raw_waker = RawWaker::new(std::ptr::null(), &VTABLE);
    let waker = unsafe { Waker::from_raw(raw_waker) };
    let mut cx = Context::from_waker(&waker);

    let mut pinned = std::pin::pin!(future);
    match pinned.as_mut().poll(&mut cx) {
        Poll::Ready(val) => Ok(val),
        Poll::Pending => {
            Err("Future yielded! compute_execute requires synchronous completion.".to_string())
        }
    }
}
// ======================================================================
// LEGACY EXPORTS REMOVED
// ======================================================================
// All compute operations now use compute_execute() for consistency.
// The previous direct exports (compute_boids_step, compute_nbody_step, etc.)
// are now routed through the ComputeEngine via their respective units.
// ======================================================================

pub struct ComputeKernel {
    reactor: Reactor,
    engine: ComputeEngine,
    epoch: Epoch,
}

// Helper to register capabilities (moved from ComputeKernel::new to be standalone)
fn register_compute_capabilities(sab: &sdk::sab::SafeSAB) {
    use sdk::registry::*;

    // Register core modules provided by this kernel
    // Dynamically register all units from the engine
    let engine = get_engine();
    let capabilities = engine.generate_capability_registry();

    let mut builder = ModuleEntryBuilder::new("compute").version(1, 6, 0);
    // Add all discovered capabilities to the 'compute' module entry
    // These are registered as "service:action" for strictly generic discovery.
    for cap_str in &capabilities {
        // cap_str is "service:action:v1"
        let parts: Vec<&str> = cap_str.split(':').collect();
        if parts.len() >= 2 {
            let cap_name = format!("{}:{}", parts[0], parts[1]);
            builder = builder.capability(&cap_name, false, 512);
        }
    }

    match builder.build() {
        Ok((mut entry, _, caps)) => {
            if let Ok(offset) = write_capability_table(sab, &caps) {
                entry.cap_table_offset = offset;
            }
            if let Ok((slot, _)) = find_slot_double_hashing(sab, "compute") {
                let _ = write_enhanced_entry(sab, slot, &entry);
                info!(
                    "[Compute] Registered 0x{:x} generic capabilities to 'compute' module",
                    capabilities.len()
                );
            }
        }
        Err(e) => info!("Failed to auto-register compute: {:?}", e),
    }

    // Signal registry change to wake any listeners (SystemStore epoch watcher)
    sdk::registry::signal_registry_change(sab);
}

impl ComputeKernel {
    pub fn new(sab: sdk::sab::SafeSAB, node_id: String) -> Self {
        sdk::init_logging();
        info!("Compute Kernel initialized on node {}", node_id);

        let engine = initialize_engine();
        let reactor = Reactor::new(sab.clone());

        // Use standardized System Epoch index from SDK
        let epoch = Epoch::new(sab, IDX_SYSTEM_EPOCH);

        // No need to call register_compute_capabilities here anymore,
        // it's already done in compute_init_with_sab using the correct safe_sab.

        Self {
            reactor,
            engine,
            epoch,
        }
    }

    /// Poll for new compute segments using Reactive Mutation
    pub async fn poll(&mut self) -> bool {
        if !self.reactor.check_inbox() {
            return false;
        }

        self.reactor.ack_inbox();

        // 1. Get Inbox data and copy to buffer
        let data = match self.reactor.read_request() {
            Some(d) => d,
            None => return false,
        };

        // 2. Execute via Engine
        // Use proper Cap'n Proto processing
        let result = self.process_job(&data).await;

        match result {
            Ok(output) => {
                // Return success result
                if let Ok(serialized) = self.serialize_result(true, &output, "") {
                    if !self.reactor.write_result(&serialized) {
                        log::error!("Output too large for outbox: {} bytes", serialized.len());
                        // Write error result
                        if let Ok(err_bytes) = self.serialize_result(false, &[], "Output too large")
                        {
                            self.reactor.write_result(&err_bytes);
                        }
                    }
                }
            }
            Err(e) => {
                log::error!("Compute job failed: {}", e);
                // Write error result
                if let Ok(err_bytes) = self.serialize_result(false, &[], &e.to_string()) {
                    self.reactor.write_result(&err_bytes);
                }
            }
        }

        // 3. Signal completion via Epoch
        self.epoch.increment();

        true
    }

    /// Process job using Cap'n Proto "Lens"
    async fn process_job(&self, data: &[u8]) -> Result<Vec<u8>, engine::ComputeError> {
        let mut reader = std::io::Cursor::new(data);
        let message_reader =
            capnp::serialize::read_message(&mut reader, capnp::message::ReaderOptions::new())
                .map_err(|e| {
                    engine::ComputeError::ExecutionFailed(format!("Capnp read error: {}", e))
                })?;

        // Access the lens
        let job = message_reader
            .get_root::<sdk::protocols::compute::compute::job_request::Reader>()
            .map_err(|e| {
                engine::ComputeError::ExecutionFailed(format!("Capnp root error: {}", e))
            })?;

        // Zero-copy field access
        let library_reader = job
            .get_library()
            .map_err(|_| engine::ComputeError::ExecutionFailed("Invalid library field".into()))?;
        let library = library_reader
            .to_str()
            .map_err(|_| engine::ComputeError::ExecutionFailed("Library not valid UTF-8".into()))?;

        let method_reader = job
            .get_method()
            .map_err(|_| engine::ComputeError::ExecutionFailed("Invalid method field".into()))?;
        let method = method_reader
            .to_str()
            .map_err(|_| engine::ComputeError::ExecutionFailed("Method not valid UTF-8".into()))?;

        let params_reader = job
            .get_params()
            .map_err(|_| engine::ComputeError::ExecutionFailed("Invalid params field".into()))?;

        let params_bytes = match params_reader.which() {
            Ok(sdk::protocols::compute::compute::job_params::Which::Binary(data)) => {
                data.map_err(|_| {
                    engine::ComputeError::ExecutionFailed("Invalid binary params".into())
                })?
            }
            Ok(sdk::protocols::compute::compute::job_params::Which::CustomParams(custom_res)) => {
                let custom = custom_res.map_err(|_| {
                    engine::ComputeError::ExecutionFailed("Invalid custom params".into())
                })?;
                custom
                    .get_shader_source()
                    .map_err(|_| {
                        engine::ComputeError::ExecutionFailed("Invalid shader field".into())
                    })?
                    .as_bytes()
            }
            _ => &[], // Other structured types fall back to empty bytes for generic engines
        };

        let params = params_bytes;

        let input = job
            .get_input()
            .map_err(|_| engine::ComputeError::ExecutionFailed("Invalid input field".into()))?;

        info!(
            "Engine execution (Capnp): unit={}, action={}, input_size={}",
            library,
            method,
            input.len()
        );

        self.engine.execute(library, method, input, params).await
    }

    /// Helper to serialize JobResult
    fn serialize_result(
        &self,
        success: bool,
        data: &[u8],
        error_msg: &str,
    ) -> Result<Vec<u8>, engine::ComputeError> {
        let mut message = capnp::message::Builder::new_default();
        let mut root = message.init_root::<sdk::protocols::compute::compute::job_result::Builder>();

        // Set status
        if success {
            root.set_status(sdk::protocols::compute::compute::Status::Success);
        } else {
            root.set_status(sdk::protocols::compute::compute::Status::Failed);
        }

        // Set output
        root.set_output(data);

        // Set error message
        root.set_error_message(error_msg);

        let mut output_bytes = Vec::new();
        capnp::serialize::write_message(&mut output_bytes, &message).map_err(|e| {
            engine::ComputeError::ExecutionFailed(format!("Serialize error: {}", e))
        })?;

        Ok(output_bytes)
    }
}