hubro-sdk 0.9.7

//! Mobile-specific functionality for the Hubro SDK.
//!
//! This module provides interfaces for interacting with mobile health data (Health Connect on Android, HealthKit on iOS),
//! platform detection, machine learning model execution on the host, and TFHE key management.
//! It relies on WASI and custom host imports defined in the `hubro_sdk` module.

use crate::records::{HealthRecord, Sampleable};
use rand::SeedableRng;
use rand::prelude::*;
use rand::rngs::SmallRng;
use serde::de;
use serde::{Deserialize, Serialize};
use std::ffi::CStr;
use std::mem;
use std::os::raw::{c_char, c_void};
use tfhe::ServerKey;
use wasip1;

/// Represents the platform where the plugin is currently running.
#[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
pub enum Platform {
    /// Apple iOS.
    IOS,
    /// Google Android.
    Android,
    /// JavaScript/Browser environment.
    JS,
    /// Unknown or unsupported platform.
    Unknown,
}

/// Host functions imported from the `hubro_sdk` WASM module.
#[link(wasm_import_module = "hubro_sdk")]
unsafe extern "C" {
    /// Retrieves health records from the host in JSON format.
    fn get_health_connect_records(record_type: i32, from: i32, to: i32) -> *mut c_char;
    /// Retrieves the number of health records available on the host for the given range.
    fn get_health_connect_number_of_records(record_type: i32, from: i32, to: i32) -> i32;
    /// Prints a line to the host's debug output.
    fn print_line(nf_name: *mut c_char);
    /// Gets the current platform name from the host.
    fn get_platform() -> *mut c_char;
    /// Uploads a TFHE server key to the host.
    fn upload_server_key(ptr: *const u8, len: usize) -> i32;
    /// Loads a machine learning model from the specified path on the host.
    fn host_load_model(path_ptr: *const u8, path_len: usize) -> i32;
    /// Runs inference on a loaded machine learning model.
    fn host_run_inference(
        model_id: i32,
        input_ptr: *const f32,
        input_len: i32,
        shape_ptr: *const i32,
        shape_len: i32,
        output_ptr: *mut f32,
        output_len: i32,
    ) -> i32;
}

/// Allocates memory on the WASM heap and returns a pointer to it.
///
/// This is typically called by the host to pass data into the guest.
/// The caller is responsible for deallocating the memory using [`dealloc`].
#[unsafe(no_mangle)]
pub extern "C" fn allocate(size: usize) -> *mut c_void {
    let mut buf = Vec::with_capacity(size);
    let ptr = buf.as_mut_ptr();
    mem::forget(buf);
    return ptr as *mut c_void;
}

/// Deallocates memory previously allocated by [`allocate`].
///
/// # Safety
///
/// This function is unsafe because it reconstructs a `Vec` from a raw pointer and capacity
/// to drop it. The `ptr` must have been allocated by `allocate`, and `cap` must match
/// the size used during allocation.
#[unsafe(no_mangle)]
pub extern "C" fn dealloc(ptr: *mut c_void, cap: usize) {
    unsafe {
        let _buf = Vec::from_raw_parts(ptr, 0, cap);
    }
}

/// Fetches the count of records from the host, retrying if the host is busy.
fn fetch_records_count(record_type: i32, from: i32, to: i32) -> i32 {
    let mut s;
    loop {
        s = unsafe { get_health_connect_number_of_records(record_type, from, to) };
        if s != -1 {
            break;
        }
    }
    s
}

/// Fetches health records of type `T` from the host and deserializes them from JSON.
fn fetch_records<T: de::DeserializeOwned + HealthRecord>(
    record_type: i32,
    from: i32,
    to: i32,
) -> Vec<T> {
    if fetch_records_count(record_type, from, to) > 0 {
        let s = unsafe { get_health_connect_records(record_type, from, to) };
        let json_data = unsafe { CStr::from_ptr(s).to_bytes().to_vec() };
        let subject_str = std::str::from_utf8(&json_data).unwrap();
        return serde_json::from_str::<Vec<T>>(subject_str).unwrap();
    } else {
        Vec::<T>::new()
    }
}

/// Retrieves health records of type `T` from the host for a given time range.
///
/// # Arguments
///
/// * `from` - Start timestamp (Unix epoch in seconds).
/// * `to` - End timestamp (Unix epoch in seconds).
pub fn get_health_records<T: de::DeserializeOwned + HealthRecord>(from: i32, to: i32) -> Vec<T> {
    fetch_records::<T>(T::IDENTIFIER, from, to)
}

/// Generates sample health records for testing or fallback purposes.
///
/// Uses WASI clock and `getrandom` to seed a `SmallRng` for data generation.
pub fn generate_sample_records<T: Serialize + Sampleable>(from: i32, to: i32) -> Vec<T> {
    let mut seed = [0u8; 32];

    // Integrate wasip1::clock_time_get into seeding
    if let Ok(time) = unsafe { wasip1::clock_time_get(wasip1::CLOCKID_REALTIME, 1_000) } {
        let b = time.to_le_bytes();
        seed[0..8].copy_from_slice(&b);
    }

    let _ = getrandom::getrandom(&mut seed[8..32]);

    let mut rng = SmallRng::from_seed(seed);

    let mut records = Vec::new();
    let start_time = chrono::DateTime::from_timestamp(from as i64, 0)
        .unwrap()
        .naive_utc()
        .and_utc();
    let end_time = chrono::DateTime::from_timestamp(to as i64, 0)
        .unwrap()
        .naive_utc()
        .and_utc();
    let duration = end_time.signed_duration_since(start_time);
    let days = duration.num_days();

    for day in 0..days {
        let base_timestamp = start_time + chrono::Duration::days(day);
        let num_records = rng.random_range(3..8);

        for hour in 0..num_records {
            let timestamp = base_timestamp + chrono::Duration::hours(hour * 4);
            let record = T::generate_sample(&mut rng, timestamp, chrono::Duration::hours(4));
            records.push(record);
        }
    }
    records
}

/// Prints a debug string to the host's output console.
pub fn debug_print_line(output: &str) {
    let size = output.len();
    let ptr = allocate(size + 1) as *mut c_char;
    unsafe {
        std::ptr::copy(output.as_ptr(), ptr as *mut u8, size);
        *ptr.add(size) = 0;
        print_line(ptr);
    }
}

/// Detects the current platform by querying the host environment.
pub fn get_current_platform() -> Platform {
    let ptr = unsafe { get_platform() };
    if ptr.is_null() {
        return Platform::Unknown;
    }
    let s = unsafe { CStr::from_ptr(ptr).to_string_lossy().to_string() };
    let s_lower = s.to_lowercase();
    if s_lower.contains("ios") {
        Platform::IOS
    } else if s_lower.contains("android") {
        Platform::Android
    } else if s_lower.contains("js") {
        Platform::JS
    } else {
        Platform::Unknown
    }
}

/// Loads a machine learning model from the specified path on the host.
///
/// Returns a model ID handle on success.
pub fn load_model(path: &str) -> i32 {
    unsafe { host_load_model(path.as_ptr(), path.len()) }
}

/// Runs inference on a loaded machine learning model.
///
/// # Arguments
///
/// * `model_id` - The model handle returned by [`load_model`].
/// * `input` - Flattened array of input features.
/// * `shape` - Dimensions of the input tensor.
/// * `output` - Buffer where the inference results will be written.
pub fn run_inference(model_id: i32, input: &[f32], shape: &[i32], output: &mut [f32]) -> i32 {
    unsafe {
        host_run_inference(
            model_id,
            input.as_ptr(),
            input.len() as i32,
            shape.as_ptr(),
            shape.len() as i32,
            output.as_mut_ptr(),
            output.len() as i32,
        )
    }
}

/// Uploads a TFHE server key to the host for use in encrypted computations.
pub fn upload_key(server_key: &ServerKey) -> i32 {
    let key = serialize_keys(server_key);
    unsafe { upload_server_key(key.as_ptr(), key.len()) }
}

/// Returns the current time in nanoseconds since the Unix epoch using the WASI realtime clock.
pub fn get_current_time_nanos() -> u64 {
    unsafe { wasip1::clock_time_get(wasip1::CLOCKID_REALTIME, 1_000).expect("WASI clock_time_get failed") }
}

/// Serializes a TFHE `ServerKey` into a byte vector using `bincode`.
pub fn serialize_keys(server_key: &ServerKey) -> Vec<u8> {
    let config = bincode::config::standard();
    bincode::serde::encode_to_vec(server_key, config)
        .expect("Failed to serialize TFHE server key using Serde bridge")
}