tairitsu 0.4.4

//! Container - Generic WASM component container
//!
//! This module provides a generic container implementation that can run any WASM component.
//! Users need to implement WIT interface bindings and initialization themselves.

use anyhow::{Context as AnyhowContext, Result};

use wasmtime::{
    Store,
    component::{Component, Linker},
    error::Context,
};
use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView};

use crate::Image;
#[cfg(feature = "dynamic")]
use crate::dynamic::host_imports::HostImportRegistry;

/// Base trait for host state
///
/// Users need to implement this trait to provide their own host functionality
pub trait HostStateImpl: WasiView + Send + 'static {
    /// Get user-defined state
    fn as_any_mut(&mut self) -> &mut dyn std::any::Any;
}

/// Default host state implementation
///
/// Provides basic WASI support, users can extend through inheritance or composition
pub struct HostState {
    wasi: WasiCtx,
    table: ResourceTable,
}

impl HostState {
    /// Create a new host state
    pub fn new() -> Result<Self> {
        let wasi = WasiCtxBuilder::new()
            .inherit_stdio()
            .inherit_network()
            .build();

        let table = ResourceTable::new();

        Ok(Self { wasi, table })
    }

    /// Create host state with custom WASI configuration
    pub fn with_wasi<F>(f: F) -> Result<Self>
    where
        F: FnOnce(&mut WasiCtxBuilder) -> &mut WasiCtxBuilder,
    {
        let mut builder = WasiCtxBuilder::new();
        f(&mut builder);
        let wasi = builder.build();

        let table = ResourceTable::new();

        Ok(Self { wasi, table })
    }
}

impl Default for HostState {
    fn default() -> Self {
        Self::new().expect("Failed to create default HostState")
    }
}

impl WasiView for HostState {
    fn ctx(&mut self) -> WasiCtxView<'_> {
        WasiCtxView {
            ctx: &mut self.wasi,
            table: &mut self.table,
        }
    }
}

impl HostStateImpl for HostState {
    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
        self
    }
}

/// Handle to a guest instance
///
/// This type wraps the instance type generated by WIT bindings.
/// Users need to interact with it through `GuestHandlerContext`.
pub struct GuestInstance {
    // The actual WIT binding instance is provided by the user when building Container
    // Box is used here to store any type
    inner: Box<dyn std::any::Any + Send + Sync>,

    // Dynamic instance reference (optional, controlled by cfg feature)
    #[cfg(feature = "dynamic")]
    dynamic_instance: Option<wasmtime::component::Instance>,
}

impl GuestInstance {
    /// Create a new guest instance
    ///
    /// # Arguments
    /// * `instance` - Instance type generated by WIT bindgen
    pub fn new<T: 'static + Send + Sync>(instance: T) -> Self {
        Self {
            inner: Box::new(instance),
            #[cfg(feature = "dynamic")]
            dynamic_instance: None,
        }
    }

    /// Get the underlying instance
    ///
    /// # Type Parameters
    /// * `T` - Instance type generated by WIT bindgen
    pub fn downcast_ref<T: 'static>(&self) -> Option<&T> {
        self.inner.downcast_ref::<T>()
    }

    /// Get mutable reference to the underlying instance
    pub fn downcast_mut<T: 'static>(&mut self) -> Option<&mut T> {
        self.inner.downcast_mut::<T>()
    }

    /// Get export function for dynamic invocation (requires `dynamic` feature)
    #[cfg(feature = "dynamic")]
    pub fn get_export_func<T: HostStateImpl>(
        &self,
        _store: &mut Store<T>,
        _func_name: &str,
    ) -> Result<Option<wasmtime::component::Func>> {
        if let Some(ref instance) = self.dynamic_instance {
            Ok(instance.get_func(_store, _func_name))
        } else {
            Ok(None)
        }
    }

    /// Get the dynamic instance reference (for internal use during Container build)
    #[cfg(feature = "dynamic")]
    pub(crate) fn get_dynamic_instance_ref(&self) -> Option<&wasmtime::component::Instance> {
        self.dynamic_instance.as_ref()
    }
}

/// Context for building guest instances
///
/// Provides Linker, Store, and Component, allowing users to register their own WIT interfaces
pub struct GuestHandlerContext<'a, T: HostStateImpl> {
    pub linker: &'a mut Linker<T>,
    pub store: &'a mut Store<T>,
    pub component: &'a Component,
}

impl<'a, T: HostStateImpl> GuestHandlerContext<'a, T> {
    /// Create new context
    pub fn new(
        linker: &'a mut Linker<T>,
        store: &'a mut Store<T>,
        component: &'a Component,
    ) -> Self {
        Self {
            linker,
            store,
            component,
        }
    }
}

/// Container builder
///
/// Used to configure and create a WASM container instance
pub struct ContainerBuilder<T: HostStateImpl> {
    image: Image,
    host_state: T,
    #[allow(clippy::type_complexity)]
    host_linker_init: Option<Box<dyn FnOnce(&mut Linker<T>) -> Result<(), anyhow::Error> + Send>>,
    #[allow(clippy::type_complexity)]
    guest_initializer: Option<
        Box<
            dyn for<'a> FnOnce(GuestHandlerContext<'a, T>) -> Result<GuestInstance, anyhow::Error>
                + Send,
        >,
    >,
}

impl<T: HostStateImpl> ContainerBuilder<T>
where
    T: Default,
{
    /// Create builder from Image
    pub fn new(image: Image) -> Self {
        Self {
            image,
            host_state: T::default(),
            host_linker_init: None,
            guest_initializer: None,
        }
    }
}

impl<T: HostStateImpl> ContainerBuilder<T> {
    /// Use custom host state
    pub fn with_host_state(mut self, state: T) -> Self {
        self.host_state = state;
        self
    }

    /// Register additional host imports with the linker before guest instantiation.
    ///
    /// This is called after WASI is added to the linker but before the guest
    /// initializer runs, allowing host imports to be available when the component
    /// is instantiated.
    ///
    /// # Example
    /// ```ignore
    /// let container = Container::builder(image)
    ///     .with_host_linker(|linker| {
    ///         MyHostImpl::add_to_linker(linker, |state| &mut state.my_data)
    ///     })
    ///     .with_guest_initializer(|ctx| { ... })
    ///     .build();
    /// ```
    pub fn with_host_linker<F>(mut self, func: F) -> Self
    where
        F: FnOnce(&mut Linker<T>) -> Result<(), anyhow::Error> + Send + 'static,
    {
        self.host_linker_init = Some(Box::new(func));
        self
    }

    /// Set guest instance initializer
    ///
    /// # Arguments
    /// * `f` - A closure that receives `GuestHandlerContext` and returns WIT-bound instance
    ///
    /// # Example
    /// ```ignore
    /// let container = Container::builder(image)?
    ///     .with_guest_initializer(|ctx| {
    ///         // Register your WIT interface to linker
    ///         MyWit::add_to_linker(ctx.linker, |state| &mut state.my_state)?;
    ///
    ///         // Instantiate component
    ///         let instance = MyWit::instantiate(ctx.store, ctx.component, ctx.linker)?;
    ///
    ///         Ok(GuestInstance::new(instance))
    ///     })?
    ///     .build();
    /// ```
    pub fn with_guest_initializer<F>(mut self, f: F) -> Self
    where
        F: for<'a> FnOnce(GuestHandlerContext<'a, T>) -> Result<GuestInstance, anyhow::Error>
            + Send
            + 'static,
    {
        self.guest_initializer = Some(Box::new(f));
        self
    }

    /// Build container
    pub fn build(self) -> Result<Container<T>> {
        let mut store = Store::new(self.image.engine(), self.host_state);

        let mut linker = Linker::new(self.image.engine());
        wasmtime_wasi::p2::add_to_linker_sync(&mut linker)
            .context("Failed to add WASI to linker")?;

        // Apply custom host linker configuration (e.g. registering host imports).
        if let Some(linker_init) = self.host_linker_init {
            linker_init(&mut linker).context("Failed to configure host linker")?;
        }

        // Clone the component for type introspection
        // Component is reference-counted internally, so this is cheap
        let component = self.image.component().clone();

        let guest_instance = if let Some(initializer) = self.guest_initializer {
            let ctx = GuestHandlerContext::new(&mut linker, &mut store, &component);
            initializer(ctx)?
        } else {
            return Err(anyhow::anyhow!(
                "Guest initializer is required. Use with_guest_initializer() to set it."
            ));
        };

        // Extract dynamic instance from guest_instance if available
        #[cfg(feature = "dynamic")]
        let dynamic_instance = guest_instance.get_dynamic_instance_ref().cloned();

        Ok(Container {
            store,
            guest: guest_instance,
            #[cfg(feature = "dynamic")]
            dynamic_instance,
            #[cfg(feature = "dynamic")]
            host_imports: None,
        })
    }
}

/// A Container represents a running instance of an Image
///
/// Similar to Docker containers, it maintains runtime state and can be started/stopped
pub struct Container<T: HostStateImpl = HostState> {
    store: Store<T>,
    guest: GuestInstance,

    /// Dynamic instance for runtime function invocation (duplicated from guest for easier access)
    #[cfg(feature = "dynamic")]
    dynamic_instance: Option<wasmtime::component::Instance>,

    /// Host import registry for dynamic host function invocation
    #[cfg(feature = "dynamic")]
    host_imports: Option<HostImportRegistry>,
}

impl Container {
    /// Create container builder from Image
    pub fn builder(image: Image) -> ContainerBuilder<HostState> {
        ContainerBuilder::new(image)
    }
}

impl<T: HostStateImpl> Container<T> {
    /// Get mutable reference to Store
    pub fn store_mut(&mut self) -> &mut Store<T> {
        &mut self.store
    }

    /// Get immutable reference to Store
    pub fn store(&self) -> &Store<T> {
        &self.store
    }

    /// Get reference to guest instance
    pub fn guest(&self) -> &GuestInstance {
        &self.guest
    }

    /// Get mutable reference to guest instance
    pub fn guest_mut(&mut self) -> &mut GuestInstance {
        &mut self.guest
    }

    /// Get mutable reference to host state
    pub fn host_state_mut(&mut self) -> &mut T {
        self.store.data_mut()
    }

    /// Get immutable reference to host state
    pub fn host_state(&self) -> &T {
        self.store.data()
    }

    /// Call a guest function by name with JSON payload
    ///
    /// This is a convenience method for dynamic invocation.
    /// Requires that the guest instance implements the appropriate interface.
    ///
    /// # Arguments
    /// * `function_name` - Name of the function to call
    /// * `json_payload` - JSON string containing function arguments
    ///
    /// # Returns
    /// JSON string containing the result
    ///
    /// # Example
    /// ```ignore
    /// let result = container.call_guest_json("process", r#"{"input":"hello"}"#)?;
    /// ```
    pub fn call_guest_json(&mut self, function_name: &str, json_payload: &str) -> Result<String> {
        // JSON invocation is superseded by the RON-based call_guest_raw_desc() which
        // preserves Rust type fidelity. This entry-point returns an error to guide
        // callers to the preferred API.
        anyhow::bail!(
            "JSON invocation is not supported. Use call_guest_raw_desc() instead with RON format for better Rust type compatibility. Function: {}, Payload: {}",
            function_name,
            json_payload
        )
    }

    /// Call a guest function by name with raw descriptor payload
    ///
    /// This is the preferred method for dynamic invocation using type descriptors at runtime.
    /// Requires the `dynamic` feature to be enabled.
    ///
    /// Uses RON format for serialization which provides better Rust type compatibility than JSON.
    ///
    /// # Arguments
    /// * `function_name` - Name of the function to call
    /// * `raw_desc_payload` - RON string containing function arguments (using type descriptors)
    ///
    /// # Returns
    /// RON string containing the result
    ///
    /// # Example
    /// ```ignore
    /// let result = container.call_guest_raw_desc("process", r#"Request { input: "hello", count: 42 }"#)?;
    /// ```
    #[cfg(feature = "dynamic")]
    pub fn call_guest_raw_desc(
        &mut self,
        function_name: &str,
        raw_desc_payload: &str,
    ) -> Result<String> {
        use crate::dynamic::{ron_to_val, val_to_ron};
        use wasmtime::component::Val;

        // Get the function (direct field access avoids borrow checker issues)
        let instance = self
            .dynamic_instance
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Dynamic instance not available"))?;

        let func = instance
            .get_func(&mut self.store, function_name)
            .ok_or_else(|| anyhow::anyhow!("Export function not found: {}", function_name))?;

        // Get function type info
        let func_ty = func.ty(&self.store);
        let param_types: Vec<_> = func_ty.params().collect();
        let result_types: Vec<_> = func_ty.results().collect();

        // Convert raw descriptor → Val
        let mut args = Vec::new();

        if param_types.len() == 1 {
            // Single parameter - direct parsing
            // Extract just the type from the (name, type) tuple
            let param_type = &param_types[0].1;
            args.push(ron_to_val(raw_desc_payload, param_type)?);
        } else {
            // Multiple parameters - parse as sequence/array
            let ron_array = if raw_desc_payload.trim().starts_with('[') {
                raw_desc_payload.to_string()
            } else if raw_desc_payload.trim().starts_with('(') {
                // RON treats tuples differently, convert to array syntax
                raw_desc_payload.to_string()
            } else {
                format!("[{}]", raw_desc_payload)
            };

            use ron::Value as RonValue;
            let ron_value: RonValue = ron::from_str(&ron_array)?;

            if let RonValue::Seq(items) = ron_value {
                if items.len() != param_types.len() {
                    anyhow::bail!(
                        "Parameter count mismatch: expected {}, got {}",
                        param_types.len(),
                        items.len()
                    );
                }
                for (ron_val, (_param_name, param_type)) in
                    items.into_iter().zip(param_types.iter())
                {
                    args.push(ron_value_to_val(ron_val, param_type)?);
                }
            } else {
                anyhow::bail!(
                    "Invalid raw descriptor payload for function with multiple parameters"
                );
            }
        }

        // Call the function
        let mut results = vec![Val::Bool(false); result_types.len()];
        func.call(&mut self.store, &args, &mut results)
            .context("Function call failed")?;

        // Convert Val → raw descriptor (RON)
        let output_ron: Result<Vec<_>> = results.iter().map(val_to_ron).collect();
        let output_ron = output_ron.context("Failed to convert result to RON")?;

        // Format output based on return value count
        let output = match output_ron.len() {
            0 => "()".to_string(),
            1 => output_ron[0].clone(),
            _ => format!("({})", output_ron.join(", ")),
        };

        Ok(output)
    }

    /// Call a guest function by name with binary payload
    ///
    /// This is the high-performance path using canonical ABI directly.
    /// Requires the `dynamic` feature to be enabled.
    ///
    /// # Arguments
    /// * `function_name` - Name of the function to call
    /// * `args` - Arguments as Val types
    ///
    /// # Returns
    /// Results as Val types
    ///
    /// # Example
    /// ```ignore
    /// use wasmtime::component::Val;
    /// let args = vec![Val::String("hello".to_string()), Val::U32(42)];
    /// let results = container.call_guest_binary("process", &args)?;
    /// ```
    #[cfg(feature = "dynamic")]
    pub fn call_guest_binary(
        &mut self,
        function_name: &str,
        args: &[wasmtime::component::Val],
    ) -> Result<Vec<wasmtime::component::Val>> {
        use wasmtime::component::Val;

        // Get the function (direct field access avoids borrow checker issues)
        let instance = self
            .dynamic_instance
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Dynamic instance not available"))?;

        let func = instance
            .get_func(&mut self.store, function_name)
            .ok_or_else(|| anyhow::anyhow!("Export function not found: {}", function_name))?;

        // Get function type info
        let func_ty = func.ty(&self.store);
        let num_results = func_ty.results().count();

        // Call the function
        let mut results = vec![Val::Bool(false); num_results];
        func.call(&mut self.store, args, &mut results)
            .context("Function call failed")?;

        Ok(results)
    }

    /// Set the host import registry
    ///
    /// # Arguments
    /// * `registry` - Host import registry to use
    #[cfg(feature = "dynamic")]
    pub fn with_host_import_registry(&mut self, registry: HostImportRegistry) {
        self.host_imports = Some(registry);
    }

    /// Get mutable reference to host import registry
    #[cfg(feature = "dynamic")]
    pub fn host_imports_mut(&mut self) -> Option<&mut HostImportRegistry> {
        self.host_imports.as_mut()
    }

    /// Call a host import function by name with raw descriptor payload
    ///
    /// This allows dynamically calling host functions that the WASM component imports.
    /// Requires the `dynamic` feature to be enabled.
    ///
    /// # Arguments
    /// * `function_name` - Name of the host import function to call
    /// * `raw_desc_payload` - RON string containing function arguments
    ///
    /// # Returns
    /// RON string containing the result
    ///
    /// # Example
    /// ```ignore
    /// let result = container.call_host_import_raw_desc("log", r#"LogMessage { level: "info", msg: "hello" }"#)?;
    /// ```
    #[cfg(feature = "dynamic")]
    pub fn call_host_import_raw_desc(
        &mut self,
        function_name: &str,
        raw_desc_payload: &str,
    ) -> Result<String> {
        use crate::dynamic::{ron_to_val, val_to_ron};

        let registry = self
            .host_imports
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Host import registry not initialized"))?;

        // Get function signature
        let (params, _results) = registry
            .get_signature(function_name)
            .ok_or_else(|| anyhow::anyhow!("Host import not found: {}", function_name))?;

        // Convert raw descriptor → Val
        let mut args = Vec::new();
        if params.len() == 1 {
            args.push(ron_to_val(raw_desc_payload, &params[0])?);
        } else {
            // Multiple parameters - parse as sequence/array
            let payload = raw_desc_payload.trim();
            let ron_array = if payload.starts_with('[') || payload.starts_with('(') {
                raw_desc_payload.to_string()
            } else {
                format!("[{}]", raw_desc_payload)
            };

            use ron::Value as RonValue;
            let ron_value: RonValue = ron::from_str(&ron_array)?;

            if let RonValue::Seq(items) = ron_value {
                if items.len() != params.len() {
                    anyhow::bail!(
                        "Parameter count mismatch: expected {}, got {}",
                        params.len(),
                        items.len()
                    );
                }
                for (ron_val, param_type) in items.into_iter().zip(params.iter()) {
                    args.push(ron_value_to_val(ron_val, param_type)?);
                }
            } else {
                anyhow::bail!(
                    "Invalid raw descriptor payload for function with multiple parameters"
                );
            }
        }

        // Call the function
        let result_vals = registry.call(function_name, &args)?;

        // Convert Val → raw descriptor (RON)
        let output_ron: Result<Vec<_>> = result_vals.iter().map(val_to_ron).collect();
        let output_ron = output_ron.context("Failed to convert result to RON")?;
        let output = match output_ron.len() {
            0 => "()".to_string(),
            1 => output_ron[0].clone(),
            _ => format!("({})", output_ron.join(", ")),
        };

        Ok(output)
    }

    /// List all guest export functions
    ///
    /// Returns information about all functions exported by the WASM component.
    /// Requires the `dynamic` feature to be enabled.
    ///
    /// # Returns
    /// Vector of export information including function names and types
    ///
    /// # Example
    /// ```ignore
    /// let exports = container.list_guest_exports()?;
    /// for export in exports {
    ///     println!("Function: {}", export.name);
    ///     println!("  Params: {:?}", export.params);
    ///     println!("  Results: {:?}", export.results);
    /// }
    /// ```
    #[cfg(feature = "dynamic")]
    pub fn list_guest_exports(&mut self) -> Result<Vec<ExportInfo>> {
        if let Some(instance) = &self.dynamic_instance {
            let mut exports = Vec::new();

            // Try to get common export function names
            // This is a pragmatic approach since we can't easily iterate all exports
            // without knowing their names beforehand
            let potential_exports = vec![
                "init",
                "process",
                "getname",
                "getversion",
                "getfeatures",
                "shutdown",
                "notify",
                "add",
                "sub",
                "mul",
                "div",
                "to-upper",
                "to-lower",
                "reverse",
                "length",
                "process-numbers",
                "transform",
            ];

            for export_name in potential_exports {
                // Try to get the function
                if let Some(func) = instance.get_func(&mut self.store, export_name) {
                    // Get function type information
                    let func_ty = func.ty(&self.store);
                    let mut params = Vec::new();
                    let mut results = Vec::new();

                    // Extract parameter types (they come as (name, type) tuples)
                    for (param_name, param_type) in func_ty.params() {
                        params.push((param_name.to_string(), param_type.clone()));
                    }

                    // Extract result types
                    for result_type in func_ty.results() {
                        results.push(result_type.clone());
                    }

                    exports.push(ExportInfo {
                        name: export_name.to_string(),
                        params,
                        results,
                    });
                }
            }

            Ok(exports)
        } else {
            // No dynamic instance available, return empty list
            Ok(Vec::new())
        }
    }

    /// List all host import functions
    ///
    /// Returns information about all registered host import functions.
    /// Requires the `dynamic` feature to be enabled.
    ///
    /// # Returns
    /// Vector of import information including function names and types
    ///
    /// # Example
    /// ```ignore
    /// if let Some(imports) = container.list_host_imports()? {
    ///     for import in imports {
    ///         println!("Function: {}", import.name);
    ///     }
    /// }
    /// ```
    #[cfg(feature = "dynamic")]
    pub fn list_host_imports(&self) -> Result<Vec<ImportInfo>> {
        if let Some(ref registry) = self.host_imports {
            let imports = registry.list_imports();
            imports
                .into_iter()
                .map(|name| {
                    let (params, results) = registry.get_signature(name).unwrap();
                    Ok(ImportInfo {
                        name: name.to_string(),
                        params,
                        results,
                    })
                })
                .collect()
        } else {
            Ok(Vec::new())
        }
    }
}

/// Export function information
#[derive(Debug, Clone)]
pub struct ExportInfo {
    pub name: String,
    pub params: Vec<(String, wasmtime::component::Type)>,
    pub results: Vec<wasmtime::component::Type>,
}

/// Import function information
#[derive(Debug, Clone)]
pub struct ImportInfo {
    pub name: String,
    pub params: Vec<wasmtime::component::Type>,
    pub results: Vec<wasmtime::component::Type>,
}

/// Helper: RON Value to Val (for use in Container)
#[cfg(feature = "dynamic")]
fn ron_value_to_val(
    ron_value: ron::Value,
    target_type: &wasmtime::component::Type,
) -> Result<wasmtime::component::Val> {
    use crate::dynamic::ron_to_val;
    // Convert ron::Value to RON string, then use ron_to_val
    let ron_str = ron::to_string(&ron_value)?;
    ron_to_val(&ron_str, target_type)
}

impl<T: HostStateImpl> std::fmt::Debug for Container<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Container").finish()
    }
}