//! Credential scripts run in a fresh, memory-limited WASI instance. Only two calls run at once.
use super::script::{Error, ErrorKind, Limits, Script};
use std::{
    sync::{Arc, Condvar, Mutex, OnceLock, mpsc},
    time::{Duration, Instant},
};
use wasmtime::{
    Config, Engine, Store, StoreLimits, StoreLimitsBuilder, UpdateDeadline,
    component::{Component, Linker, ResourceTable},
};
use wasmtime_wasi::{
    WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView,
    clocks::HostWallClock,
    p2::{
        bindings::sync::CommandPre,
        pipe::{MemoryInputPipe, MemoryOutputPipe},
    },
    random::Deterministic,
};
use zeroize::Zeroizing;

struct State {
    wasi: WasiCtx,
    table: ResourceTable,
    limits: StoreLimits,
}
impl WasiView for State {
    fn ctx(&mut self) -> WasiCtxView<'_> {
        WasiCtxView {
            ctx: &mut self.wasi,
            table: &mut self.table,
        }
    }
}
struct FixedClock;
impl HostWallClock for FixedClock {
    fn resolution(&self) -> Duration {
        Duration::from_millis(1)
    }
    fn now(&self) -> Duration {
        Duration::ZERO
    }
}
struct Compiled {
    engine: Engine,
    pre: CommandPre<State>,
    active: Mutex<usize>,
    available: Condvar,
}
struct Slot<'a>(&'a Compiled);
impl Drop for Slot<'_> {
    fn drop(&mut self) {
        *self.0.active.lock().unwrap() -= 1;
        self.0.available.notify_one();
    }
}
impl Compiled {
    fn new() -> Result<Arc<Self>, Error> {
        let mut config = Config::new();
        config.epoch_interruption(true);
        // The cache also keeps warm test runs from compiling the component again.
        config.cache(wasmtime::Cache::new(Default::default()).ok());
        let engine = Engine::new(&config).map_err(|_| failure(ErrorKind::Engine))?;
        let component = Component::new(&engine, include_bytes!("../../cred-script/script.wasm"))
            .map_err(|_| failure(ErrorKind::Engine))?;
        let mut linker = Linker::new(&engine);
        wasmtime_wasi::p2::add_to_linker_sync(&mut linker)
            .map_err(|_| failure(ErrorKind::Engine))?;
        let pre = CommandPre::new(
            linker
                .instantiate_pre(&component)
                .map_err(|_| failure(ErrorKind::Engine))?,
        )
        .map_err(|_| failure(ErrorKind::Engine))?;
        Ok(Arc::new(Self {
            engine,
            pre,
            active: Mutex::new(0),
            available: Condvar::new(),
        }))
    }
    fn acquire(&self, deadline: Instant) -> Result<Slot<'_>, Error> {
        let mut active = self.active.lock().map_err(|_| failure(ErrorKind::Engine))?;
        while *active >= 2 {
            let remaining = deadline
                .checked_duration_since(Instant::now())
                .ok_or_else(|| failure(ErrorKind::Limit))?;
            (active, _) = self
                .available
                .wait_timeout(active, remaining)
                .map_err(|_| failure(ErrorKind::Engine))?;
        }
        if Instant::now() >= deadline {
            return Err(failure(ErrorKind::Limit));
        }
        *active += 1;
        Ok(Slot(self))
    }
}
static COMPILED: OnceLock<Result<Arc<Compiled>, Error>> = OnceLock::new();

/// Construct once at startup. Compilation is cached; each evaluation gets a fresh Store.
#[derive(Clone)]
pub struct QuickJs {
    compiled: Result<Arc<Compiled>, Error>,
}
impl Default for QuickJs {
    fn default() -> Self {
        Self {
            compiled: COMPILED.get_or_init(Compiled::new).clone(),
        }
    }
}
fn failure(kind: ErrorKind) -> Error {
    Error {
        kind,
        message: match kind {
            ErrorKind::Limit => "script time, memory or stack limit exceeded",
            ErrorKind::Syntax => "invalid script or JSON",
            ErrorKind::Throw => "script threw or rejected",
            ErrorKind::NonJson => "result is not JSON",
            ErrorKind::Engine => "script engine failure",
        }
        .into(),
    }
}
impl Script for QuickJs {
    fn evaluate(
        &self,
        source: &str,
        input: &str,
        export: Option<&str>,
        limits: Limits,
    ) -> Result<String, Error> {
        let compiled = self.compiled.as_ref().map_err(Clone::clone)?;
        let time = limits.max_time_ms.unwrap_or(100).clamp(1, 1000);
        let heap = limits
            .max_memory_bytes
            .unwrap_or(8 * 1024 * 1024)
            .clamp(256 * 1024, 32 * 1024 * 1024);
        let deadline = Instant::now() + Duration::from_millis(time);
        let _slot = compiled.acquire(deadline)?;
        if source
            .len()
            .saturating_add(input.len())
            .saturating_add(export.map_or(0, str::len))
            > heap
        {
            return Err(failure(ErrorKind::Limit));
        }
        // WASI's synchronous adapter must not enter a caller's Tokio runtime.
        std::thread::scope(|scope| {
            let worker = std::thread::Builder::new()
                .name("cred-script".into())
                .spawn_scoped(scope, move || {
                    #[derive(serde::Serialize)]
                    struct Request<'a> {
                        source: &'a str,
                        input: &'a str,
                        export: Option<&'a str>,
                        max_time_ms: u64,
                        max_memory_bytes: usize,
                    }
                    let wire = Zeroizing::new(
                        serde_json::to_vec(&Request {
                            source,
                            input,
                            export,
                            max_time_ms: time,
                            max_memory_bytes: heap,
                        })
                        .map_err(|_| failure(ErrorKind::Engine))?,
                    );
                    let memory = heap + 8 * 1024 * 1024;
                    let output = MemoryOutputPipe::new(memory);
                    let mut wasi = WasiCtxBuilder::new();
                    wasi.stdin(MemoryInputPipe::new(bytes::Bytes::from_owner(wire)))
                        .stdout(output.clone())
                        .wall_clock(FixedClock)
                        .secure_random(Deterministic::new(vec![0]))
                        .insecure_random(Deterministic::new(vec![0]))
                        .insecure_random_seed(0)
                        .allow_tcp(false)
                        .allow_udp(false)
                        .allow_ip_name_lookup(false);
                    let mut store = Store::new(
                        &compiled.engine,
                        State {
                            wasi: wasi.build(),
                            table: ResourceTable::new(),
                            limits: StoreLimitsBuilder::new()
                                .memory_size(memory)
                                .memories(1)
                                .trap_on_grow_failure(true)
                                .build(),
                        },
                    );
                    store.limiter(|state| &mut state.limits);
                    store.epoch_deadline_callback(move |_| {
                        if Instant::now() >= deadline {
                            Err(wasmtime::Trap::Interrupt.into())
                        } else {
                            Ok(UpdateDeadline::Continue(1))
                        }
                    });
                    store.set_epoch_deadline(1);
                    let (done, wait) = mpsc::channel();
                    let engine = compiled.engine.clone();
                    let watchdog = std::thread::Builder::new()
                        .name("cred-script-deadline".into())
                        .spawn(move || {
                            if wait
                                .recv_timeout(deadline.saturating_duration_since(Instant::now()))
                                .is_err()
                            {
                                engine.increment_epoch();
                            }
                        })
                        .map_err(|_| failure(ErrorKind::Engine))?;
                    let result = compiled
                        .pre
                        .instantiate(&mut store)
                        .and_then(|command| command.wasi_cli_run().call_run(&mut store));
                    let _ = done.send(());
                    let _ = watchdog.join();
                    if Instant::now() >= deadline || !matches!(result, Ok(Ok(()))) {
                        return Err(failure(ErrorKind::Limit));
                    }
                    #[derive(serde::Deserialize)]
                    struct Response {
                        result: Option<String>,
                        kind: Option<String>,
                    }
                    let response: Response = serde_json::from_slice(&output.contents())
                        .map_err(|_| failure(ErrorKind::Engine))?;
                    if Instant::now() >= deadline {
                        return Err(failure(ErrorKind::Limit));
                    }
                    if let Some(result) = response.result {
                        return Ok(result);
                    }
                    Err(failure(match response.kind.as_deref() {
                        Some("Limit") => ErrorKind::Limit,
                        Some("Syntax") => ErrorKind::Syntax,
                        Some("Throw") => ErrorKind::Throw,
                        Some("NonJson") => ErrorKind::NonJson,
                        _ => ErrorKind::Engine,
                    }))
                })
                .map_err(|_| failure(ErrorKind::Engine))?;
            worker.join().map_err(|_| failure(ErrorKind::Engine))?
        })
    }
}

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

    #[test]
    fn concurrency_is_bounded_and_queue_uses_the_deadline() {
        let runtime = QuickJs::default();
        let compiled = runtime.compiled.as_ref().unwrap();
        let deadline = Instant::now() + Duration::from_secs(1);
        let first = compiled.acquire(deadline).unwrap();
        let second = compiled.acquire(deadline).unwrap();
        assert!(
            matches!(compiled.acquire(Instant::now()), Err(error) if error.kind == ErrorKind::Limit)
        );
        drop(first);
        let third = compiled.acquire(deadline).unwrap();
        assert_eq!(*compiled.active.lock().unwrap(), 2);
        drop((second, third));
    }
}