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

/// `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 {
    match helper(spec_path) {
        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()))
}

struct State {
    wasi: wasmtime_wasi::p1::WasiP1Ctx,
    limits: wasmtime::StoreLimits,
}

fn helper(spec_path: &Path) -> Result<i32> {
    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())?;
    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 = 0;
    for path in &spec.modules {
        let code = run_module(&engine, &linker, &spec, path)?;
        if spec.modules.len() > 1 {
            eprintln!("{NOTE}{path}: exit {code}");
        }
        if first_failure == 0 && code != 0 {
            first_failure = code;
        }
    }
    Ok(first_failure)
}

fn run_module(engine: &wasmtime::Engine, linker: &wasmtime::Linker<State>, spec: &Spec, path: &str) -> Result<i32> {
    use wasmtime_wasi::{FsPerms, WasiCtxBuilder};
    if spec.modules.len() > 1 {
        eprintln!("{NOTE}running {path}");
    }
    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();
    wasi.inherit_stdout().inherit_stderr().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 = wasmtime::StoreLimitsBuilder::new().memory_size(spec.memory_mib as usize * 1024 * 1024).build();
    let mut store = wasmtime::Store::new(engine, State { wasi: wasi.build_p1(), limits });
    store.limiter(|s| &mut s.limits);
    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?"))?;
    Ok(match start.call(&mut store, ()) {
        Ok(()) => 0,
        Err(e) => match e.downcast_ref::<wasmtime_wasi::I32Exit>() {
            Some(exit) => exit.0,
            None => {
                eprintln!("{NOTE}{path} trapped: {e:#}");
                TRAPPED
            }
        },
    })
}

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

    #[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"));
    }
}