//! ctx.wasm.run (javascript/src/plugin-context wasmAccess + wasm-worker.js) on wasmtime: the
//! same limits and WASM_* error codes, a fresh instance per call, the sha256 checked, compiled
//! modules cached by digest, and a wall-clock limit per call through epoch interruption.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use serde_json::Value;
use sha2::{Digest, Sha256};
use wasmtime::{
    Config, Engine, ExternType, Instance, Module, Store, StoreLimits, StoreLimitsBuilder, Trap,
};

pub const MODULE_BYTES: usize = 4 << 20;
/// The default input and output cap per call; a host raises it per run (budget ioBytes) up to IO_BYTES_MAX.
pub const IO_BYTES: usize = 1 << 20;
pub const IO_BYTES_MAX: usize = 16 << 20;
pub const MEMORY_PAGES: usize = 1024;
pub const WALL_TIME_MS: u64 = 5000;
const TICK_MS: u64 = 5;

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

pub fn entry_ok(entry: &str) -> bool {
    let mut chars = entry.chars();
    matches!(chars.next(), Some(c) if c.is_ascii_alphabetic() || c == '_')
        && entry.len() <= 256
        && chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
}

/// ATLIBS_WASM_CONCURRENCY: how many runs execute at once in a process (default: its cores). Hosts
/// check it at startup (atlibs: before any command), so a bad value is an error there, not mid-run.
pub fn concurrency() -> Result<usize, String> {
    parse_concurrency(std::env::var("ATLIBS_WASM_CONCURRENCY").ok())
}

fn parse_concurrency(value: Option<String>) -> Result<usize, String> {
    match value.filter(|v| !v.trim().is_empty()) {
        None => Ok(std::thread::available_parallelism().map_or(1, |n| n.get())),
        Some(v) => v.trim().parse().ok().filter(|n| *n > 0).ok_or_else(|| {
            format!("ATLIBS_WASM_CONCURRENCY must be a positive integer, not {v:?}")
        }),
    }
}

/// One per process (atlibs's run::Host, the agents service, a bridge run): its semaphore bounds the
/// runs executing at once, so overlapping requests cannot multiply per-instance memory without limit.
pub struct Wasm {
    engine: Engine,
    cache: Mutex<HashMap<String, Module>>,
    timeout_ms: u64,
    slots: Arc<tokio::sync::Semaphore>,
}

impl Wasm {
    /// Panics on a bad ATLIBS_WASM_CONCURRENCY: check `concurrency()` first to report it.
    pub fn new(timeout_ms: u64) -> Arc<Self> {
        Self::with_concurrency(timeout_ms, concurrency().unwrap_or_else(|e| panic!("{e}")))
    }

    pub fn with_concurrency(timeout_ms: u64, slots: usize) -> Arc<Self> {
        let mut config = Config::new();
        config.epoch_interruption(true);
        // The memory cap (StoreLimits.memory_size) is per memory: one 32-bit memory per module, as
        // the JS host's capMemory refuses a second memory or memory64.
        config.wasm_multi_memory(false);
        config.wasm_memory64(false);
        // Compiled code cached on disk (wasmtime's default cache dir), so a CLI run skips Cranelift.
        config.cache(wasmtime::Cache::from_file(None).ok());
        let engine = Engine::new(&config).expect("wasmtime engine");
        let ticker = engine.weak();
        std::thread::spawn(move || {
            while let Some(engine) = ticker.upgrade() {
                engine.increment_epoch();
                drop(engine);
                std::thread::sleep(Duration::from_millis(TICK_MS));
            }
        });
        Arc::new(Self {
            engine,
            cache: Mutex::default(),
            timeout_ms,
            slots: Arc::new(tokio::sync::Semaphore::new(
                slots.clamp(1, tokio::sync::Semaphore::MAX_PERMITS),
            )),
        })
    }

    /// Run `entry` with `input` as JSON; errors are "WASM_<CODE>: message", as the JS host throws.
    pub async fn run(
        self: &Arc<Self>,
        module: Arc<Vec<u8>>,
        sha256: Option<String>,
        entry: String,
        input: &Value,
        io_bytes: usize,
    ) -> Result<Value, String> {
        let io = if io_bytes == 0 { IO_BYTES } else { io_bytes.min(IO_BYTES_MAX) };
        if module.len() > MODULE_BYTES {
            return Err(format!(
                "WASM_TOO_LARGE: module exceeds {MODULE_BYTES} bytes"
            ));
        }
        if !entry_ok(&entry) {
            return Err("WASM_ABI_INVALID: entry must be an export name".into());
        }
        let input = serde_json::to_vec(input).map_err(|e| format!("WASM_ABI_INVALID: {e}"))?;
        if input.len() > io {
            return Err(format!("WASM_TOO_LARGE: input exceeds {io} bytes"));
        }
        // Queued here (FIFO) rather than on the blocking pool; the wall-time limit starts in call(),
        // so time spent waiting for a slot is not counted against it.
        let permit = self.slots.clone().acquire_owned().await.map_err(|e| format!("WASM_UNAVAILABLE: {e}"))?;
        let this = self.clone();
        tokio::task::spawn_blocking(move || {
            // held until the blocking task ends, even when the caller stops waiting (its memory is live till then)
            let _permit = permit;
            let digest = sha256_hex(&module);
            if sha256.is_some_and(|s| s.to_lowercase() != digest) {
                return Err("WASM_DIGEST_MISMATCH: module digest does not match sha256".into());
            }
            let compiled = this.compile(&digest, &module)?;
            this.call(&compiled, &entry, &input, io)
        })
        .await
        .map_err(|e| format!("WASM_UNAVAILABLE: {e}"))?
    }

    /// The entries `module` exports (its functions but alloc), so a host can answer "no entry X (a, b)"
    /// before a run. Compiles it (cached, as run does).
    pub fn entries(&self, module: &[u8]) -> Result<Vec<String>, String> {
        let compiled = self.compile(&sha256_hex(module), module)?;
        Ok(compiled.exports().filter(|e| matches!(e.ty(), ExternType::Func(_)) && e.name() != "alloc").map(|e| e.name().to_string()).collect())
    }

    fn compile(&self, digest: &str, bytes: &[u8]) -> Result<Module, String> {
        let cached = self.cache.lock().unwrap().get(digest).cloned();
        let module = match cached {
            Some(module) => module,
            None => {
                let module = Module::new(&self.engine, bytes)
                    .map_err(|e| format!("WASM_ABI_INVALID: {e}"))?;
                self.cache
                    .lock()
                    .unwrap()
                    .insert(digest.into(), module.clone());
                module
            }
        };
        if module.imports().len() > 0 {
            return Err("WASM_IMPORTS_NOT_ALLOWED: modules must not declare imports".into());
        }
        // A minimum over the cap is refused here, as the JS host's capMemory does.
        if module.exports().any(
            |e| matches!(e.ty(), ExternType::Memory(m) if m.minimum() > MEMORY_PAGES as u64),
        ) {
            return Err(format!(
                "WASM_TOO_LARGE: memory exceeds {MEMORY_PAGES} pages"
            ));
        }
        Ok(module)
    }

    fn call(&self, module: &Module, entry: &str, input: &[u8], io: usize) -> Result<Value, String> {
        // The memory cap holds inside the call: memory.grow past it returns -1 (the JS host caps
        // the module's declared maximum, capMemory, to the same effect).
        let limits: StoreLimits = StoreLimitsBuilder::new()
            .memory_size(MEMORY_PAGES * 65536)
            .memories(1)
            .instances(1)
            .build();
        let mut store = Store::new(&self.engine, limits);
        store.limiter(|limits| limits);
        store.set_epoch_deadline(self.timeout_ms.div_ceil(TICK_MS).max(1));
        store.epoch_deadline_trap();
        let timeout = |e: wasmtime::Error| -> String {
            match e.downcast_ref::<Trap>() {
                Some(Trap::Interrupt) => format!(
                    "WASM_TIMEOUT: {entry} ran longer than {} ms",
                    self.timeout_ms
                ),
                Some(trap) => format!("WASM_TRAP: {}", v8_message(*trap)),
                None => format!("WASM_TRAP: {e}"),
            }
        };
        let instance = Instance::new(&mut store, module, &[]).map_err(timeout)?;
        let (Some(memory), Some(alloc), Some(func)) = (
            instance.get_memory(&mut store, "memory"),
            instance.get_func(&mut store, "alloc"),
            instance.get_func(&mut store, entry),
        ) else {
            return Err(format!(
                "WASM_ABI_INVALID: module must export memory, alloc and {entry}"
            ));
        };
        let alloc = alloc
            .typed::<i32, i32>(&store)
            .map_err(|e| format!("WASM_TRAP: {e}"))?;
        let func = func
            .typed::<(i32, i32), i64>(&store)
            .map_err(|_| format!("WASM_ABI_INVALID: {entry} must return i64"))?;
        let pointer = alloc
            .call(&mut store, input.len() as i32)
            .map_err(timeout)? as u32 as usize;
        memory
            .write(&mut store, pointer, input)
            .map_err(|e| format!("WASM_TRAP: {e}"))?;
        let packed = func
            .call(&mut store, (pointer as i32, input.len() as i32))
            .map_err(timeout)? as u64;
        let size = memory.data_size(&store);
        if size > MEMORY_PAGES * 65536 {
            return Err(format!(
                "WASM_TOO_LARGE: memory exceeds {MEMORY_PAGES} pages"
            ));
        }
        let (out_ptr, out_len) = ((packed >> 32) as usize, (packed & 0xffff_ffff) as usize);
        if out_len > io {
            return Err(format!("WASM_OUTPUT_INVALID: output exceeds {io} bytes"));
        }
        if out_ptr + out_len > size {
            return Err("WASM_OUTPUT_INVALID: output points outside memory".into());
        }
        let bytes = &memory.data(&store)[out_ptr..out_ptr + out_len];
        std::str::from_utf8(bytes)
            .ok()
            .and_then(|text| serde_json::from_str(text).ok())
            .ok_or_else(|| "WASM_OUTPUT_INVALID: output is not JSON".into())
    }
}

/// ctx.wasm.run bound to one module's bytes and expected sha256.
pub struct Bound {
    pub wasm: Arc<Wasm>,
    pub bytes: Arc<Vec<u8>>,
    pub sha256: Option<String>,
}

impl crate::run::Runner for Bound {
    fn run(&self, entry: String, input: Value, io_bytes: usize) -> crate::BoxFut<'_, Result<Value, String>> {
        Box::pin(async move {
            self.wasm
                .run(self.bytes.clone(), self.sha256.clone(), entry, &input, io_bytes)
                .await
        })
    }
}

/// V8's wording for wasm traps, so WASM_TRAP messages match the JS host.
fn v8_message(trap: Trap) -> String {
    match trap {
        Trap::UnreachableCodeReached => "unreachable",
        Trap::MemoryOutOfBounds | Trap::HeapMisaligned => "memory access out of bounds",
        Trap::IntegerDivisionByZero => "divide by zero",
        Trap::IntegerOverflow => "divide result unrepresentable",
        Trap::BadConversionToInteger => "float unrepresentable in integer range",
        Trap::TableOutOfBounds => "table index is out of bounds",
        Trap::IndirectCallToNull | Trap::BadSignature => {
            "null function or function signature mismatch"
        }
        Trap::StackOverflow => "Maximum call stack size exceeded",
        other => return other.to_string(),
    }
    .into()
}

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

    // (module (memory (export "memory") 1) (func (export "alloc") (param i32) (result i32) i32.const 16)
    //   (func (export "spin") (param i32 i32) (result i64) (loop br 0) i64.const 0)
    //   (func (export "boom") (param i32 i32) (result i64) unreachable)
    //   (func (export "echo") (param i32 i32) (result i64) i64.const 16 i64.const 32 i64.shl local.get 1 i64.extend_i32_u i64.or))
    // Hand-assembled binary of the module above (no wat dependency).
    fn toy() -> Vec<u8> {
        let mut m = vec![0, 0x61, 0x73, 0x6d, 1, 0, 0, 0];
        let section = |m: &mut Vec<u8>, id: u8, body: Vec<u8>| {
            m.push(id);
            m.push(body.len() as u8);
            m.extend(body);
        };
        section(
            &mut m,
            1,
            vec![2, 0x60, 1, 0x7f, 1, 0x7f, 0x60, 2, 0x7f, 0x7f, 1, 0x7e],
        );
        section(&mut m, 3, vec![4, 0, 1, 1, 1]);
        section(&mut m, 5, vec![1, 0, 1]);
        let mut exports = vec![5];
        for (name, kind, index) in [
            ("memory", 2u8, 0u8),
            ("alloc", 0, 0),
            ("spin", 0, 1),
            ("boom", 0, 2),
            ("echo", 0, 3),
        ] {
            exports.push(name.len() as u8);
            exports.extend(name.as_bytes());
            exports.extend([kind, index]);
        }
        section(&mut m, 7, exports);
        let bodies: [Vec<u8>; 4] = [
            vec![0, 0x41, 16, 0x0b],
            vec![0, 0x03, 0x40, 0x0c, 0, 0x0b, 0x42, 0, 0x0b],
            vec![0, 0x00, 0x0b],
            vec![0, 0x42, 16, 0x42, 32, 0x86, 0x20, 1, 0xad, 0x84, 0x0b],
        ];
        let mut code = vec![4];
        for body in bodies {
            code.push(body.len() as u8);
            code.extend(body);
        }
        section(&mut m, 10, code);
        m
    }

    // exports memory (min pages, no maximum, as rustc emits), alloc, and grow: memory.grow(2000),
    // then {"grew":false} when that returned -1, else {"grew":true}. The same bytes are in
    // javascript/src/plugin-context/wasm.test.mjs.
    fn grow(min: u16) -> Vec<u8> {
        let memory: &[u8] = if min == 1 {
            &[0x05, 0x03, 0x01, 0x00, 0x01]
        } else {
            &[0x05, 0x04, 0x01, 0x00, 0xd0, 0x0f]
        };
        [
            &[
                0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00, 0x01, 0x0c, 0x02, 0x60, 0x01, 0x7f,
                0x01, 0x7f, 0x60, 0x02, 0x7f, 0x7f, 0x01, 0x7e, 0x03, 0x03, 0x02, 0x00, 0x01,
            ][..],
            memory,
            &[
                0x07, 0x19, 0x03, 0x06, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x02, 0x00, 0x05, 0x61,
                0x6c, 0x6c, 0x6f, 0x63, 0x00, 0x00, 0x04, 0x67, 0x72, 0x6f, 0x77, 0x00, 0x01, 0x0a,
                0x1f, 0x02, 0x05, 0x00, 0x41, 0xc0, 0x00, 0x0b, 0x17, 0x00, 0x41, 0xd0, 0x0f, 0x40,
                0x00, 0x41, 0x7f, 0x46, 0x04, 0x7e, 0x42, 0x0e, 0x05, 0x42, 0x8d, 0x80, 0x80, 0x80,
                0x80, 0x02, 0x0b, 0x0b, 0x0b, 0x23, 0x01, 0x00, 0x41, 0x00, 0x0b, 0x1d,
            ],
            br#"{"grew":false}"#,
            &[0, 0],
            br#"{"grew":true}"#,
        ]
        .concat()
    }

    #[tokio::test]
    async fn memory_cap_holds_inside_the_call() {
        let wasm = Wasm::new(1000);
        let input = serde_json::json!({});
        assert_eq!(
            wasm.run(Arc::new(grow(1)), None, "grow".into(), &input, 0)
                .await
                .unwrap(),
            serde_json::json!({ "grew": false })
        );
        assert_eq!(
            wasm.run(Arc::new(grow(2000)), None, "grow".into(), &input, 0)
                .await
                .unwrap_err(),
            "WASM_TOO_LARGE: memory exceeds 1024 pages"
        );
        // A second memory (or memory64) would escape the per-memory cap: refused at compile.
        let one = grow(1);
        let two = [&one[..27], &[0x05, 0x05, 0x02, 0x00, 0x01, 0x00, 0x01], &one[32..]].concat();
        let mut wide = one.clone();
        wide[30] = 0x04;
        for module in [two, wide] {
            let e = wasm.run(Arc::new(module), None, "grow".into(), &input, 0).await.unwrap_err();
            assert!(e.starts_with("WASM_ABI_INVALID: "), "{e}");
        }
    }

    // Each spin executes for its full wall time T, so with 2 slots the (2k+1)th and (2k+2)th runs end
    // no earlier than (k+1)·T (never more than 2 at once), and the first two end together (2 at once).
    // The echo queued behind them waits ~3T, longer than T, and still succeeds.
    #[tokio::test]
    async fn concurrency_is_bounded_and_queued_time_is_not_wall_time() {
        const T: u64 = 300;
        let wasm = Wasm::with_concurrency(T, 2);
        let module = Arc::new(toy());
        let started = std::time::Instant::now();
        let run = |entry: &'static str| {
            let (wasm, module) = (wasm.clone(), module.clone());
            tokio::spawn(async move {
                let result = wasm.run(module, None, entry.into(), &serde_json::json!([1]), 0).await;
                (result, started.elapsed())
            })
        };
        let spins: Vec<_> = (0..6).map(|_| run("spin")).collect();
        let echo = run("echo");
        let mut ends = vec![];
        for spin in spins {
            let (result, end) = spin.await.unwrap();
            assert_eq!(result.unwrap_err(), format!("WASM_TIMEOUT: spin ran longer than {T} ms"));
            ends.push(end);
        }
        ends.sort();
        for (i, end) in ends.iter().enumerate() {
            assert!(*end >= Duration::from_millis((i as u64 / 2 + 1) * (T - 10)), "run {i}: {ends:?}");
        }
        assert!(ends[1] < Duration::from_millis(2 * T - 50), "two at once: {ends:?}");
        let (result, end) = echo.await.unwrap();
        assert_eq!(result.unwrap(), serde_json::json!([1]));
        assert!(end >= Duration::from_millis(3 * (T - 10)), "queued behind the spins: {end:?}");
    }

    #[test]
    fn concurrency_env_must_be_a_positive_integer() {
        let parse = |v: &str| parse_concurrency(Some(v.into()));
        assert_eq!((parse("2"), parse(" 3 ")), (Ok(2), Ok(3)));
        assert!(parse("").unwrap() >= 1 && parse_concurrency(None).unwrap() >= 1, "unset: the cores");
        for bad in ["0", "-1", "two", "1.5"] {
            assert_eq!(parse(bad).unwrap_err(), format!("ATLIBS_WASM_CONCURRENCY must be a positive integer, not {bad:?}"));
        }
    }

    #[tokio::test]
    async fn limits_and_errors_match_the_js_host() {
        let wasm = Wasm::new(200);
        let module = Arc::new(toy());
        let input = serde_json::json!({ "a": [1, "b"] });
        assert_eq!(
            wasm.run(module.clone(), None, "echo".into(), &input, 0)
                .await
                .unwrap(),
            input
        );
        let digest = sha256_hex(&module);
        assert!(
            wasm.run(
                module.clone(),
                Some(digest.to_uppercase()),
                "echo".into(),
                &input, 0
            )
            .await
            .is_ok()
        );
        // input past 1 MiB: refused by default, let through by a larger io (then it traps: one page of
        // memory), and never past IO_BYTES_MAX
        let big = serde_json::json!({ "x": "x".repeat(IO_BYTES) });
        assert_eq!(wasm.run(module.clone(), None, "echo".into(), &big, 0).await.unwrap_err(), format!("WASM_TOO_LARGE: input exceeds {IO_BYTES} bytes"));
        assert!(wasm.run(module.clone(), None, "echo".into(), &big, 2 << 20).await.unwrap_err().starts_with("WASM_TRAP"));
        let huge = serde_json::json!({ "x": "x".repeat(IO_BYTES_MAX) });
        assert_eq!(wasm.run(module.clone(), None, "echo".into(), &huge, 1 << 30).await.unwrap_err(), format!("WASM_TOO_LARGE: input exceeds {IO_BYTES_MAX} bytes"));
        assert_eq!(
            wasm.run(module.clone(), Some("00".into()), "echo".into(), &input, 0)
                .await
                .unwrap_err(),
            "WASM_DIGEST_MISMATCH: module digest does not match sha256"
        );
        assert_eq!(
            wasm.run(module.clone(), None, "boom".into(), &input, 0)
                .await
                .unwrap_err(),
            "WASM_TRAP: unreachable"
        );
        assert_eq!(
            wasm.run(module.clone(), None, "nope".into(), &input, 0)
                .await
                .unwrap_err(),
            "WASM_ABI_INVALID: module must export memory, alloc and nope"
        );
        let entries = wasm.entries(&module).unwrap();
        assert!(entries.contains(&"echo".to_string()) && entries.contains(&"boom".to_string()) && !entries.iter().any(|e| e == "alloc" || e == "memory"), "{entries:?}");
        assert_eq!(
            wasm.run(module.clone(), None, "1x".into(), &input, 0)
                .await
                .unwrap_err(),
            "WASM_ABI_INVALID: entry must be an export name"
        );
        let started = std::time::Instant::now();
        assert_eq!(
            wasm.run(module, None, "spin".into(), &input, 0)
                .await
                .unwrap_err(),
            "WASM_TIMEOUT: spin ran longer than 200 ms"
        );
        assert!(started.elapsed() < Duration::from_secs(2));
    }
}