tairitsu 0.2.0

//! 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, Result};

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

use crate::Image;

/// 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>,
}

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),
        }
    }

    /// 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>()
    }
}

/// 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)]
    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(),
            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
    }

    /// 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")?;

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

        Ok(Container {
            store,
            guest: guest_instance,
        })
    }
}

/// 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,
}

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()
    }
}

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()
    }
}