//! 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(crate::tier::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| {
        ww(wasi.preopened_dir(dir, at, FsPerms::ReadWrite).map(|_| ()), || format!("open {} as {at}", dir.display()))
    };
    open(&mut wasi, &spec.workspace, "/src")?;
    // Relative paths resolve against the workspace too.
    open(&mut wasi, &spec.workspace, ".")?;
    for (dir, target) in &spec.preopens {
        open(&mut wasi, dir, target)?;
    }
    let limits = Limits { memory: spec.memory_mib as usize * 1024 * 1024, peak: 0 };
    let mut store = wasmtime::Store::new(engine, State { wasi: wasi.build_p1(), limits });
    store.limiter(|s| &mut s.limits);
    // With epoch interruption on (in process), one tick from now traps:
    // the runner ticks only to stop the module.
    store.set_epoch_deadline(1);
    let instance = ww(linker.instantiate(&mut store, &module), || format!("module {path} cannot start here"))?;
    let start =
        ww(instance.get_typed_func::<(), ()>(&mut store, "_start"), || format!("module {path} has no _start: is it a WASI command?"))?;
    let code = match start.call(&mut store, ()) {
        Ok(()) => 0,
        Err(e) => match e.downcast_ref::<wasmtime_wasi::I32Exit>() {
            Some(exit) => exit.0,
            None => {
                io.note(format!("{path} trapped: {e:#}"));
                TRAPPED
            }
        },
    };
    Ok((code, store.data().limits.peak))
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A spec running `text` (WAT) as the only module, with 64 MiB.
    fn spec_of(dir: &Path, text: &str) -> Spec {
        std::fs::write(dir.join("m.wasm"), wat::parse_str(text).unwrap()).unwrap();
        Spec { workspace: dir.to_path_buf(), modules: vec!["m.wasm".into()], args: vec![], env: vec![], memory_mib: 64, preopens: vec![] }
    }

    fn env() -> Environment {
        Environment { runtime: "wasm".into(), parts: vec![], rootfs: String::new(), argv: vec![], env: vec![] }
    }

    /// Records, and whether to say the run was cancelled.
    struct Collect(Vec<Record>, bool);
    impl Sink for Collect {
        fn record(&mut self, r: Record) {
            self.0.push(r);
        }
        fn cancelled(&mut self) -> bool {
            self.1
        }
    }

    const PRINT: &str = r#"(module
      (import "wasi_snapshot_preview1" "fd_write" (func $w (param i32 i32 i32 i32) (result i32)))
      (memory (export "memory") 2)
      (data (i32.const 16) "line one\nline two")
      (func (export "_start")
        (i32.store (i32.const 0) (i32.const 16)) (i32.store (i32.const 4) (i32.const 17))
        (drop (call $w (i32.const 1) (i32.const 0) (i32.const 1) (i32.const 8)))))"#;

    const SPIN: &str = r#"(module (func (export "_start") (loop $l (br $l))))"#;

    #[test]
    fn in_process_output_streams_as_lines() {
        let dir = tempfile::tempdir().unwrap();
        let mut sink = Collect(vec![], false);
        let o = run_in_thread(&spec_of(dir.path(), PRINT), env(), Duration::from_secs(30), &mut sink).unwrap();
        assert_eq!(o.exit_code, Some(0));
        let out: Vec<&str> = sink.0.iter().filter(|r| r.stream == "out").map(|r| r.text.as_str()).collect();
        assert_eq!(out, ["line one", "line two"], "a last line without its newline still comes");
        assert_eq!(o.peak_memory_mib, Some(1), "two 64 KiB pages, in whole MiB");
    }

    #[test]
    fn in_process_stops_at_the_time_limit_and_when_cancelled() {
        let dir = tempfile::tempdir().unwrap();
        let mut sink = Collect(vec![], false);
        let started = Instant::now();
        let o = run_in_thread(&spec_of(dir.path(), SPIN), env(), Duration::from_millis(300), &mut sink).unwrap();
        assert!(o.timed_out && o.exit_code.is_none(), "{:?}", sink.0);
        assert!(started.elapsed() < Duration::from_secs(5), "the endless loop was interrupted");
        let mut sink = Collect(vec![], true);
        let o = run_in_thread(&spec_of(dir.path(), SPIN), env(), Duration::from_secs(30), &mut sink).unwrap();
        assert!(o.cancelled && o.exit_code.is_none());
        assert!(sink.0.iter().any(|r| r.text == "stopped: the run was cancelled"));
    }

    #[test]
    fn globs_match_names_with_stars() {
        assert!(wildcard("*.wasm", "oxy_core-1a2b.wasm"));
        assert!(wildcard("oxy_core-*.wasm", "oxy_core-1a2b.wasm"));
        assert!(!wildcard("oxy_core-*.wasm", "oxy_core-1a2b.d"));
        assert!(wildcard("a*b*c", "aXbYc") && !wildcard("a*b*c", "aXcYb"));
        assert!(wildcard("plain.wasm", "plain.wasm") && !wildcard("plain.wasm", "other.wasm"));
    }

    #[test]
    fn modules_are_files_of_the_job() {
        let dir = tempfile::tempdir().unwrap();
        let deps = dir.path().join("target/wasm32-wasip1/debug/deps");
        std::fs::create_dir_all(&deps).unwrap();
        for f in ["b-2.wasm", "a-1.wasm", "a-1.d"] {
            std::fs::write(deps.join(f), b"").unwrap();
        }
        let ws = dir.path();
        assert_eq!(
            modules(ws, "/src/target/wasm32-wasip1/debug/deps/*.wasm").unwrap(),
            ["target/wasm32-wasip1/debug/deps/a-1.wasm", "target/wasm32-wasip1/debug/deps/b-2.wasm"]
        );
        assert_eq!(modules(ws, "./target/wasm32-wasip1/debug/deps/b-2.wasm").unwrap(), ["target/wasm32-wasip1/debug/deps/b-2.wasm"]);
        assert!(modules(ws, "target/wasm32-wasip1/debug/deps/*.so").unwrap_err().to_string().contains("no module matches"));
        assert!(modules(ws, "../etc/passwd").is_err());
        assert!(modules(ws, "/etc/x.wasm").is_err());
        assert!(modules(ws, "target/*/debug/a.wasm").unwrap_err().to_string().contains("only the file name"));
    }
}