//! The WASM tier (spec §4.1, oxygenflow#75): a job's `exec` names a WASI

//! preview 1 module, or a glob of modules run one after another, which

//! wasmtime runs with the job's workspace at `/src` and nothing else: no

//! network, no other file of the runner. Every runner can offer it, and it

//! is the tier phones and browsers run (docs/handover.md).

//!

//! As in the syscall tier the module runs in a child of the runner

//! (`flow __wasm SPEC`), so the runner streams its output as it comes,

//! stops it at the time limit or when the run is cancelled, and measures

//! its memory the same way; wasmtime itself also holds it to the sandbox's

//! memory limit.

use crate::microvm::{Outcome, Sink};

use crate::syscall::{lines, self_hash, FAILED};

use anyhow::{bail, Context, Result};

use flow_core::api::Record;

use flow_core::{Environment, Job, MountSource, Sandbox};

use serde::{Deserialize, Serialize};

use std::path::{Path, PathBuf};

use std::sync::mpsc;

use std::time::{Duration, Instant};

/// The runtime, named in every receipt this tier writes.

pub const WASMTIME: &str = "49.0.4";

/// A line the helper writes to say something about the job itself (which

/// module starts), shown as the runner's own note rather than the job's.

const NOTE: &str = "\u{1e}N ";

/// The exit code for a module that trapped (crashed) rather than exiting.

const TRAPPED: i32 = 134;

/// What the helper runs, written by the runner for `flow __wasm`.

#[derive(Serialize, Deserialize)]

struct Spec {

    /// The job's files, preopened as `/src` (and as the current directory).

    workspace: PathBuf,

    /// The modules to run in turn, relative to the workspace.

    modules: Vec<String>,

    /// The arguments after the module.

    args: Vec<String>,

    env: Vec<(String, String)>,

    memory_mib: u32,

    /// The job's mounts: a directory of the runner and where the module sees it.

    preopens: Vec<(PathBuf, String)>,

}

/// The receipt environment of a job in this tier: the runtime and its

/// version, this runner's binary, the memory allowed, and every module run

/// with its hash, so the receipt says exactly what code ran.

pub fn environment(sandbox: &Sandbox, job: &Job, modules: &[(String, String)]) -> Environment {

    Environment {

        runtime: "wasm".into(),

        parts: vec![

            ("wasmtime".into(), WASMTIME.into()),

            ("flow".into(), self_hash().to_string()),

            ("memory_mib".into(), sandbox.memory_mib.to_string()),

        ]

        .into_iter()

        .chain(modules.iter().map(|(path, hash)| (format!("module {path}"), hash.clone())))

        .collect(),

        rootfs: String::new(),

        argv: job.exec.clone(),

        env: job.env.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),

    }

}

/// The modules `program` names in `workspace`: one file, or the files a

/// glob matches (`*` in its last part only), sorted. A path may start with

/// `/src/`, where the module sees the workspace.

pub fn modules(workspace: &Path, program: &str) -> Result<Vec<String>> {

    let rel = program.strip_prefix("/src/").unwrap_or(program).trim_start_matches("./");

    if rel.is_empty() || rel.starts_with('/') || rel.split('/').any(|p| p == "..") {

        bail!("module {program:?}: give a path inside the job's files, such as target/wasm32-wasip1/debug/app.wasm");

    }

    let (dir, last) = rel.rsplit_once('/').unwrap_or(("", rel));

    if dir.contains('*') {

        bail!("module {program:?}: only the file name may hold `*`");

    }

    if !last.contains('*') {

        if !workspace.join(rel).is_file() {

            bail!("module {rel} is not among the job's files (its scope or an earlier job's artifact)");

        }

        return Ok(vec![rel.to_string()]);

    }

    let mut found: Vec<String> = std::fs::read_dir(workspace.join(dir))

        .with_context(|| format!("module {program:?}: the directory {dir} is not among the job's files"))?

        .filter_map(|e| e.ok())

        .filter(|e| e.file_type().is_ok_and(|t| t.is_file()))

        .filter_map(|e| e.file_name().into_string().ok())

        .filter(|name| wildcard(last, name))

        .map(|name| if dir.is_empty() { name } else { format!("{dir}/{name}") })

        .collect();

    found.sort();

    if found.is_empty() {

        bail!("no module matches {program:?}");

    }

    Ok(found)

}

/// Whether `name` matches `pattern`, where `*` stands for any run of characters.

fn wildcard(pattern: &str, name: &str) -> bool {

    let parts: Vec<&str> = pattern.split('*').collect();

    let (first, last) = (parts[0], parts[parts.len() - 1]);

    if parts.len() == 1 {

        return pattern == name;

    }

    if !name.starts_with(first) || !name[first.len()..].ends_with(last) || name.len() < first.len() + last.len() {

        return false;

    }

    let mut rest = &name[first.len()..name.len() - last.len()];

    for middle in &parts[1..parts.len() - 1] {

        match rest.find(middle) {

            Some(i) => rest = &rest[i + middle.len()..],

            None => return false,

        }

    }

    true

}

/// Run `job` in this tier with its files in `workspace` (as the syscall

/// tier lays them out); `scratch` holds the helper's spec and the job's

/// mounts.

pub fn run(sandbox: &Sandbox, job: &Job, workspace: &Path, scratch: &Path, timeout: Duration, sink: &mut dyn Sink) -> Result<Outcome> {

    if !job.secrets.is_empty() || !job.identities.is_empty() {

        bail!("job {}: the WASM tier takes no secrets or identities yet; run it in a microvm or syscall sandbox", job.name);

    }

    let Some(program) = job.exec.first() else { bail!("job {}: exec names no module", job.name) };

    let modules = modules(workspace, program)?;

    let hashed: Vec<(String, String)> =

        modules.iter().map(|m| Ok((m.clone(), crate::home::hash_file(&workspace.join(m))?))).collect::<Result<_>>()?;

    let environment = environment(sandbox, job, &hashed);

    let mut preopens = Vec::new();

    let mut disks = 0;

    for (i, m) in job.mounts.iter().enumerate() {

        let dir = match m.source {

            MountSource::Disk => {

                disks += 1;

                crate::syscall::disk_dir(scratch, disks - 1)

            }

            MountSource::Tmpfs => scratch.join(format!("tmpfs-{i}")),

        };

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

        preopens.push((dir, m.target.clone()));

    }

    let spec = Spec {

        workspace: workspace.to_path_buf(),

        modules,

        args: job.exec[1..].to_vec(),

        env: job.env.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),

        memory_mib: sandbox.memory_mib,

        preopens,

    };

    if in_process() {

        return run_in_thread(&spec, environment, timeout, sink);

    }

    let spec_path = scratch.join("wasm.json");

    std::fs::write(&spec_path, serde_json::to_vec(&spec)?)?;

    let started = Instant::now();

    let mut child = std::process::Command::new("/proc/self/exe")

        .arg("__wasm")

        .arg(&spec_path)

        .env_clear()

        .current_dir(scratch)

        .stdin(std::process::Stdio::null())

        .stdout(std::process::Stdio::piped())

        .stderr(std::process::Stdio::piped())

        .spawn()

        .context("start the WASM helper")?;

    let (tx, rx) = mpsc::channel::<(bool, Vec<u8>)>();

    let out = child.stdout.take().unwrap();

    let err = child.stderr.take().unwrap();

    let tx2 = tx.clone();

    let readers = [std::thread::spawn(move || lines(out, true, tx)), std::thread::spawn(move || lines(err, false, tx2))];

    let t = |s: &Instant| s.elapsed().as_millis() as u64;

    let (mut timed_out, mut cancelled, mut failed_to_start) = (false, false, false);

    let mut peak_memory_mib = None;

    let deadline = started + timeout;

    loop {

        if let Some(m) = crate::facts::peak_rss_mib(child.id()) {

            peak_memory_mib = Some(m);

        }

        match rx.recv_timeout(Duration::from_millis(200)) {

            Ok((stdout, line)) => {

                let text = String::from_utf8_lossy(&line).into_owned();

                let rec = |stream: &str, text: &str| Record { t: t(&started), stream: stream.into(), text: text.into() };

                if let Some(note) = text.strip_prefix(NOTE).filter(|_| !stdout) {

                    sink.record(rec("flow", note));

                } else if let Some(why) = text.strip_prefix(FAILED).filter(|_| !stdout) {

                    failed_to_start = true;

                    sink.record(rec("flow", &format!("the WASM tier could not run the job: {why}")));

                } else {

                    sink.record(rec(if stdout { "out" } else { "err" }, &text));

                }

            }

            Err(mpsc::RecvTimeoutError::Timeout) => {}

            Err(mpsc::RecvTimeoutError::Disconnected) => break,

        }

        if !timed_out && !cancelled && sink.cancelled() {

            cancelled = true;

            let _ = child.kill();

            sink.record(Record { t: t(&started), stream: "flow".into(), text: "stopped: the run was cancelled".into() });

        }

        if Instant::now() >= deadline && !timed_out && !cancelled {

            timed_out = true;

            let _ = child.kill();

            sink.record(Record {

                t: t(&started),

                stream: "flow".into(),

                text: format!("stopped at the time limit of {} s", timeout.as_secs()),

            });

        }

    }

    let status = child.wait()?;

    for r in readers {

        let _ = r.join();

    }

    let _ = std::fs::remove_file(&spec_path);

    let exit_code = status.code();

    if exit_code.is_none() && !timed_out && !cancelled && !failed_to_start {

        sink.record(Record {

            t: t(&started),

            stream: "flow".into(),

            text: format!("the WASM helper ended without an exit code ({status})"),

        });

    }

    Ok(Outcome {

        exit_code: if timed_out || cancelled { None } else { exit_code },

        duration_ms: t(&started),

        environment,

        timed_out,

        cancelled,

        peak_memory_mib,

        pruned: Vec::new(),

    })

}

/// Whether modules run in a thread of this process rather than in a child

/// `flow __wasm`: on Android, where the runner is a library inside an app

/// that cannot start itself as a command (oxygenhub-android#8), and

/// anywhere `FLOW_WASM_IN_PROCESS` is set (how the tests check this way).

pub fn in_process() -> bool {

    cfg!(target_os = "android") || std::env::var_os("FLOW_WASM_IN_PROCESS").is_some()

}

/// The most of a module's output kept before it is read, per stream.

const PIPE_BYTES: usize = 64 << 20;

/// Run the spec's modules in a thread of this process, streaming their

/// output as it comes: the pipes are read every 100 ms, and a time limit

/// or a cancelled run interrupts the module through wasmtime's epoch.

fn run_in_thread(spec: &Spec, environment: Environment, timeout: Duration, sink: &mut dyn Sink) -> Result<Outcome> {

    use wasmtime_wasi::p2::pipe::MemoryOutputPipe;

    let mut config = wasmtime::Config::new();

    config.epoch_interruption(true);

    let engine = ww(wasmtime::Engine::new(&config), || "start wasmtime".into())?;

    let (stdout, stderr) = (MemoryOutputPipe::new(PIPE_BYTES), MemoryOutputPipe::new(PIPE_BYTES));

    let (notes_tx, notes) = mpsc::channel::<String>();

    let io = Io::Pipes { stdout: stdout.clone(), stderr: stderr.clone(), notes: notes_tx };

    let started = Instant::now();

    let t = |s: &Instant| s.elapsed().as_millis() as u64;

    let deadline = started + timeout;

    let (mut timed_out, mut cancelled) = (false, false);

    // Bytes of each stream already passed on, and the part line not ended yet.

    let mut said = [(0usize, Vec::<u8>::new()), (0usize, Vec::<u8>::new())];

    let mut pass_on = |sink: &mut dyn Sink, last: bool| {

        for (i, pipe) in [&stdout, &stderr].into_iter().enumerate() {

            let all = pipe.contents();

            let (offset, partial) = &mut said[i];

            partial.extend_from_slice(&all[*offset..]);

            *offset = all.len();

            while let Some(n) = partial.iter().position(|b| *b == b'\n') {

                let line: Vec<u8> = partial.drain(..=n).collect();

                let text = String::from_utf8_lossy(&line[..line.len() - 1]).into_owned();

                sink.record(Record { t: t(&started), stream: if i == 0 { "out" } else { "err" }.into(), text });

            }

            if last && !partial.is_empty() {

                let text = String::from_utf8_lossy(partial).into_owned();

                partial.clear();

                sink.record(Record { t: t(&started), stream: if i == 0 { "out" } else { "err" }.into(), text });

            }

        }

        while let Ok(note) = notes.try_recv() {

            sink.record(Record { t: t(&started), stream: "flow".into(), text: note });

        }

    };

    let ran = std::thread::scope(|scope| {

        let worker = scope.spawn(|| run_modules(&engine, spec, &io));

        while !worker.is_finished() {

            std::thread::sleep(Duration::from_millis(100));

            pass_on(sink, false);

            if !timed_out && !cancelled && sink.cancelled() {

                cancelled = true;

                engine.increment_epoch();

                sink.record(Record { t: t(&started), stream: "flow".into(), text: "stopped: the run was cancelled".into() });

            }

            if Instant::now() >= deadline && !timed_out && !cancelled {

                timed_out = true;

                engine.increment_epoch();

                sink.record(Record {

                    t: t(&started),

                    stream: "flow".into(),

                    text: format!("stopped at the time limit of {} s", timeout.as_secs()),

                });

            }

        }

        worker.join().unwrap_or_else(|_| Err(anyhow::anyhow!("the WASM worker panicked")))

    });

    pass_on(sink, true);

    let (exit_code, peak) = match ran {

        Ok((code, peak)) => (Some(code), Some(peak)),

        Err(e) => {

            sink.record(Record { t: t(&started), stream: "flow".into(), text: format!("the WASM tier could not run the job: {e:#}") });

            (Some(125), None)

        }

    };

    Ok(Outcome {

        exit_code: if timed_out || cancelled { None } else { exit_code },

        duration_ms: t(&started),

        environment,

        timed_out,

        cancelled,

        peak_memory_mib: peak.map(|b| (b as u64).div_ceil(1 << 20)),

        pruned: Vec::new(),

    })

}

/// `flow __wasm SPEC`: run the spec's modules in turn, each in a fresh

/// instance; the exit code is the first non-zero one, after all have run,

/// so one failing module does not hide another's results.

pub fn wasm_main(spec_path: &Path) -> i32 {

    let ran = (|| {

        let spec: Spec = serde_json::from_slice(&std::fs::read(spec_path)?).context("read the spec")?;

        let engine = ww(wasmtime::Engine::new(&wasmtime::Config::new()), || "start wasmtime".into())?;

        run_modules(&engine, &spec, &Io::Inherit)

    })();

    match ran {

        Ok((code, _)) => code,

        Err(e) => {

            eprintln!("{FAILED}{e:#}");

            125

        }

    }

}

/// wasmtime's own error, as this crate's: with its whole chain of causes.

fn ww<T>(r: wasmtime::Result<T>, what: impl FnOnce() -> String) -> Result<T> {

    r.map_err(|e| anyhow::anyhow!("{}: {e:#}", what()))

}

/// Where a module's output and the tier's notes go: inherited by the

/// helper process (notes on stderr, marked), or pipes the runner reads.

enum Io {

    Inherit,

    Pipes {

        stdout: wasmtime_wasi::p2::pipe::MemoryOutputPipe,

        stderr: wasmtime_wasi::p2::pipe::MemoryOutputPipe,

        notes: mpsc::Sender<String>,

    },

}

impl Io {

    fn note(&self, text: String) {

        match self {

            Io::Inherit => eprintln!("{NOTE}{text}"),

            Io::Pipes { notes, .. } => {

                let _ = notes.send(text);

            }

        }

    }

}

/// The sandbox's memory limit, and the most a module asked for within it.

struct Limits {

    memory: usize,

    peak: usize,

}

impl wasmtime::ResourceLimiter for Limits {

    fn memory_growing(&mut self, _current: usize, desired: usize, maximum: Option<usize>) -> wasmtime::Result<bool> {

        let ok = desired <= self.memory && maximum.is_none_or(|m| desired <= m);

        if ok {

            self.peak = self.peak.max(desired);

        }

        Ok(ok)

    }

    fn table_growing(&mut self, _current: usize, desired: usize, maximum: Option<usize>) -> wasmtime::Result<bool> {

        Ok(maximum.is_none_or(|m| desired <= m))

    }

}

struct State {

    wasi: wasmtime_wasi::p1::WasiP1Ctx,

    limits: Limits,

}

/// Run each module in turn: the first non-zero exit code after all ran,

/// and the most memory any of them held, in bytes.

fn run_modules(engine: &wasmtime::Engine, spec: &Spec, io: &Io) -> Result<(i32, usize)> {

    let mut linker: wasmtime::Linker<State> = wasmtime::Linker::new(engine);

    ww(wasmtime_wasi::p1::add_to_linker_sync(&mut linker, |s: &mut State| &mut s.wasi), || "link WASI".into())?;

    let (mut first_failure, mut peak) = (0, 0);

    for path in &spec.modules {

        if spec.modules.len() > 1 {

            io.note(format!("running {path}"));

        }

        let (code, held) = run_module(engine, &linker, spec, path, io)?;

        peak = peak.max(held);

        if spec.modules.len() > 1 {

            io.note(format!("{path}: exit {code}"));

        }

        if first_failure == 0 && code != 0 {

            first_failure = code;

        }

    }

    Ok((first_failure, peak))

}

fn run_module(engine: &wasmtime::Engine, linker: &wasmtime::Linker<State>, spec: &Spec, path: &str, io: &Io) -> Result<(i32, usize)> {

    use wasmtime_wasi::{FsPerms, WasiCtxBuilder};

    let module = ww(wasmtime::Module::from_file(engine, spec.workspace.join(path)), || format!("module {path} does not load"))?;

    let name = path.rsplit('/').next().unwrap_or(path);

    let mut argv = vec![name.to_string()];

    argv.extend(spec.args.iter().cloned());

    let mut wasi = WasiCtxBuilder::new();

    match io {

        Io::Inherit => {

            wasi.inherit_stdout().inherit_stderr();

        }

        Io::Pipes { stdout, stderr, .. } => {

            wasi.stdout(stdout.clone()).stderr(stderr.clone());

        }

    }

    wasi.args(&argv).envs(&spec.env).env("PWD", "/src").allow_blocking_current_thread(true);

    let open = |wasi: &mut WasiCtxBuilder, dir: &Path, at: &str| {