//! `wai-int-module-host`: the module of record on native engines, as a sink runs it.
//!
//!     wai-int-module-host [--engines cranelift,winch,wasmi] [--external DIR] [--only PREFIX]
//!                         CORPUS MODULE...
//!
//! For each module and each engine named, every case of the module corpus `CORPUS`
//! (`../corpus.json`, written by `../corpus.py`) on an instance of its own, as waism §5 and §6
//! checks 11 to 13 state: the module instantiated with no imports, `wai_abi`, the capability,
//! the payload and the set each written at an offset from `wai_alloc`, `wai_run` under a fuel
//! budget, then the output read back. Each output is held to the corpus: the SHA-256 it
//! states, or the refusal code. `--external DIR` is the directory the corpus's external files
//! were fetched to (`live-conformance/dred/fetch.sh`); a case that names one is a failure
//! without it.
//!
//! The engines, each set as waism §7 sets an engine at the integer level:
//!
//! - `cranelift`, `winch`: wasmtime 50.0.1 with that compiler; shared memory and atomics off
//!   (`threads-off`), relaxed vector instructions off (`relaxed-off`), fuel, and the memory
//!   limited to the module's declared maximum. No engine setting fixes a NaN or a relaxed
//!   instruction's behaviour: the integer level reaches neither (waism §7).
//! - `wasmi`: Wasmi 2.0.0, an interpreter, set the same way, with its floating-point
//!   instructions and types turned off as well: a module with any floating point in it does
//!   not compile, a second guard behind waism R3.
//!
//! Each run prints one line: its module, engine, the cases it held, and the **run digest**,
//! the SHA-256 of a line per case (`<name> <outcome>`, the outcome `sha256:<hex>` of the
//! output or `refuse:<code>`). The run digest of every run must be the same: a module's
//! scalar and SIMD128 builds, and every engine, give each case the same outcome.
//!
//! Exit status 0 when every case of every run is the corpus's and every run digest is the
//! same; 1 otherwise; 2 for a usage or input error.
//!
//!     wai-int-module-host --count MODULE...
//!
//! For each module, how many instructions its function bodies hold, and how many of them
//! are 128-bit vector instructions: what tells the SIMD128 build from the scalar one.

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::process::ExitCode;
use std::time::Instant;

use serde_json::Value;
use sha2::{Digest, Sha256};

/// The instructions an engine may count for one run of `wai_run` before it is stopped
/// (`module_timeout`, waism §6 check 12).
const FUEL: u64 = 1 << 42;

fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}

fn sha256(b: &[u8]) -> String {
    hex(&Sha256::digest(b))
}

// ------------------------------------------------------------------------------ the corpus

enum Expect {
    Sha256(String),
    Refuse(String),
}

struct Case {
    name: String,
    cap: String,
    payload: Vec<u8>,
    set: Vec<u8>,
    expect: Expect,
}

fn unhex(s: &str) -> Result<Vec<u8>, String> {
    if s.len() % 2 != 0 {
        return Err("odd hexadecimal".into());
    }
    (0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16).map_err(|e| e.to_string())).collect()
}

/// A file a case names: a path below wai/, `{"hex"}`, or `{"external"}`.
fn file(v: &Value, wai: &Path, external: Option<&Path>) -> Result<Vec<u8>, String> {
    if let Some(p) = v.as_str() {
        return std::fs::read(wai.join(p)).map_err(|e| format!("{p}: {e}"));
    }
    if let Some(h) = v.get("hex").and_then(Value::as_str) {
        return unhex(h);
    }
    if let Some(p) = v.get("file").and_then(Value::as_str) {
        return std::fs::read(wai.join(p)).map_err(|e| format!("{p}: {e}"));
    }
    if let Some(n) = v.get("external").and_then(Value::as_str) {
        let dir = external.ok_or_else(|| format!("{n} is an external file, and no --external directory was given"))?;
        return std::fs::read(dir.join(n)).map_err(|e| format!("{n}: {e}"));
    }
    Err(format!("not a file: {v}"))
}

/// The set a case names, as waism §5 passes it: nothing, one file's bytes, or the file
/// sequence of a list of `[name, file]`.
fn set(v: &Value, wai: &Path, external: Option<&Path>) -> Result<Vec<u8>, String> {
    match v {
        Value::Null => Ok(Vec::new()),
        Value::Array(files) => {
            let mut out = Vec::new();
            for f in files {
                let (Some(name), Some(what)) = (f.get(0).and_then(Value::as_str), f.get(1)) else { return Err(format!("not [name, file]: {f}")) };
                let b = file(what, wai, external)?;
                out.push(u8::try_from(name.len()).map_err(|_| "a name past 255 bytes")?);
                out.extend_from_slice(name.as_bytes());
                out.extend_from_slice(&u32::try_from(b.len()).map_err(|_| "a file past 4 GiB")?.to_le_bytes());
                out.extend_from_slice(&b);
            }
            Ok(out)
        }
        other => file(other, wai, external),
    }
}

fn corpus(path: &Path, external: Option<&Path>, only: Option<&str>) -> Result<Vec<Case>, String> {
    let text = std::fs::read(path).map_err(|e| format!("{}: {e}", path.display()))?;
    let doc: Value = serde_json::from_slice(&text).map_err(|e| format!("{}: {e}", path.display()))?;
    if doc["format"] != "waism-module-corpus/1" {
        return Err(format!("{}: not a waism-module-corpus/1", path.display()));
    }
    let abs = std::fs::canonicalize(path).map_err(|e| e.to_string())?;
    let wai = abs.parent().and_then(Path::parent).ok_or("the corpus is not two levels below wai/")?;
    let mut cases = Vec::new();
    for c in doc["cases"].as_array().ok_or("no cases")? {
        let name = c["name"].as_str().ok_or("a case with no name")?.to_string();
        if only.is_some_and(|o| !name.starts_with(o)) {
            continue;
        }
        let expect = if let Some(s) = c["expect"]["sha256"].as_str() {
            Expect::Sha256(s.to_string())
        } else if let Some(r) = c["expect"]["refuse"].as_str() {
            Expect::Refuse(r.to_string())
        } else {
            return Err(format!("{name}: no expect"));
        };
        let cx = |e: String| format!("{name}: {e}");
        cases.push(Case {
            cap: c["capability"].as_str().ok_or("a case with no capability")?.to_string(),
            payload: file(&c["payload"], wai, external).map_err(cx)?,
            set: set(&c["set"], wai, external).map_err(cx)?,
            expect,
            name,
        });
    }
    Ok(cases)
}

// --------------------------------------------------------------------------- the sink

/// What a run of one case gave: the module's output or its refusal, or the sink's refusal
/// of the run (waism §8: `module_resources`, `module_trap`, `module_timeout`,
/// `module_result_invalid`).
enum Outcome {
    Output(Vec<u8>),
    Refused(String),
    Sink(&'static str, String),
}

impl Outcome {
    fn line(&self) -> String {
        match self {
            Outcome::Output(b) => format!("sha256:{}", sha256(b)),
            Outcome::Refused(c) => format!("refuse:{c}"),
            Outcome::Sink(c, _) => format!("sink:{c}"),
        }
    }
}

/// A refusal code a module returns: ASCII lowercase letters, digits and `_`.
fn code(b: &[u8]) -> Option<String> {
    (!b.is_empty() && b.iter().all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || *c == b'_')).then(|| String::from_utf8_lossy(b).into_owned())
}

/// The calls of waism §5 on one instance, for any engine: `call` runs an export, `write`
/// and `read` reach the memory.
trait Instance {
    fn call0(&mut self, name: &str) -> Result<i32, Outcome>;
    fn call1(&mut self, name: &str, a: i32) -> Result<i32, Outcome>;
    fn run(&mut self, args: [i32; 6]) -> Result<i32, Outcome>;
    fn write(&mut self, at: usize, b: &[u8]) -> bool;
    fn read(&mut self, at: usize, n: usize) -> Option<Vec<u8>>;
}

fn sink(inst: &mut dyn Instance, c: &Case) -> Outcome {
    match inst.call0("wai_abi") {
        Ok(1) => {}
        Ok(v) => return Outcome::Sink("module_result_invalid", format!("wai_abi gave {v}")),
        Err(o) => return o,
    }
    let mut args = [0i32; 6];
    for (i, input) in [c.cap.as_bytes(), &c.payload, &c.set].into_iter().enumerate() {
        if input.is_empty() {
            continue;
        }
        let Ok(len) = i32::try_from(input.len()) else { return Outcome::Sink("module_resources", "an input past 2 GiB".into()) };
        let at = match inst.call1("wai_alloc", len) {
            Ok(0) => return Outcome::Sink("module_resources", "wai_alloc gave 0".into()),
            Ok(at) => at,
            Err(o) => return o,
        };
        if !inst.write(at as u32 as usize, input) {
            return Outcome::Sink("module_result_invalid", "wai_alloc gave a region outside the memory".into());
        }
        args[2 * i] = at;
        args[2 * i + 1] = len;
    }
    let status = match inst.run(args) {
        Ok(s) => s,
        Err(o) => return o,
    };
    let (ptr, len) = match (inst.call0("wai_out_ptr"), inst.call0("wai_out_len")) {
        (Ok(p), Ok(l)) => (p as u32 as usize, l as u32 as usize),
        (Err(o), _) | (_, Err(o)) => return o,
    };
    let Some(out) = inst.read(ptr, len) else { return Outcome::Sink("module_result_invalid", "the output lies outside the memory".into()) };
    match status {
        0 => Outcome::Output(out),
        1 => match code(&out) {
            Some(c) => Outcome::Refused(c),
            None => Outcome::Sink("module_result_invalid", "a refusal that is not a code".into()),
        },
        s => Outcome::Sink("module_result_invalid", format!("wai_run gave {s}")),
    }
}

// --------------------------------------------------------------------------- wasmtime

mod wt {
    use super::{FUEL, Instance as Inst, Outcome};
    use wasmtime::{Config, Engine, Instance, Memory, Module, Store, StoreLimits, StoreLimitsBuilder, Strategy, Trap};

    pub struct Data {
        limits: StoreLimits,
    }

    pub fn engine(strategy: Strategy) -> wasmtime::Result<Engine> {
        let mut c = Config::new();
        c.strategy(strategy);
        c.wasm_threads(false);
        c.wasm_relaxed_simd(false);
        c.consume_fuel(true);
        Engine::new(&c)
    }

    pub struct Run {
        store: Store<Data>,
        inst: Instance,
        mem: Memory,
    }

    fn fail(e: wasmtime::Error) -> Outcome {
        match e.downcast_ref::<Trap>() {
            Some(Trap::OutOfFuel) => Outcome::Sink("module_timeout", e.to_string()),
            _ => Outcome::Sink("module_trap", format!("{e:#}")),
        }
    }

    pub fn instantiate(engine: &Engine, module: &Module, max_memory: usize) -> Result<Run, Outcome> {
        let mut store = Store::new(engine, Data { limits: StoreLimitsBuilder::new().memory_size(max_memory).build() });
        store.limiter(|d| &mut d.limits);
        store.set_fuel(FUEL).map_err(|e| Outcome::Sink("module_resources", e.to_string()))?;
        let inst = Instance::new(&mut store, module, &[]).map_err(|e| Outcome::Sink("module_resources", format!("{e:#}")))?;
        let mem = inst.get_memory(&mut store, "memory").ok_or_else(|| Outcome::Sink("module_result_invalid", "no memory".into()))?;
        Ok(Run { store, inst, mem })
    }

    impl Inst for Run {
        fn call0(&mut self, name: &str) -> Result<i32, Outcome> {
            let f = self.inst.get_typed_func::<(), i32>(&mut self.store, name).map_err(|e| Outcome::Sink("module_result_invalid", e.to_string()))?;
            f.call(&mut self.store, ()).map_err(fail)
        }
        fn call1(&mut self, name: &str, a: i32) -> Result<i32, Outcome> {
            let f = self.inst.get_typed_func::<i32, i32>(&mut self.store, name).map_err(|e| Outcome::Sink("module_result_invalid", e.to_string()))?;
            f.call(&mut self.store, a).map_err(fail)
        }
        fn run(&mut self, a: [i32; 6]) -> Result<i32, Outcome> {
            let f = self
                .inst
                .get_typed_func::<(i32, i32, i32, i32, i32, i32), i32>(&mut self.store, "wai_run")
                .map_err(|e| Outcome::Sink("module_result_invalid", e.to_string()))?;
            f.call(&mut self.store, (a[0], a[1], a[2], a[3], a[4], a[5])).map_err(fail)
        }
        fn write(&mut self, at: usize, b: &[u8]) -> bool {
            self.mem.write(&mut self.store, at, b).is_ok()
        }
        fn read(&mut self, at: usize, n: usize) -> Option<Vec<u8>> {
            let mut out = vec![0u8; n];
            self.mem.read(&self.store, at, &mut out).ok().map(|_| out)
        }
    }
}

// ------------------------------------------------------------------------------ wasmi

mod wi {
    use super::{FUEL, Instance as Inst, Outcome};
    use wasmi::{Config, Engine, Instance, Memory, Module, Store, StoreLimits, StoreLimitsBuilder, TrapCode};

    pub struct Data {
        limits: StoreLimits,
    }

    pub fn engine() -> Engine {
        let mut c = Config::default();
        c.wasm_relaxed_simd(false);
        c.floats(false);
        c.consume_fuel(true);
        Engine::new(&c)
    }

    pub struct Run {
        store: Store<Data>,
        inst: Instance,
        mem: Memory,
    }

    fn fail(e: wasmi::Error) -> Outcome {
        match e.as_trap_code() {
            Some(TrapCode::OutOfFuel) => Outcome::Sink("module_timeout", e.to_string()),
            _ => Outcome::Sink("module_trap", e.to_string()),
        }
    }

    pub fn instantiate(engine: &Engine, module: &Module, max_memory: usize) -> Result<Run, Outcome> {
        let mut store = Store::new(engine, Data { limits: StoreLimitsBuilder::new().memory_size(max_memory).build() });
        store.limiter(|d| &mut d.limits);
        store.set_fuel(FUEL).map_err(|e| Outcome::Sink("module_resources", e.to_string()))?;
        let inst = Instance::new(&mut store, module, &[]).map_err(|e| Outcome::Sink("module_resources", e.to_string()))?;
        let mem = inst.get_memory(&store, "memory").ok_or_else(|| Outcome::Sink("module_result_invalid", "no memory".into()))?;
        Ok(Run { store, inst, mem })
    }

    impl Inst for Run {
        fn call0(&mut self, name: &str) -> Result<i32, Outcome> {
            let f = self.inst.get_typed_func::<(), i32>(&self.store, name).map_err(|e| Outcome::Sink("module_result_invalid", e.to_string()))?;
            f.call(&mut self.store, ()).map_err(fail)
        }
        fn call1(&mut self, name: &str, a: i32) -> Result<i32, Outcome> {
            let f = self.inst.get_typed_func::<i32, i32>(&self.store, name).map_err(|e| Outcome::Sink("module_result_invalid", e.to_string()))?;
            f.call(&mut self.store, a).map_err(fail)
        }
        fn run(&mut self, a: [i32; 6]) -> Result<i32, Outcome> {
            let f = self
                .inst
                .get_typed_func::<(i32, i32, i32, i32, i32, i32), i32>(&self.store, "wai_run")
                .map_err(|e| Outcome::Sink("module_result_invalid", e.to_string()))?;
            f.call(&mut self.store, (a[0], a[1], a[2], a[3], a[4], a[5])).map_err(fail)
        }
        fn write(&mut self, at: usize, b: &[u8]) -> bool {
            self.mem.write(&mut self.store, at, b).is_ok()
        }
        fn read(&mut self, at: usize, n: usize) -> Option<Vec<u8>> {
            let mut out = vec![0u8; n];
            self.mem.read(&self.store, at, &mut out).ok().map(|_| out)
        }
    }
}

// ------------------------------------------------------------------------------- runs

/// The module's declared maximum memory, in bytes (waism R6), read from its memory section.
fn max_memory(bytes: &[u8]) -> Result<usize, String> {
    use wasmparser::{Parser, Payload};
    for p in Parser::new(0).parse_all(bytes) {
        if let Payload::MemorySection(r) = p.map_err(|e| e.to_string())? {
            let m = r.into_iter().next().ok_or("no memory")?.map_err(|e| e.to_string())?;
            let pages = m.maximum.ok_or("the memory declares no maximum")?;
            return usize::try_from(pages * 65_536).map_err(|e| e.to_string());
        }
    }
    Err("no memory".into())
}

/// One run: every case on one engine, each on an instance of its own. Returns each case's
/// outcome line, and how many cases differ from the corpus.
fn run_all(cases: &[Case], mut make: impl FnMut() -> Result<Box<dyn Instance>, Outcome>) -> (Vec<String>, usize) {
    let (mut lines, mut wrong) = (Vec::new(), 0);
    for c in cases {
        let got = match make() {
            Ok(mut inst) => sink(inst.as_mut(), c),
            Err(o) => o,
        };
        let ok = match (&c.expect, &got) {
            (Expect::Sha256(want), Outcome::Output(b)) => sha256(b) == *want,
            (Expect::Refuse(want), Outcome::Refused(code)) => code == want,
            _ => false,
        };
        if !ok {
            wrong += 1;
            let want = match &c.expect {
                Expect::Sha256(s) => format!("sha256:{s}"),
                Expect::Refuse(r) => format!("refuse:{r}"),
            };
            let detail = match &got {
                Outcome::Sink(_, d) => format!(" ({d})"),
                _ => String::new(),
            };
            println!("  {}: {}{detail}, not {want}", c.name, got.line());
        }
        lines.push(format!("{} {}\n", c.name, got.line()));
    }
    (lines, wrong)
}

/// The instructions of a module's function bodies, and how many are 128-bit vector
/// instructions (those of WebAssembly's `0xfd` prefix).
fn count(bytes: &[u8]) -> Result<(u64, u64), String> {
    use wasmparser::{Parser, Payload};
    let (mut all, mut vector) = (0u64, 0u64);
    for p in Parser::new(0).parse_all(bytes) {
        if let Payload::CodeSectionEntry(body) = p.map_err(|e| e.to_string())? {
            let mut ops = body.get_operators_reader().map_err(|e| e.to_string())?;
            while !ops.eof() {
                let at = usize::try_from(ops.original_position()).map_err(|e| e.to_string())?;
                ops.read().map_err(|e| e.to_string())?;
                all += 1;
                if bytes.get(at) == Some(&0xfd) {
                    vector += 1;
                }
            }
        }
    }
    Ok((all, vector))
}

fn main() -> ExitCode {
    let args: Vec<String> = std::env::args().skip(1).collect();
    if args.first().map(String::as_str) == Some("--count") && args.len() > 1 {
        for path in &args[1..] {
            match std::fs::read(path).map_err(|e| e.to_string()).and_then(|b| count(&b).map(|c| (sha256(&b), c))) {
                Ok((sum, (all, vector))) => println!("{path} (SHA-256 {sum}): {all} instructions, {vector} of them 128-bit vector instructions"),
                Err(e) => {
                    eprintln!("wai-int-module-host: {path}: {e}");
                    return ExitCode::from(2);
                }
            }
        }
        return ExitCode::SUCCESS;
    }
    let usage = || {
        eprintln!("usage: wai-int-module-host [--engines cranelift,winch,wasmi] [--external DIR] [--only PREFIX] CORPUS MODULE...");
        ExitCode::from(2)
    };
    let (mut engines, mut external, mut only, mut rest) = ("cranelift,winch,wasmi".to_string(), None::<PathBuf>, None::<String>, Vec::new());
    let mut it = args.into_iter();
    while let Some(a) = it.next() {
        match a.as_str() {
            "--engines" => match it.next() {
                Some(e) => engines = e,
                None => return usage(),
            },
            "--external" => match it.next() {
                Some(d) => external = Some(PathBuf::from(d)),
                None => return usage(),
            },
            "--only" => match it.next() {
                Some(p) => only = Some(p),
                None => return usage(),
            },
            s if s.starts_with("--") => return usage(),
            _ => rest.push(a),
        }
    }
    if rest.len() < 2 {
        return usage();
    }
    let cases = match corpus(Path::new(&rest[0]), external.as_deref(), only.as_deref()) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("wai-int-module-host: {e}");
            return ExitCode::from(2);
        }
    };
    let mut digests: BTreeMap<String, String> = BTreeMap::new();
    let mut failed = false;
    for path in &rest[1..] {
        let bytes = match std::fs::read(path) {
            Ok(b) => b,
            Err(e) => {
                eprintln!("wai-int-module-host: {path}: {e}");
                return ExitCode::from(2);
            }
        };
        let module_sha = sha256(&bytes);
        for name in engines.split(',') {
            let t0 = Instant::now();
            let (lines, wrong) = match name {
                "cranelift" | "winch" => {
                    let strategy = if name == "winch" { wasmtime::Strategy::Winch } else { wasmtime::Strategy::Cranelift };
                    let engine = match wt::engine(strategy) {
                        Ok(e) => e,
                        Err(e) => {
                            println!("{path} on wasmtime 50.0.1/{name}: the engine does not start: {e:#}");
                            failed = true;
                            continue;
                        }
                    };
                    let max = match max_memory(&bytes) {
                        Ok(m) => m,
                        Err(e) => {
                            println!("{path} on wasmtime 50.0.1/{name}: {e}");
                            failed = true;
                            continue;
                        }
                    };
                    let module = match wasmtime::Module::new(&engine, &bytes) {
                        Ok(m) => m,
                        Err(e) => {
                            println!("{path} on wasmtime 50.0.1/{name}: the module does not compile: {e:#}");
                            failed = true;
                            continue;
                        }
                    };
                    run_all(&cases, || wt::instantiate(&engine, &module, max).map(|r| Box::new(r) as Box<dyn Instance>))
                }
                "wasmi" => {
                    let max = match max_memory(&bytes) {
                        Ok(m) => m,
                        Err(e) => {
                            println!("{path} on wasmi 2.0.0: {e}");
                            failed = true;
                            continue;
                        }
                    };
                    let engine = wi::engine();
                    let module = match wasmi::Module::new(&engine, &bytes) {
                        Ok(m) => m,
                        Err(e) => {
                            println!("{path} on wasmi 2.0.0: the module does not compile: {e}");
                            failed = true;
                            continue;
                        }
                    };
                    run_all(&cases, || wi::instantiate(&engine, &module, max).map(|r| Box::new(r) as Box<dyn Instance>))
                }
                other => {
                    eprintln!("wai-int-module-host: no engine {other}");
                    return ExitCode::from(2);
                }
            };
            let run = sha256(lines.concat().as_bytes());
            let label = match name {
                "wasmi" => "wasmi@2.0.0".to_string(),
                n => format!("wasmtime@50.0.1/{n}"),
            };
            println!(
                "{path} (SHA-256 {module_sha}) on {label}, {}: {} of {} cases the corpus's; run digest {run}; {:.1} s",
                std::env::consts::ARCH,
                cases.len() - wrong,
                cases.len(),
                t0.elapsed().as_secs_f64()
            );
            failed |= wrong > 0;
            digests.insert(format!("{path} {label}"), run);
        }
    }
    let distinct: std::collections::BTreeSet<&String> = digests.values().collect();
    if distinct.len() > 1 {
        println!("wai-int-module-host: the runs differ: {} run digests", distinct.len());
        failed = true;
    } else if let Some(d) = distinct.iter().next() {
        println!("wai-int-module-host: {} runs, one run digest {d}", digests.len());
    }
    if failed { ExitCode::FAILURE } else { ExitCode::SUCCESS }
}