//! What must fail does (extensions/quantum-wasmq.md §3, §4, §6): a module that could
//! reach outside itself or let a host choose its result is refused before it runs; a
//! bundle whose bytes are not the ones its contract pins is refused; a callout that runs
//! past its fuel, its call depth or its result's width ends the run in a signed fault
//! receipt, the same on every engine; and each protection is shown to be load-bearing by
//! the failure that appears when it is taken away.

use serde_json::Value;
use wai_quantum::ed25519_dalek::SigningKey;
use wai_wasmq::engine::{Engine, Interpreter, Trap};
use wai_wasmq::envelope::{Bundle, Callout, Module, Run, open, pack};
use wai_wasmq::interp::Metering;
use wai_wasmq::meter::{self, Parts};
use wai_wasmq::run::{EngineKind, run};

fn wasm(wat: &str) -> Vec<u8> {
    wat::parse_str(wat).unwrap()
}

/// A program that measures eight qubits into `b` and calls `f` on them.
const PROGRAM: &str = "OPENQASM 3.0;\ninclude \"stdgates.inc\";\nextern f(bit[8]) -> bit[8];\nqubit[8] q;\nbit[8] b;\nbit[8] r;\nh q;\nb = measure q;\nr = f(b);\n";

fn bundle(module: Vec<u8>, run: Run) -> Bundle {
    Bundle { program: PROGRAM.into(), inputs: Default::default(), modules: vec![Module { name: "m".into(), bytes: module, priors: vec![] }], callouts: vec![Callout { name: "f".into(), module: 0, export: "f".into() }], run }
}

fn refused(module: &str) -> &'static str {
    match pack(&bundle(wasm(module), Run::new(4, 1))) {
        Ok(_) => "packed",
        Err(r) => r.code,
    }
}

fn kinds() -> Vec<EngineKind> {
    #[allow(unused_mut)] // the native engine is added where this build has it
    let mut k = vec![EngineKind::Interpreter, EngineKind::InterpreterPrepared];
    #[cfg(all(feature = "wasmtime", not(target_arch = "wasm32")))]
    k.insert(0, EngineKind::Native);
    k
}

/// Runs a bundle on every engine: the job bodies are byte-identical; the first.
fn everywhere(b: &Bundle) -> Value {
    let o = open(&pack(b).unwrap()).unwrap();
    let key = SigningKey::from_bytes(&[7; 32]);
    let mut first: Option<Vec<u8>> = None;
    let mut body = Value::Null;
    for k in kinds() {
        let (v, _) = run(&o, k, &key, false).unwrap();
        let bytes = wai_quantum::engine::receipt::canonical(&v);
        match &first {
            None => first = Some(bytes),
            Some(f) => assert_eq!(f, &bytes, "{k:?} gives other bytes"),
        }
        body = v;
    }
    let verdict = wai_wasmq::verify(&body, Some(&o), kinds()[0]).unwrap();
    assert!(verdict.ok(), "{verdict:?}");
    body
}

const IDENTITY: &str = r#"(module (memory 1 1) (func (export "f") (param i64) (result i64) local.get 0))"#;

#[test]
fn the_positive_control_packs_and_runs() {
    assert_eq!(refused(IDENTITY), "packed");
    let body = everywhere(&bundle(wasm(IDENTITY), Run::new(16, 3)));
    assert_eq!(body["result"]["status"], "ok");
    // r = f(b) = b in every shot: the high byte of each key equals the low byte.
    for (k, _) in body["result"]["counts"].as_object().unwrap() {
        assert_eq!(&k[..8], &k[8..], "{k}");
    }
}

#[test]
fn a_module_that_imports_a_clock_is_refused() {
    let clock = r#"(module (import "wasi:clocks/monotonic-clock" "now" (func $now (result i64)))
                    (func (export "f") (param i64) (result i64) call $now))"#;
    assert_eq!(refused(clock), "W002");
    let random = r#"(module (import "wasi_snapshot_preview1" "random_get" (func (param i32 i32) (result i32))) (func (export "f") (param i64) (result i64) local.get 0))"#;
    assert_eq!(refused(random), "W002");
    let memory = r#"(module (import "env" "memory" (memory 1 1)) (func (export "f") (param i64) (result i64) local.get 0))"#;
    assert_eq!(refused(memory), "W002");
}

#[test]
fn a_memory_without_a_maximum_or_past_the_profiles_is_refused() {
    assert_eq!(refused(r#"(module (memory 1) (func (export "f") (param i64) (result i64) local.get 0))"#), "W003");
    assert_eq!(refused(r#"(module (memory 1 65536) (func (export "f") (param i64) (result i64) local.get 0))"#), "W003");
    assert_eq!(refused(r#"(module (memory 1 1 shared) (func (export "f") (param i64) (result i64) local.get 0))"#), "W005");
}

#[test]
fn threads_vectors_and_a_start_function_are_refused() {
    assert_eq!(refused(r#"(module (memory 1 1) (func (export "f") (param i64) (result i64) i32.const 0 local.get 0 i64.atomic.rmw.add))"#), "W005");
    assert_eq!(refused(r#"(module (func (export "f") (param i64) (result i64) local.get 0 i64x2.splat i64x2.extract_lane 0))"#), "W006");
    assert_eq!(refused(r#"(module (func $s) (start $s) (func (export "f") (param i64) (result i64) local.get 0))"#), "W004");
}

#[test]
fn growing_memory_past_the_maximum_fails_the_same_way_everywhere() {
    // Grows by the low bits of its input: past the maximum of 3 pages, memory.grow gives -1.
    let m = wasm(r#"(module (memory 1 3) (func (export "f") (param i64) (result i64)
        local.get 0 i32.wrap_i64 i32.const 7 i32.and memory.grow i64.extend_i32_s i64.const 255 i64.and))"#);
    let body = everywhere(&bundle(m, Run::new(32, 5)));
    let counts = body["result"]["counts"].as_object().unwrap();
    let mut saw_fail = false;
    for k in counts.keys() {
        let b = u8::from_str_radix(&k[8..], 2).unwrap() & 7;
        let r = u8::from_str_radix(&k[..8], 2).unwrap();
        if b <= 2 {
            assert_eq!(r, 1, "{k}: grew from one page");
        } else {
            assert_eq!(r, 255, "{k}: past the maximum");
            saw_fail = true;
        }
    }
    assert!(saw_fail);
}

/// Returns the low 8 bits of a NaN's payload made by arithmetic on a NaN with a payload.
const NAN_PATH: &str = r#"(module (memory 1 1) (func (export "f") (param i64) (result i64)
    f64.const nan:0x4000000000abc local.get 0 f64.convert_i64_u f64.add i64.reinterpret_f64 i64.const 255 i64.and))"#;

#[test]
fn a_nan_path_reads_the_canonical_nan_everywhere() {
    let body = everywhere(&bundle(wasm(NAN_PATH), Run::new(16, 2)));
    for k in body["result"]["counts"].as_object().unwrap().keys() {
        assert_eq!(&k[..8], "00000000", "the canonical NaN's payload is zero: {k}");
    }
}

#[cfg(all(feature = "wasmtime", not(target_arch = "wasm32")))]
#[test]
fn without_canonicalisation_a_nan_path_reads_what_the_host_chose() {
    // The payload propagates through the addition on this host's processor (arm64 and
    // x86_64 alike), so with neither the preparation's canonicalisation nor the engine's
    // the module returns the payload's low byte, 0xbc; with either, zero.
    let m = wasm(NAN_PATH);
    let call = |parts: Parts, engine_nan: bool| -> u64 {
        let p = meter::prepare_with(&m, parts).unwrap();
        let e = wai_wasmq::engine::Native::with(&p, engine_nan).unwrap();
        let mut i = e.instantiate().unwrap();
        i.set_fuel(1000);
        i.call("f", &[3]).unwrap()[0]
    };
    let off = Parts { nan: false, ..Parts::ALL };
    assert_eq!(call(off, false), 0xbc, "this test can fail: without canonicalisation the host's NaN shows");
    assert_eq!(call(Parts::ALL, false), 0, "the preparation alone canonicalises");
    assert_eq!(call(off, true), 0, "the engine alone canonicalises");
    assert_eq!(call(Parts::ALL, true), 0);
}

#[test]
fn a_callout_past_its_fuel_ends_in_a_signed_fault_receipt() {
    // Spins for as many iterations as its input says.
    let m = wasm(r#"(module (func (export "f") (param i64) (result i64) (local i64)
        block loop local.get 0 i64.eqz br_if 1 local.get 0 i64.const 1 i64.sub local.set 0 br 0 end end local.get 1))"#);
    let mut r = Run::new(64, 9);
    r.per_call = 400;
    let body = everywhere(&bundle(m.clone(), r.clone()));
    assert_eq!(body["status"], "fault");
    let f = &body["result"]["fault"];
    assert_eq!(f["reason"], "fuel_exhausted");
    assert_eq!(f["budget"], "call");
    assert_eq!(f["fuel"], 400);
    assert!(body["result"]["shots_run"].as_u64().unwrap() < 64);
    // The run's budget, not the call's, when it is the smaller.
    r.per_call = 1_000_000;
    r.per_job = 3000;
    let body = everywhere(&bundle(m, r));
    assert_eq!(body["result"]["fault"]["budget"], "run");
    assert!(body["result"]["fuel"]["used"].as_u64().unwrap() <= 3000);
}

#[test]
fn an_endless_loop_stops_at_its_budget() {
    let m = wasm(r#"(module (func (export "f") (param i64) (result i64) loop br 0 end i64.const 0))"#);
    let mut r = Run::new(4, 1);
    r.per_call = 10_000;
    let body = everywhere(&bundle(m, r));
    assert_eq!(body["result"]["fault"]["reason"], "fuel_exhausted");
    assert_eq!(body["result"]["fault"]["shot"], 0);
}

#[test]
fn recursion_past_the_stack_ends_in_a_fault_receipt() {
    let m = wasm(r#"(module (func $r (param i64) (result i64) local.get 0 i64.const 1 i64.add call $r)
        (func (export "f") (param i64) (result i64) local.get 0 call $r))"#);
    let body = everywhere(&bundle(m, Run::new(2, 1)));
    assert_eq!(body["result"]["fault"]["reason"], "stack_exhausted");
}

#[test]
fn a_trap_and_a_result_past_its_width_end_in_fault_receipts() {
    let trap = wasm(r#"(module (func (export "f") (param i64) (result i64) local.get 0 i64.const 0 i64.div_u))"#);
    assert_eq!(everywhere(&bundle(trap, Run::new(2, 1)))["result"]["fault"]["reason"], "integer_divide_by_zero");
    let wide = wasm(r#"(module (func (export "f") (param i64) (result i64) i64.const 256))"#);
    let body = everywhere(&bundle(wide, Run::new(2, 1)));
    assert_eq!(body["result"]["fault"]["reason"], "bad_result");
}

#[test]
fn a_callout_with_the_wrong_signature_is_refused() {
    assert_eq!(refused(r#"(module (func (export "f") (param i32) (result i32) local.get 0))"#), "W015");
    assert_eq!(refused(r#"(module (func (export "g") (param i64) (result i64) local.get 0))"#), "W015");
}

fn envelope() -> Vec<u8> {
    let mut b = bundle(wasm(IDENTITY), Run::new(4, 1));
    b.modules[0].priors.push(("table".into(), b"a prior's bytes".to_vec()));
    // A module with priors exports wasmq-prior.
    b.modules[0].bytes = wasm(r#"(module (memory 1 1) (export "memory" (memory 0))
        (func (export "cabi_realloc") (param i32 i32 i32 i32) (result i32) i32.const 1024)
        (func (export "wasmq-prior") (param i32 i32 i32))
        (func (export "f") (param i64) (result i64) local.get 0))"#);
    pack(&b).unwrap()
}

fn tamper(env: &[u8], find: &[u8], with: &[u8]) -> Vec<u8> {
    let at = env.windows(find.len()).position(|w| w == find).expect("the bytes are in the envelope");
    let mut e = env.to_vec();
    e[at..at + with.len()].copy_from_slice(with);
    e
}

#[test]
fn a_prior_whose_hash_is_not_the_pinned_one_is_refused() {
    let env = envelope();
    assert!(open(&env).is_ok(), "the positive control opens");
    let bad = tamper(&env, b"a prior's bytes", b"A prior's bytes");
    assert_eq!(open(&bad).unwrap_err().code, "W014");
}

/// The code `env` is refused with, by a refusal that says `needle`: the rule this test
/// is of, and no other.
fn refusal(env: &[u8], needle: &str) -> &'static str {
    let e = open(env).expect_err("the envelope opened");
    assert!(e.message.contains(needle), "refused for another rule than the one that says {needle:?}: {e}");
    e.code
}

#[test]
fn a_changed_module_or_program_is_refused_for_its_digest() {
    let env = envelope();
    assert_eq!(refusal(&tamper(&env, b"wasmq-prior", b"wasmq-PRIOR"), "module m: its bytes hash to"), "W014");
    assert_eq!(refusal(&tamper(&env, b"h q;", b"x q;"), "the program: its bytes hash to"), "W014");
}

#[test]
fn each_rule_of_the_envelope_s_form_refuses_what_breaks_it() {
    use wai_wasmq::envelope::{MAGIC, from_parts, parts};
    let env = envelope();
    let (manifest, contract, frames) = parts(&env).unwrap();
    let text = String::from_utf8(contract.clone()).unwrap();
    let with = |manifest: &[u8], magic: &[u8; 4], contract: &[u8], frames: &[u8]| from_parts(manifest, magic, contract, frames).unwrap();
    assert_eq!(with(&manifest, MAGIC, &contract, &frames), env, "the parts, unchanged, are the envelope");
    let edit = |from: &str, to: &str| {
        let edited = text.replacen(from, to, 1);
        assert_ne!(edited, text, "{from} is in the contract");
        edited.into_bytes()
    };
    // The contract's own run changed, canonical still: it no longer agrees with the manifest.
    assert_eq!(refusal(&with(&manifest, MAGIC, &edit("\"shots\":4", "\"shots\":5"), &frames), "manifest"), "W013");
    // A member the profile does not write, in its sorted place, the lengths made right:
    // another envelope for the same bundle (§2.4).
    assert_eq!(refusal(&with(&manifest, MAGIC, &edit(",\"modules\":", ",\"extra\":0,\"modules\":"), &frames), "members this profile does not write"), "W013");
    // The same members, not in the canonical form: a space, and the keys in another order.
    assert_eq!(refusal(&with(&manifest, MAGIC, &edit("{", "{ "), &frames), "canonical"), "W013");
    let v: serde_json::Value = serde_json::from_slice(&contract).unwrap();
    let reordered = format!("{{\"run\":{},\"callouts\":{},\"modules\":{},\"profile\":{},\"program\":{}}}", v["run"], v["callouts"], v["modules"], v["profile"], v["program"]);
    assert_eq!(refusal(&with(&manifest, MAGIC, reordered.as_bytes(), &frames), "canonical"), "W013");
    // Another payload.
    assert_eq!(refusal(&with(&manifest, b"WQM2", &contract, &frames), "not a WQM1 bundle"), "W013");
    // Frames that do not end where the payload does: a byte more, a byte less.
    let mut more = frames.clone();
    more.push(0);
    assert_eq!(refusal(&with(&manifest, MAGIC, &contract, &more), "after the frames"), "W013");
    assert_eq!(refusal(&with(&manifest, MAGIC, &contract, &frames[..frames.len() - 1]), "ends inside a frame"), "W013");
    // The manifest: a byte of it, and another capability.
    let m = String::from_utf8(manifest.clone()).unwrap();
    assert_eq!(refusal(&with(m.replace("\"shots\":4", "\"shots\":5").as_bytes(), MAGIC, &contract, &frames), "manifest"), "W013");
    assert_eq!(refusal(&with(m.replace("wai.quantum.wasmq", "wai.quantum.wasmx").as_bytes(), MAGIC, &contract, &frames), "capability"), "W013");
    // The container: bytes after the payload, a byte short, another magic.
    let mut long = env.clone();
    long.push(0);
    assert_eq!(refusal(&long, "after the payload"), "W013");
    assert_eq!(refusal(&env[..env.len() - 1], "ends inside a frame"), "W013");
    assert_eq!(refusal(&tamper(&env, b"WAI1", b"WAI2"), "not a WAI1 envelope"), "W013");
}

#[test]
fn a_receipt_changed_after_sealing_or_held_to_another_bundle_does_not_verify() {
    let b = bundle(wasm(IDENTITY), Run::new(8, 4));
    let o = open(&pack(&b).unwrap()).unwrap();
    let key = SigningKey::from_bytes(&[7; 32]);
    let (body, _) = run(&o, EngineKind::Interpreter, &key, false).unwrap();
    assert!(wai_wasmq::verify(&body, Some(&o), EngineKind::Interpreter).unwrap().ok());
    let mut forged = body.clone();
    forged["result"]["fuel"]["used"] = serde_json::json!(1);
    assert!(!wai_wasmq::verify(&forged, None, EngineKind::Interpreter).unwrap().signature);
    let mut other = b.clone();
    other.run.seed = 5;
    let o2 = open(&pack(&other).unwrap()).unwrap();
    assert_eq!(wai_wasmq::verify(&body, Some(&o2), EngineKind::Interpreter).unwrap().bundle, Some(false));
}

#[test]
fn without_its_meter_a_prepared_module_counts_nothing_and_without_its_stack_weight_it_is_not_bounded() {
    let m = wasm(r#"(module (func $r (param i64) (result i64) local.get 0 i64.eqz if (result i64) i64.const 0 else local.get 0 i64.const 1 i64.sub call $r end)
        (func (export "f") (param i64) (result i64) local.get 0 call $r))"#);
    let go = |parts: Parts, arg: u64| {
        let p = meter::prepare_with(&m, parts).unwrap();
        let e = Interpreter::new(&p, Metering::Module).unwrap();
        let mut i = e.instantiate().unwrap();
        i.set_fuel(1 << 40);
        i.set_stack(0);
        let r = i.call("f", &[arg]);
        (r, (1i64 << 40) - i.fuel())
    };
    let (r, used) = go(Parts::ALL, 100);
    assert!(r.is_ok() && used > 0);
    assert_eq!(go(Parts { fuel: false, ..Parts::ALL }, 100).1, 0, "this test can fail: no meter, no fuel");
    // Deeper than the stack's weight lets a function of this size go.
    let deep = 20_000;
    assert_eq!(go(Parts::ALL, deep).0, Err(Trap::Stack));
    assert!(go(Parts { stack: false, ..Parts::ALL }, deep).0.is_ok(), "this test can fail: no weight, no bound");
}