//! Common functionality shared between command implementations.

use clap::Parser;
use std::net::TcpListener;
use std::{fs::File, path::Path, time::Duration};
use wasmtime::{
    Engine, Module, Precompiled, Result, StoreLimits, StoreLimitsBuilder, bail,
    error::Context as _, format_err,
};
use wasmtime_cli_flags::{CommonOptions, opt::WasmtimeOptionValue};
use wasmtime_wasi::WasiCtxBuilder;

#[cfg(feature = "component-model")]
use wasmtime::component::Component;

/// Whether or not WASIp3 is enabled by default.
///
/// Currently this is disabled (the `&& false`), but that'll get removed in the
/// future.
pub const P3_DEFAULT: bool = cfg!(feature = "component-model-async") && false;

#[derive(Clone)]
pub enum RunTarget {
    Core(Module),

    #[cfg(feature = "component-model")]
    Component(Component),
}

impl RunTarget {
    pub fn unwrap_core(&self) -> &Module {
        match self {
            RunTarget::Core(module) => module,
            #[cfg(feature = "component-model")]
            RunTarget::Component(_) => panic!("expected a core wasm module, not a component"),
        }
    }

    #[cfg(feature = "component-model")]
    pub fn unwrap_component(&self) -> &Component {
        match self {
            RunTarget::Component(c) => c,
            RunTarget::Core(_) => panic!("expected a component, not a core wasm module"),
        }
    }
}

/// Common command line arguments for run commands.
#[derive(Parser)]
pub struct RunCommon {
    #[command(flatten)]
    pub common: CommonOptions,

    /// Allow executing precompiled WebAssembly modules as `*.cwasm` files.
    ///
    /// Note that this option is not safe to pass if the module being passed in
    /// is arbitrary user input. Only `wasmtime`-precompiled modules generated
    /// via the `wasmtime compile` command or equivalent should be passed as an
    /// argument with this option specified.
    #[arg(long = "allow-precompiled")]
    pub allow_precompiled: bool,

    /// Profiling strategy (valid options are: perfmap, jitdump, vtune, guest)
    ///
    /// The perfmap, jitdump, and vtune profiling strategies integrate Wasmtime
    /// with external profilers such as `perf`. The guest profiling strategy
    /// enables in-process sampling and will write the captured profile to
    /// `wasmtime-guest-profile.json` by default which can be viewed at
    /// https://profiler.firefox.com/.
    ///
    /// The `guest` option can be additionally configured as:
    ///
    ///     --profile=guest[,path[,interval]]
    ///
    /// where `path` is where to write the profile and `interval` is the
    /// duration between samples. When used with `--wasm-timeout` the timeout
    /// will be rounded up to the nearest multiple of this interval.
    #[arg(
        long,
        value_name = "STRATEGY",
        value_parser = Profile::parse,
    )]
    pub profile: Option<Profile>,

    /// Grant access of a host directory to a guest.
    ///
    /// If specified as just `HOST_DIR` then the same directory name on the
    /// host is made available within the guest. If specified as `HOST::GUEST`
    /// then the `HOST` directory is opened and made available as the name
    /// `GUEST` in the guest.
    #[arg(long = "dir", value_name = "HOST_DIR[::GUEST_DIR]", value_parser = parse_dirs)]
    pub dirs: Vec<(String, String)>,

    /// Pass an environment variable to the program.
    ///
    /// The `--env FOO=BAR` form will set the environment variable named `FOO`
    /// to the value `BAR` for the guest program using WASI. The `--env FOO`
    /// form will set the environment variable named `FOO` to the same value it
    /// has in the calling process for the guest, or in other words it will
    /// cause the environment variable `FOO` to be inherited.
    #[arg(long = "env", number_of_values = 1, value_name = "NAME[=VAL]", value_parser = parse_env_var)]
    pub vars: Vec<(String, Option<String>)>,

    /// Attach the built-in gdbstub debugger component, listening on
    /// the given TCP address. Accepts a port number (e.g. `1234`) or
    /// a full `address:port`. A bare port number will bind on
    /// localhost only (`127.0.0.1`). A debugger (e.g. LLDB) can then
    /// connect via `process connect --plugin=wasm
    /// connect://<ADDR>:<PORT>`.
    #[cfg(feature = "gdbstub")]
    #[arg(short = 'g', long = "gdbstub", value_name = "[ADDR:]PORT")]
    pub gdbstub: Option<String>,
}

fn parse_env_var(s: &str) -> Result<(String, Option<String>)> {
    let mut parts = s.splitn(2, '=');
    Ok((
        parts.next().unwrap().to_string(),
        parts.next().map(|s| s.to_string()),
    ))
}

fn parse_dirs(s: &str) -> Result<(String, String)> {
    let mut parts = s.split("::");
    let host = parts.next().unwrap();
    let guest = match parts.next() {
        Some(guest) => guest,
        None => host,
    };
    Ok((host.into(), guest.into()))
}

impl RunCommon {
    pub fn store_limits(&self) -> StoreLimits {
        let mut limits = StoreLimitsBuilder::new();
        if let Some(max) = self.common.wasm.max_memory_size {
            limits = limits.memory_size(max);
        }
        if let Some(max) = self.common.wasm.max_table_elements {
            limits = limits.table_elements(max);
        }
        if let Some(max) = self.common.wasm.max_instances {
            limits = limits.instances(max);
        }
        if let Some(max) = self.common.wasm.max_tables {
            limits = limits.tables(max);
        }
        if let Some(max) = self.common.wasm.max_memories {
            limits = limits.memories(max);
        }
        if let Some(enable) = self.common.wasm.trap_on_grow_failure {
            limits = limits.trap_on_grow_failure(enable);
        }

        limits.build()
    }

    pub fn ensure_allow_precompiled(&self) -> Result<()> {
        if self.allow_precompiled {
            Ok(())
        } else {
            bail!("running a precompiled module requires the `--allow-precompiled` flag")
        }
    }

    #[cfg(feature = "component-model")]
    fn ensure_allow_components(&self) -> Result<()> {
        if self.common.wasm.component_model == Some(false) {
            bail!("cannot execute a component without `--wasm component-model`");
        }

        Ok(())
    }

    pub fn load_module(
        &self,
        engine: &Engine,
        path: &Path,
        preloaded_bytes: Option<&[u8]>,
    ) -> Result<RunTarget> {
        let path = match path.to_str() {
            #[cfg(unix)]
            Some("-") => "/dev/stdin".as_ref(),
            _ => path,
        };
        if let Some(bytes) = preloaded_bytes {
            self.load_module_contents(
                engine,
                path,
                &bytes,
                || unsafe { Module::deserialize(engine, &bytes) },
                #[cfg(feature = "component-model")]
                || unsafe { Component::deserialize(engine, &bytes) },
            )
        } else {
            let file =
                File::open(path).with_context(|| format!("failed to open wasm module {path:?}"))?;

            // First attempt to load the module as an mmap. If this succeeds then
            // detection can be done with the contents of the mmap and if a
            // precompiled module is detected then `deserialize_file` can be used
            // which is a slightly more optimal version than `deserialize` since we
            // can leave most of the bytes on disk until they're referenced.
            //
            // If the mmap fails, for example if stdin is a pipe, then fall back to
            // `std::fs::read` to load the contents. At that point precompiled
            // modules must go through the `deserialize` functions.
            //
            // Note that this has the unfortunate side effect for precompiled
            // modules on disk that they're opened once to detect what they are and
            // then again internally in Wasmtime as part of the `deserialize_file`
            // API. Currently there's no way to pass the `MmapVec` here through to
            // Wasmtime itself (that'd require making `MmapVec` a public type, both
            // which isn't ready to happen at this time). It's hoped though that
            // opening a file twice isn't too bad in the grand scheme of things with
            // respect to the CLI.
            match wasmtime::_internal::MmapVec::from_file(file) {
                Ok(map) => self.load_module_contents(
                    engine,
                    path,
                    &map,
                    || unsafe { Module::deserialize_file(engine, path) },
                    #[cfg(feature = "component-model")]
                    || unsafe { Component::deserialize_file(engine, path) },
                ),
                Err(_) => {
                    let bytes = std::fs::read(path)
                        .with_context(|| format!("failed to read file: {}", path.display()))?;
                    self.load_module_contents(
                        engine,
                        path,
                        &bytes,
                        || unsafe { Module::deserialize(engine, &bytes) },
                        #[cfg(feature = "component-model")]
                        || unsafe { Component::deserialize(engine, &bytes) },
                    )
                }
            }
        }
    }

    pub fn load_module_contents(
        &self,
        engine: &Engine,
        path: &Path,
        bytes: &[u8],
        deserialize_module: impl FnOnce() -> Result<Module>,
        #[cfg(feature = "component-model")] deserialize_component: impl FnOnce() -> Result<Component>,
    ) -> Result<RunTarget> {
        Ok(match Engine::detect_precompiled(bytes) {
            Some(Precompiled::Module) => {
                self.ensure_allow_precompiled()?;
                RunTarget::Core(deserialize_module()?)
            }
            #[cfg(feature = "component-model")]
            Some(Precompiled::Component) => {
                self.ensure_allow_precompiled()?;
                self.ensure_allow_components()?;
                RunTarget::Component(deserialize_component()?)
            }
            #[cfg(not(feature = "component-model"))]
            Some(Precompiled::Component) => {
                bail!("support for components was not enabled at compile time");
            }
            #[cfg(any(feature = "cranelift", feature = "winch"))]
            None => {
                let mut code = wasmtime::CodeBuilder::new(engine);
                code.wasm_binary_or_text(bytes, Some(path))?;
                match code.hint() {
                    Some(wasmtime::CodeHint::Component) => {
                        #[cfg(feature = "component-model")]
                        {
                            self.ensure_allow_components()?;
                            RunTarget::Component(code.compile_component()?)
                        }
                        #[cfg(not(feature = "component-model"))]
                        {
                            bail!("support for components was not enabled at compile time");
                        }
                    }
                    Some(wasmtime::CodeHint::Module) | None => {
                        RunTarget::Core(code.compile_module()?)
                    }
                }
            }

            #[cfg(not(any(feature = "cranelift", feature = "winch")))]
            None => {
                let _ = (path, engine);
                bail!("support for compiling modules was disabled at compile time");
            }
        })
    }

    pub fn configure_wasip2(&self, builder: &mut WasiCtxBuilder) -> Result<()> {
        // It's ok to block the current thread since we're the only thread in
        // the program as the CLI. This helps improve the performance of some
        // blocking operations in WASI, for example, by skipping the
        // back-and-forth between sync and async.
        //
        // However, do not set this if a timeout is configured, as that would
        // cause the timeout to be ignored if the guest does, for example,
        // something like `sleep(FOREVER)`.
        builder.allow_blocking_current_thread(self.common.wasm.timeout.is_none());

        if self.common.wasi.inherit_env == Some(true) {
            for (k, v) in std::env::vars() {
                builder.env(&k, &v);
            }
        }
        for (key, value) in self.vars.iter() {
            let value = match value {
                Some(value) => value.clone(),
                None => match std::env::var_os(key) {
                    Some(val) => val
                        .into_string()
                        .map_err(|_| format_err!("environment variable `{key}` not valid utf-8"))?,
                    None => {
                        // leave the env var un-set in the guest
                        continue;
                    }
                },
            };
            builder.env(key, &value);
        }

        for (host, guest) in self.dirs.iter() {
            builder.preopened_dir(
                host,
                guest,
                wasmtime_wasi::DirPerms::all(),
                wasmtime_wasi::FilePerms::all(),
            )?;
        }

        if self.common.wasi.listenfd == Some(true) {
            bail!("components do not support --listenfd");
        }
        for _ in self.compute_preopen_sockets()? {
            bail!("components do not support --tcplisten");
        }

        if self.common.wasi.inherit_network == Some(true) {
            builder.inherit_network();
        }
        if let Some(enable) = self.common.wasi.allow_ip_name_lookup {
            builder.allow_ip_name_lookup(enable);
        }
        if let Some(enable) = self.common.wasi.tcp {
            builder.allow_tcp(enable);
        }
        if let Some(enable) = self.common.wasi.udp {
            builder.allow_udp(enable);
        }
        if let Some(max_size) = self.common.wasi.max_random_size {
            builder.max_random_size(max_size);
        }

        Ok(())
    }

    #[cfg(feature = "wasi-http")]
    pub fn wasi_http_ctx(&self) -> Result<wasmtime_wasi_http::WasiHttpCtx> {
        let mut http = wasmtime_wasi_http::WasiHttpCtx::new();
        if let Some(limit) = self.common.wasi.max_http_fields_size {
            http.set_field_size_limit(limit);
        }
        Ok(http)
    }

    #[cfg(feature = "wasi-http")]
    pub fn wasi_http_hooks(&self) -> HttpHooks {
        HttpHooks {
            p2_outgoing_body_buffer_chunks: self
                .common
                .wasi
                .http_outgoing_body_buffer_chunks
                .unwrap_or_else(|| wasmtime_wasi_http::p2::DEFAULT_OUTGOING_BODY_BUFFER_CHUNKS),
            p2_outgoing_body_chunk_size: self
                .common
                .wasi
                .http_outgoing_body_chunk_size
                .unwrap_or_else(|| wasmtime_wasi_http::p2::DEFAULT_OUTGOING_BODY_CHUNK_SIZE),
        }
    }

    pub fn compute_preopen_sockets(&self) -> Result<Vec<TcpListener>> {
        let mut listeners = vec![];

        for address in &self.common.wasi.tcplisten {
            let stdlistener = std::net::TcpListener::bind(address)
                .with_context(|| format!("failed to bind to address '{address}'"))?;

            let _ = stdlistener.set_nonblocking(true)?;

            listeners.push(stdlistener)
        }
        Ok(listeners)
    }

    pub fn validate_p3_option(&self) -> Result<()> {
        let p3 = self.common.wasi.p3.unwrap_or(P3_DEFAULT);
        if p3 && !cfg!(feature = "component-model-async") {
            bail!("support for WASIp3 disabled at compile time");
        }
        Ok(())
    }

    pub fn validate_cli_enabled(&self) -> Result<Option<bool>> {
        let mut cli = self.common.wasi.cli;

        // Accept -Scommon as a deprecated alias for -Scli.
        if let Some(common) = self.common.wasi.common {
            if cli.is_some() {
                bail!(
                    "The -Scommon option should not be use with -Scli as it is a deprecated alias"
                );
            } else {
                // In the future, we may add a warning here to tell users to use
                // `-S cli` instead of `-S common`.
                cli = Some(common);
            }
        }

        Ok(cli)
    }

    /// Adds `wasmtime-wasi` interfaces (dubbed "-Scli" in the flags to the
    /// `wasmtime` command) to the `linker` provided.
    ///
    /// This will handle adding various WASI standard versions to the linker
    /// internally.
    #[cfg(feature = "component-model")]
    pub fn add_wasmtime_wasi_to_linker<T>(
        &self,
        linker: &mut wasmtime::component::Linker<T>,
    ) -> Result<()>
    where
        T: wasmtime_wasi::WasiView,
    {
        let mut p2_options = wasmtime_wasi::p2::bindings::LinkOptions::default();
        p2_options.cli_exit_with_code(self.common.wasi.cli_exit_with_code.unwrap_or(false));
        p2_options.network_error_code(self.common.wasi.network_error_code.unwrap_or(false));
        wasmtime_wasi::p2::add_to_linker_with_options_async(linker, &p2_options)?;

        #[cfg(feature = "component-model-async")]
        if self.common.wasi.p3.unwrap_or(P3_DEFAULT) {
            let mut p3_options = wasmtime_wasi::p3::bindings::LinkOptions::default();
            p3_options.cli_exit_with_code(self.common.wasi.cli_exit_with_code.unwrap_or(false));
            wasmtime_wasi::p3::add_to_linker_with_options(linker, &p3_options)
                .context("failed to link `wasi:cli@0.3.x`")?;
        }

        Ok(())
    }
}

#[derive(Clone, PartialEq)]
pub enum Profile {
    Native(wasmtime::ProfilingStrategy),
    Guest { path: String, interval: Duration },
}

impl Profile {
    /// Parse the `profile` argument to either the `run` or `serve` commands.
    pub fn parse(s: &str) -> Result<Profile> {
        let parts = s.split(',').collect::<Vec<_>>();
        match &parts[..] {
            ["perfmap"] => Ok(Profile::Native(wasmtime::ProfilingStrategy::PerfMap)),
            ["jitdump"] => Ok(Profile::Native(wasmtime::ProfilingStrategy::JitDump)),
            ["vtune"] => Ok(Profile::Native(wasmtime::ProfilingStrategy::VTune)),
            ["pulley"] => Ok(Profile::Native(wasmtime::ProfilingStrategy::Pulley)),
            ["guest"] => Ok(Profile::Guest {
                path: "wasmtime-guest-profile.json".to_string(),
                interval: Duration::from_millis(10),
            }),
            ["guest", path] => Ok(Profile::Guest {
                path: path.to_string(),
                interval: Duration::from_millis(10),
            }),
            ["guest", path, dur] => Ok(Profile::Guest {
                path: path.to_string(),
                interval: WasmtimeOptionValue::parse(Some(dur))?,
            }),
            _ => bail!("unknown profiling strategy: {s}"),
        }
    }
}

#[derive(Copy, Clone, Debug)]
#[cfg(feature = "wasi-http")]
pub struct HttpHooks {
    p2_outgoing_body_buffer_chunks: usize,
    p2_outgoing_body_chunk_size: usize,
}

#[cfg(feature = "wasi-http")]
impl Default for HttpHooks {
    fn default() -> Self {
        Self {
            p2_outgoing_body_buffer_chunks:
                wasmtime_wasi_http::p2::DEFAULT_OUTGOING_BODY_BUFFER_CHUNKS,
            p2_outgoing_body_chunk_size: wasmtime_wasi_http::p2::DEFAULT_OUTGOING_BODY_CHUNK_SIZE,
        }
    }
}

#[cfg(feature = "wasi-http")]
impl wasmtime_wasi_http::p2::WasiHttpHooks for HttpHooks {
    fn outgoing_body_buffer_chunks(&mut self) -> usize {
        self.p2_outgoing_body_buffer_chunks
    }

    fn outgoing_body_chunk_size(&mut self) -> usize {
        self.p2_outgoing_body_chunk_size
    }
}

#[cfg(feature = "wasi-http")]
impl wasmtime_wasi_http::p3::WasiHttpHooks for HttpHooks {}