//! A Spin execution context for applications.

#![deny(missing_docs)]

/// Input / Output redirects.
pub mod io;

use anyhow::{bail, Context, Result};
use io::IoStreamRedirects;
use spin_config::{CoreComponent, DirectoryMount, ModuleSource};
use std::{collections::HashMap, io::Write, path::PathBuf, sync::Arc};
use tracing::{instrument, log};
use wasi_common::WasiCtx;
use wasmtime::{Engine, Instance, InstancePre, Linker, Module, Store};
use wasmtime_wasi::{ambient_authority, Dir, WasiCtxBuilder};

const SPIN_HOME: &str = ".spin";

/// Builder-specific configuration.
#[derive(Clone, Debug)]
pub struct ExecutionContextConfiguration {
    /// Spin application configuration.
    pub app: spin_config::Configuration<CoreComponent>,
    /// Wasmtime engine configuration.
    pub wasmtime: wasmtime::Config,
    /// Log directory on host
    pub log_dir: Option<PathBuf>,
}

impl ExecutionContextConfiguration {
    /// Creates a new execution context configuration.
    pub fn new(app: spin_config::Configuration<CoreComponent>, log_dir: Option<PathBuf>) -> Self {
        // In order for Wasmtime to run WebAssembly components, multi memory
        // and module linking must always be enabled.
        // See https://github.com/bytecodealliance/wit-bindgen/blob/main/crates/wasmlink.
        let mut wasmtime = wasmtime::Config::default();
        wasmtime.wasm_multi_memory(true);
        wasmtime.wasm_module_linking(true);

        log::trace!("Created execution context configuration.");

        Self {
            app,
            wasmtime,
            log_dir,
        }
    }
}

/// Top-level runtime context data to be passed to a component.
#[derive(Default)]
pub struct RuntimeContext<T> {
    /// WASI context data.
    pub wasi: Option<WasiCtx>,
    /// Experimental outbound HTTP configuration.
    pub experimental_http: Option<wasi_experimental_http_wasmtime::HttpCtx>,
    /// Outbound HTTP configuration.
    pub outbound_http: Option<wasi_outbound_http::OutboundHttp>,
    /// Generic runtime data that can be configured by specialized engines.
    pub data: Option<T>,
}

/// An execution context builder.
pub struct Builder<T: Default> {
    /// Top level configuration for
    pub config: ExecutionContextConfiguration,
    /// Linker used to configure the execution context.
    pub linker: Linker<RuntimeContext<T>>,
    /// Store used to configure the execution context.
    pub store: Store<RuntimeContext<T>>,
    /// Wasmtime engine.
    pub engine: Engine,
}

impl<T: Default> Builder<T> {
    /// Creates a new instance of the execution builder.
    pub fn new(config: ExecutionContextConfiguration) -> Result<Builder<T>> {
        let data = RuntimeContext::default();
        let engine = Engine::new(&config.wasmtime)?;
        let store = Store::new(&engine, data);
        let linker = Linker::new(&engine);

        Ok(Self {
            config,
            linker,
            store,
            engine,
        })
    }

    /// Configures the WASI linker imports for the current execution context.
    pub fn link_wasi(&mut self) -> Result<&Self> {
        wasmtime_wasi::add_to_linker(&mut self.linker, |ctx| ctx.wasi.as_mut().unwrap())?;
        Ok(self)
    }

    /// Configures the ability to execute outbound HTTP requests.
    pub fn link_http(&mut self) -> Result<&Self> {
        wasi_experimental_http_wasmtime::HttpState::new()?
            .add_to_linker(&mut self.linker, |ctx| {
                ctx.experimental_http.as_ref().unwrap()
            })?;

        wasi_outbound_http::add_to_linker(&mut self.linker, |ctx| {
            ctx.outbound_http.as_mut().unwrap()
        })?;

        Ok(self)
    }

    /// Builds a new instance of the execution context.
    #[instrument(skip(self))]
    pub async fn build(&mut self) -> Result<ExecutionContext<T>> {
        let mut components = HashMap::new();
        for c in &self.config.app.components {
            let core = c.clone();
            let module = match c.source.clone() {
                ModuleSource::FileReference(p) => {
                    let module = Module::from_file(&self.engine, &p).with_context(|| {
                        format!(
                            "Cannot create module for component {} from file {}",
                            &c.id,
                            &p.display()
                        )
                    })?;
                    log::trace!("Created module for component {} from file {:?}", &c.id, &p);
                    module
                }
                ModuleSource::Buffer(bytes, info) => {
                    let module = Module::from_binary(&self.engine, &bytes).with_context(|| {
                        format!("Cannot create module for component {} from {}", &c.id, info)
                    })?;
                    log::trace!(
                        "Created module for component {} from {} with size {}",
                        &c.id,
                        info,
                        bytes.len()
                    );
                    module
                }
            };

            let pre = Arc::new(self.linker.instantiate_pre(&mut self.store, &module)?);
            log::trace!("Created pre-instance from module for component {}.", &c.id);

            components.insert(c.id.clone(), Component { core, pre });
        }

        let config = self.config.clone();
        let engine = self.engine.clone();

        log::trace!("Execution context initialized.");

        Ok(ExecutionContext {
            config,
            engine,
            components,
        })
    }

    /// Builds a new default instance of the execution context.
    pub async fn build_default(
        config: ExecutionContextConfiguration,
    ) -> Result<ExecutionContext<T>> {
        let mut builder = Self::new(config)?;
        builder.link_wasi()?;
        builder.link_http()?;

        builder.build().await
    }
}

/// Component for the execution context.
#[derive(Clone)]
pub struct Component<T: Default> {
    /// Configuration for the component.
    pub core: CoreComponent,
    /// The pre-instance of the component
    pub pre: Arc<InstancePre<RuntimeContext<T>>>,
}

/// A generic execution context for WebAssembly components.
#[derive(Clone)]
pub struct ExecutionContext<T: Default> {
    /// Top-level runtime configuration.
    pub config: ExecutionContextConfiguration,
    /// Wasmtime engine.
    pub engine: Engine,
    /// Collection of pre-initialized (and already linked) components.
    pub components: HashMap<String, Component<T>>,
}

impl<T: Default> ExecutionContext<T> {
    /// Creates a store for a given component given its configuration and runtime data.
    #[instrument(skip(self, data, io))]
    pub fn prepare_component(
        &self,
        component: &str,
        data: Option<T>,
        io: Option<IoStreamRedirects>,
        env: Option<HashMap<String, String>>,
        args: Option<Vec<String>>,
    ) -> Result<(Store<RuntimeContext<T>>, Instance)> {
        log::trace!("Preparing component {}", component);
        let component = match self.components.get(component) {
            Some(c) => c,
            None => bail!("Cannot find component {}", component),
        };

        let mut store = self.store(component, data, io, env, args)?;
        let instance = component.pre.instantiate(&mut store)?;

        Ok((store, instance))
    }

    /// Save logs for a given component in the log directory on the host
    pub fn save_output_to_logs(
        &self,
        io_redirects: IoStreamRedirects,
        component: &str,
        save_stdout: bool,
        save_stderr: bool,
    ) -> Result<()> {
        let sanitized_app_name = sanitize(&self.config.app.info.name);
        let sanitized_component_name = sanitize(&component);

        let log_dir = match &self.config.log_dir {
            Some(l) => l.clone(),
            None => match dirs::home_dir() {
                Some(h) => h.join(SPIN_HOME).join(sanitized_app_name).join("logs"),
                None => PathBuf::from(sanitized_app_name).join("logs"),
            },
        };

        let stdout_filename = log_dir.join(sanitize(format!(
            "{}_{}.txt",
            sanitized_component_name, "stdout",
        )));

        let stderr_filename = log_dir.join(sanitize(format!(
            "{}_{}.txt",
            sanitized_component_name, "stderr"
        )));

        std::fs::create_dir_all(&log_dir)?;

        log::trace!("Saving logs to {:?} {:?}", stdout_filename, stderr_filename);

        if save_stdout {
            let mut file = std::fs::OpenOptions::new()
                .write(true)
                .append(true)
                .create(true)
                .open(stdout_filename)?;
            let contents = io_redirects.stdout.lock.read().unwrap();
            file.write_all(&contents)?;
        }

        if save_stderr {
            let contents = io_redirects.stderr.lock.read().unwrap();
            let mut file = std::fs::OpenOptions::new()
                .write(true)
                .append(true)
                .create(true)
                .open(stderr_filename)?;
            file.write_all(&contents)?;
        }

        Ok(())
    }
    /// Creates a store for a given component given its configuration and runtime data.
    fn store(
        &self,
        component: &Component<T>,
        data: Option<T>,
        io: Option<IoStreamRedirects>,
        env: Option<HashMap<String, String>>,
        args: Option<Vec<String>>,
    ) -> Result<Store<RuntimeContext<T>>> {
        log::trace!("Creating store.");
        let (env, dirs) = Self::wasi_config(component, env)?;
        let mut ctx = RuntimeContext::default();
        let mut wasi_ctx = WasiCtxBuilder::new()
            .args(&args.unwrap_or_default())?
            .envs(&env)?;
        match io {
            Some(r) => {
                wasi_ctx = wasi_ctx
                    .stderr(Box::new(r.stderr.out))
                    .stdout(Box::new(r.stdout.out))
                    .stdin(Box::new(r.stdin));
            }
            None => wasi_ctx = wasi_ctx.inherit_stdio(),
        };

        for dir in dirs {
            let guest = dir.guest;
            let host = dir.host;
            wasi_ctx =
                wasi_ctx.preopened_dir(Dir::open_ambient_dir(host, ambient_authority())?, guest)?;
        }

        // We basically have to support both versions of outbound HTTP for the time being.
        let experimental_http = wasi_experimental_http_wasmtime::HttpCtx {
            allowed_hosts: Some(component.core.wasm.allowed_http_hosts.clone()),
            max_concurrent_requests: None,
        };
        let outbound_http = wasi_outbound_http::OutboundHttp {
            allowed_hosts: Some(component.core.wasm.allowed_http_hosts.clone()),
        };

        ctx.wasi = Some(wasi_ctx.build());
        ctx.experimental_http = Some(experimental_http);
        ctx.outbound_http = Some(outbound_http);
        ctx.data = data;

        let store = Store::new(&self.engine, ctx);
        Ok(store)
    }

    #[allow(clippy::type_complexity)]
    fn wasi_config(
        component: &Component<T>,
        env: Option<HashMap<String, String>>,
    ) -> Result<(Vec<(String, String)>, Vec<DirectoryMount>)> {
        let mut res = vec![];

        for (k, v) in &component.core.wasm.environment {
            res.push((k.clone(), v.clone()));
        }

        // Custom environment variables currently take precedence over component-defined
        // environment variables. This might change in the future.
        if let Some(envs) = env {
            for (k, v) in envs {
                res.push((k.clone(), v.clone()));
            }
        };

        let dirs = component.core.wasm.mounts.clone();

        Ok((res, dirs))
    }
}

fn sanitize(name: impl AsRef<str>) -> String {
    // options block copied from sanitize_filename project readme
    let options = sanitize_filename::Options {
        // true by default, truncates to 255 bytes
        truncate: true,
        // default value depends on the OS, removes reserved names like `con` from start of strings on Windows
        windows: true,
        // str to replace sanitized chars/strings
        replacement: "",
    };

    // filename logic defined in the project works for directory names as well
    // refer to: https://github.com/kardeiz/sanitize-filename/blob/f5158746946ed81015c3a33078dedf164686da19/src/lib.rs#L76-L165
    sanitize_filename::sanitize_with_options(name, options)
}
