viam_bridge_sdk/

loader.rs

use anyhow::{anyhow, Result};
use libloading::{Library, Symbol};
use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use std::sync::Arc;
use serde_json::Value;

use crate::bridge::{
    BridgeTypeRegistry, ExtensionBridge, ExtensionBridgeWithRegistry, TypedRegistryCallback,
};

use super::bridge::BridgeContext;

/// Type to store bridge in trait object for type erasure
pub type BoxedBridge = Box<dyn ExtensionBridgeWithRegistry>;

/// Function type for creating a bridge
pub type CreateBridgeFn = fn() -> BoxedBridge;

/// Registry for managing bridge callbacks
pub struct BridgeRegistry {
    /// Map of TypeId to callbacks for that type
    callbacks: HashMap<TypeId, Vec<TypedRegistryCallback>>,
}

impl BridgeRegistry {
    /// Create a new empty bridge registry
    pub fn new() -> Self {
        Self {
            callbacks: HashMap::new(),
        }
    }

    /// Register a bridge with the registry
    pub fn register_bridge(&mut self, bridge: &mut dyn ExtensionBridgeWithRegistry) -> Result<()> {
        // The bridge itself will handle adding its callbacks to the registry
        // We're using a trait object to get dynamic dispatch
        let registry: &mut dyn BridgeTypeRegistry = bridge;

        // Get all the callbacks from the bridge
        let callbacks = registry.get_callbacks_for(TypeId::of::<Self>());

        // Add each callback to our registry
        for callback in callbacks {
            let host_type_id = callback.host_type_id;

            // Create entry if it doesn't exist
            if !self.callbacks.contains_key(&host_type_id) {
                self.callbacks.insert(host_type_id, Vec::new());
            }

            // Access the entry and add the callback
            if let Some(callbacks_for_type) = self.callbacks.get_mut(&host_type_id) {
                // Clone is not supported for TypedRegistryCallback, so we'd need to implement
                // a proper Clone or copy mechanism from the callback
                // For now, this would require bridge implementations to provide a way to get
                // multiple callbacks
                // callbacks_for_type.push(callback.clone());
                // This is a TODO: properly get callbacks from bridges
            }
        }

        Ok(())
    }

    /// Get callbacks for a specific host type
    pub fn get_callbacks<T: 'static>(&self) -> Vec<&TypedRegistryCallback> {
        let host_type_id = TypeId::of::<T>();
        match self.callbacks.get(&host_type_id) {
            Some(callbacks) => callbacks.iter().collect(),
            None => Vec::new(),
        }
    }
}

/// Trait for loading bridges
pub trait BridgeLoader {
    /// Load bridges from a source
    fn load_bridges(&mut self) -> Result<Vec<BoxedBridge>>;
}

/// Dynamic bridge loader for loading bridges from shared libraries
pub struct DynamicBridgeLoader {
    /// Directory to scan for bridge libraries
    pub directory: String,
    /// Registry for managing bridge callbacks
    registry: BridgeRegistry,
    context: Option<Arc<BridgeContext>>,
}

impl DynamicBridgeLoader {
    /// Create a new dynamic bridge loader
    pub fn new(directory: &str) -> Self {
        Self {
            directory: directory.to_string(),
            registry: BridgeRegistry::new(),
            context: None,
        }
    }

    /// Get a reference to the registry
    pub fn registry(&self) -> &BridgeRegistry {
        &self.registry
    }

    /// Get a mutable reference to the registry
    pub fn registry_mut(&mut self) -> &mut BridgeRegistry {
        &mut self.registry
    }

    pub fn with_context(mut self, context: Arc<BridgeContext>) -> Self {
        self.context = Some(context);
        self
    }
}

type CreateBridgeWithContextFn = fn(Option<Arc<BridgeContext>>) -> BoxedBridge;

impl BridgeLoader for DynamicBridgeLoader {
    fn load_bridges(&mut self) -> Result<Vec<BoxedBridge>> {
        println!("[BRIDGE_LOADER] Starting load_bridges");
        let mut bridges = Vec::new();

        // Ensure directory exists
        let directory = Path::new(&self.directory);
        if !directory.exists() {
            println!("[BRIDGE_LOADER] Bridge directory does not exist: {}", self.directory);
            return Err(anyhow!(
                "Bridge directory does not exist: {}",
                self.directory
            ));
        }

        // Collect libraries and bridges
        let mut libraries_to_forget = Vec::new();

        // Scan directory for bridge libraries
        println!("[BRIDGE_LOADER] Scanning directory for bridge libraries: {}", self.directory);
        for entry in fs::read_dir(directory)? {
            let entry = entry?;
            let entry_path = entry.path();
            println!("[BRIDGE_LOADER] Found entry: {:?}", entry_path);

            // determine library path and optional JSON config
            let (path, config_value): (std::path::PathBuf, Value) = if entry_path.is_dir() {
                let config_file = entry_path.join("config.json");
                let config_value = if config_file.exists() && config_file.is_file() {
                    let s = fs::read_to_string(&config_file)?;
                    serde_json::from_str(&s)?
                } else {
                    Value::Null
                };
                // find library inside
                let mut libs = Vec::new();
                for sub in fs::read_dir(&entry_path)? {
                    let sub = sub?;
                    let p = sub.path();
                    #[cfg(unix)]
                    if p.extension().and_then(|e| e.to_str()) == Some("so") {
                        libs.push(p.clone());
                    }
                    #[cfg(windows)]
                    if p.extension().and_then(|e| e.to_str()) == Some("dll") {
                        libs.push(p.clone());
                    }
                    #[cfg(target_os = "macos")]
                    if p.extension().and_then(|e| e.to_str()) == Some("dylib") {
                        libs.push(p.clone());
                    }
                }
                if libs.is_empty() {
                    println!("[BRIDGE_LOADER] No library found in directory: {:?}", entry_path);
                    continue;
                }
                (libs.remove(0), config_value)
            } else {
                #[cfg(unix)]
                if !entry_path.is_file() || entry_path.extension().and_then(|e| e.to_str()) != Some("so") {
                    println!("[BRIDGE_LOADER] Skipping non-so file: {:?}", entry_path);
                    continue;
                }
                #[cfg(windows)]
                if !entry_path.is_file() || entry_path.extension().and_then(|e| e.to_str()) != Some("dll") {
                    println!("[BRIDGE_LOADER] Skipping non-dll file: {:?}", entry_path);
                    continue;
                }
                #[cfg(target_os = "macos")]
                if !entry_path.is_file() || entry_path.extension().and_then(|e| e.to_str()) != Some("dylib") {
                    println!("[BRIDGE_LOADER] Skipping non-dylib file: {:?}", entry_path);
                    continue;
                }
                (entry_path, Value::Null)
            };

            // Load the library
            #[cfg(unix)]
            let library = unsafe {
                println!("[BRIDGE_LOADER] Attempting to load library: {:?}", path);
                let lib_result = libloading::os::unix::Library::open(
                    Some(path.to_str().unwrap()),
                    libloading::os::unix::RTLD_NOW | libloading::os::unix::RTLD_GLOBAL,
                );
                
                match lib_result {
                    Ok(lib) => {
                        println!("[BRIDGE_LOADER] Successfully loaded library: {:?}", path);
                        libloading::Library::from(lib)
                    },
                    Err(e) => {
                        println!("[BRIDGE_LOADER] Failed to load library: {:?}, error: {:?}", path, e);
                        return Err(anyhow!("Failed to load library: {:?}, error: {:?}", path, e));
                    }
                }
            };

            #[cfg(not(unix))]
            let library = unsafe { 
                println!("[BRIDGE_LOADER] Attempting to load library: {:?}", path);
                let lib_result = Library::new(&path);
                
                match lib_result {
                    Ok(lib) => {
                        println!("[BRIDGE_LOADER] Successfully loaded library: {:?}", path);
                        lib
                    },
                    Err(e) => {
                        println!("[BRIDGE_LOADER] Failed to load library: {:?}, error: {:?}", path, e);
                        return Err(anyhow!("Failed to load library: {:?}, error: {:?}", path, e));
                    }
                }
            };

            // Try to get the new create_bridge symbol that accepts context
            println!("[BRIDGE_LOADER] Looking for create_bridge_with_context symbol in: {:?}", path);
            let create_bridge_result: Result<Symbol<CreateBridgeWithContextFn>, libloading::Error> = 
                unsafe { library.get(b"create_bridge_with_context") };
            
            let mut bridge = if let Ok(create_bridge) = create_bridge_result {
                // Call the new function with context
                println!("[BRIDGE_LOADER] Found create_bridge_with_context symbol, calling it with context: {:?}", self.context.is_some());
                let bridge_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                    create_bridge(self.context.clone())
                }));
                
                match bridge_result {
                    Ok(bridge) => {
                        println!("[BRIDGE_LOADER] Successfully created bridge with context");
                        bridge
                    },
                    Err(e) => {
                        println!("[BRIDGE_LOADER] Panic in create_bridge_with_context: {:?}", e);
                        return Err(anyhow!("Panic in create_bridge_with_context"));
                    }
                }
            } else {
                // Fall back to the original create_bridge function
                println!("[BRIDGE_LOADER] create_bridge_with_context not found, looking for create_bridge");
                let create_bridge_result: Result<Symbol<CreateBridgeFn>, libloading::Error> = 
                    unsafe { library.get(b"create_bridge") };
                
                match create_bridge_result {
                    Ok(create_bridge) => {
                        println!("[BRIDGE_LOADER] Found create_bridge symbol, calling it");
                        let bridge_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                            create_bridge()
                        }));
                        
                        match bridge_result {
                            Ok(bridge) => {
                                println!("[BRIDGE_LOADER] Successfully created bridge");
                                bridge
                            },
                            Err(e) => {
                                println!("[BRIDGE_LOADER] Panic in create_bridge: {:?}", e);
                                return Err(anyhow!("Panic in create_bridge"));
                            }
                        }
                    },
                    Err(e) => {
                        println!("[BRIDGE_LOADER] Failed to find create_bridge symbol: {:?}", e);
                        return Err(anyhow!("Failed to find create_bridge symbol: {:?}", e));
                    }
                }
            };

            println!("[BRIDGE_LOADER] Initializing bridge with config");
            bridge.initialize(config_value)?;
            println!("[BRIDGE_LOADER] Adding bridge to list");
            bridges.push(bridge);
            libraries_to_forget.push(library);
        }

        // Now register all bridges
        println!("[BRIDGE_LOADER] Registering {} bridges", bridges.len());
        for (i, bridge) in bridges.iter_mut().enumerate() {
            println!("[BRIDGE_LOADER] Registering bridge {}", i);
            match self.registry_mut().register_bridge(&mut **bridge) {
                Ok(_) => println!("[BRIDGE_LOADER] Successfully registered bridge {}", i),
                Err(e) => println!("[BRIDGE_LOADER] Failed to register bridge {}: {:?}", i, e)
            }
        }

        // We intentionally leak the libraries to keep them loaded
        // This is necessary because the bridge instances depend on the libraries
        println!("[BRIDGE_LOADER] Leaking {} libraries to keep them loaded", libraries_to_forget.len());
        for library in libraries_to_forget {
            std::mem::forget(library);
        }

        println!("[BRIDGE_LOADER] Completed load_bridges successfully");
        Ok(bridges)
    }
}

/// Static bridge loader for loading bridges from compiled code
pub struct StaticBridgeLoader {
    /// Functions that create bridge instances
    factories: Vec<Box<dyn Fn() -> BoxedBridge + Send + Sync>>,
    /// Registry for managing bridge callbacks
    registry: BridgeRegistry,
}

impl StaticBridgeLoader {
    /// Create a new static bridge loader
    pub fn new() -> Self {
        Self {
            factories: Vec::new(),
            registry: BridgeRegistry::new(),
        }
    }

    /// Register a factory function that creates a bridge
    pub fn register_factory<F>(&mut self, factory: F)
    where
        F: Fn() -> BoxedBridge + Send + Sync + 'static,
    {
        self.factories.push(Box::new(factory));
    }

    /// Get a reference to the registry
    pub fn registry(&self) -> &BridgeRegistry {
        &self.registry
    }

    /// Get a mutable reference to the registry
    pub fn registry_mut(&mut self) -> &mut BridgeRegistry {
        &mut self.registry
    }
}

impl BridgeLoader for StaticBridgeLoader {
    fn load_bridges(&mut self) -> Result<Vec<BoxedBridge>> {
        // First, create all bridges without registering them
        let mut bridges = Vec::new();
        for factory in &self.factories {
            let bridge = factory();
            bridges.push(bridge);
        }

        // Now register all bridges
        for bridge in &mut bridges {
            self.registry_mut().register_bridge(&mut **bridge)?;
        }

        Ok(bridges)
    }
}