//! Deterministic Lua sandbox.
//!
//! The sandbox is the verifier (DATA.md): a real Lua interpreter, stripped to a
//! whitelist of the language core plus always-safe builtins. No `io`, `os`,
//! `require`, `package`, no non-deterministic sources. The "timeout" is an
//! instruction budget enforced by a VM hook — deterministic, unlike wall-clock,
//! so verification is reproducible.

use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::Arc;

use mlua::{HookTriggers, Lua, LuaOptions, MultiValue, StdLib, Table, Value as LuaValue, VmState};

use crate::value::LValue;

/// How many VM instructions elapse between hook fires. Coarser = faster.
const HOOK_GRANULARITY: u32 = 256;

#[derive(Debug, Clone, PartialEq)]
pub enum RunError {
    /// Program/function raised a Lua error.
    Lua(String),
    /// Instruction budget exhausted (our deterministic timeout).
    Budget,
    /// Returned value lay outside the observable value space.
    NonObservable,
}

pub struct Sandbox {
    lua: Lua,
    env: Table,
    budget: Arc<AtomicI64>,
}

impl Sandbox {
    pub fn new() -> mlua::Result<Self> {
        // Load only deterministic, side-effect-free standard libraries. No IO,
        // OS, PACKAGE, DEBUG, COROUTINE.
        let lua = Lua::new_with(
            StdLib::MATH | StdLib::STRING | StdLib::TABLE,
            LuaOptions::default(),
        )?;

        let env = build_env(&lua)?;

        // Instruction-budget hook. Shared counter, reset before each operation.
        let budget = Arc::new(AtomicI64::new(0));
        let hook_budget = budget.clone();
        lua.set_hook(
            HookTriggers::new().every_nth_instruction(HOOK_GRANULARITY),
            move |_lua, _debug| {
                let left = hook_budget.fetch_sub(HOOK_GRANULARITY as i64, Ordering::Relaxed);
                if left <= 0 {
                    Err(mlua::Error::RuntimeError("instruction budget exceeded".into()))
                } else {
                    Ok(VmState::Continue)
                }
            },
        );

        Ok(Sandbox { lua, env, budget })
    }

    fn arm(&self, instructions: i64) {
        self.budget.store(instructions, Ordering::Relaxed);
    }

    fn is_budget_error(e: &mlua::Error) -> bool {
        // The hook error is wrapped (CallbackError) as it propagates through the
        // VM, so match on the message rather than the variant.
        e.to_string().contains("instruction budget")
    }

    /// Load a program (defining a global function `f` in the sandbox env) and
    /// execute its top level. Returns Ok if it loads and runs cleanly.
    pub fn load_program(&self, source: &str, budget: i64) -> Result<(), RunError> {
        self.arm(budget);
        let chunk = self
            .lua
            .load(source)
            .set_name("program")
            .set_environment(self.env.clone());
        chunk.exec().map_err(|e| {
            if Self::is_budget_error(&e) {
                RunError::Budget
            } else {
                RunError::Lua(e.to_string())
            }
        })
    }

    /// Call `f(args...)` and return its single observable result.
    pub fn call_f(&self, args: &[LValue], budget: i64) -> Result<LValue, RunError> {
        let f: mlua::Function = self
            .env
            .get("f")
            .map_err(|e| RunError::Lua(format!("no function f: {e}")))?;

        let mut lua_args = MultiValue::new();
        for a in args {
            let v = a
                .to_lua(&self.lua)
                .map_err(|e| RunError::Lua(e.to_string()))?;
            lua_args.push_back(v);
        }

        self.arm(budget);
        let ret: LuaValue = f.call(lua_args).map_err(|e| {
            if Self::is_budget_error(&e) {
                RunError::Budget
            } else {
                RunError::Lua(e.to_string())
            }
        })?;

        LValue::from_lua(&ret).ok_or(RunError::NonObservable)
    }
}

/// Build the restricted global environment: a fresh table holding only
/// whitelisted names copied from the real globals, with non-deterministic
/// entries pruned.
fn build_env(lua: &Lua) -> mlua::Result<Table> {
    let g = lua.globals();
    let env = lua.create_table()?;

    // Safe base functions. Deliberately excluded: print (side effect),
    // load/loadstring/dofile/loadfile (code injection), collectgarbage,
    // require, rawset (mutation escape hatches kept minimal), _G access.
    const BASE: &[&str] = &[
        "assert", "error", "ipairs", "pairs", "next", "select", "tonumber",
        "tostring", "type", "pcall", "xpcall", "rawequal", "rawget", "rawlen",
        "setmetatable", "getmetatable", "unpack",
    ];
    for name in BASE {
        let v: LuaValue = g.get(*name)?;
        if !v.is_nil() {
            env.set(*name, v)?;
        }
    }

    // math, minus non-deterministic random sources.
    if let Ok(math) = g.get::<Table>("math") {
        let m = clone_table_shallow(lua, &math)?;
        m.set("random", LuaValue::Nil)?;
        m.set("randomseed", LuaValue::Nil)?;
        env.set("math", m)?;
    }

    // string and table are fully deterministic.
    if let Ok(string) = g.get::<Table>("string") {
        env.set("string", clone_table_shallow(lua, &string)?)?;
    }
    if let Ok(table) = g.get::<Table>("table") {
        env.set("table", clone_table_shallow(lua, &table)?)?;
    }

    Ok(env)
}

fn clone_table_shallow(lua: &Lua, src: &Table) -> mlua::Result<Table> {
    let dst = lua.create_table()?;
    for pair in src.clone().pairs::<LuaValue, LuaValue>() {
        let (k, v) = pair?;
        dst.set(k, v)?;
    }
    Ok(dst)
}

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

    const B: i64 = 100_000;

    #[test]
    fn forbidden_globals_absent() {
        let sb = Sandbox::new().unwrap();
        for name in ["os", "io", "require", "package", "print", "load", "dofile"] {
            let src = format!("function f(x) return {name} == nil end");
            sb.load_program(&src, B).unwrap();
            assert_eq!(
                sb.call_f(&[LValue::Int(0)], B).unwrap(),
                LValue::Bool(true),
                "{name} should be nil in sandbox"
            );
        }
    }

    #[test]
    fn nondeterministic_sources_removed() {
        let sb = Sandbox::new().unwrap();
        sb.load_program("function f(x) return math.random == nil end", B).unwrap();
        assert_eq!(sb.call_f(&[LValue::Int(0)], B).unwrap(), LValue::Bool(true));
    }

    #[test]
    fn infinite_loop_hits_budget() {
        let sb = Sandbox::new().unwrap();
        sb.load_program("function f(x) while true do end end", B).unwrap();
        assert_eq!(sb.call_f(&[LValue::Int(0)], B), Err(RunError::Budget));
    }

    #[test]
    fn deterministic_arithmetic() {
        let sb = Sandbox::new().unwrap();
        sb.load_program("function f(x) return (x * 2) + 1 end", B).unwrap();
        assert_eq!(sb.call_f(&[LValue::Int(20)], B).unwrap(), LValue::Int(41));
    }

    #[test]
    fn nonobservable_function_return() {
        let sb = Sandbox::new().unwrap();
        sb.load_program("function f(x) return function() end end", B).unwrap();
        assert_eq!(sb.call_f(&[LValue::Int(0)], B), Err(RunError::NonObservable));
    }
}