use mlua::{Lua, Result as LuaResult, VmState};

const MEMORY_LIMIT: usize = 8 * 1024 * 1024;
const INSTRUCTION_LIMIT: u32 = 1_000_000;

pub fn create_sandbox() -> LuaResult<Lua> {
    let lua = Lua::new();

    lua.set_memory_limit(MEMORY_LIMIT)?;

    lua.set_hook(
        mlua::HookTriggers::new().every_nth_instruction(10000),
        move |lua, _| {
            let count: u32 = lua
                .named_registry_value("__instruction_count")
                .unwrap_or(0);
            let new_count = count + 10000;
            lua.set_named_registry_value("__instruction_count", new_count)?;
            if new_count >= INSTRUCTION_LIMIT {
                Err(mlua::Error::RuntimeError(
                    "instruction limit exceeded".to_string(),
                ))
            } else {
                Ok(VmState::Continue)
            }
        },
    )?;

    {
        let globals = lua.globals();
        globals.raw_remove("os")?;
        globals.raw_remove("io")?;
        globals.raw_remove("loadfile")?;
        globals.raw_remove("dofile")?;
        globals.raw_remove("require")?;
        globals.raw_remove("package")?;
        globals.raw_remove("debug")?;
        globals.raw_remove("load")?;
        globals.raw_remove("coroutine")?;
        globals.raw_remove("rawset")?;
        globals.raw_remove("rawget")?;
    }

    Ok(lua)
}

pub fn reset_instruction_counter(lua: &Lua) -> LuaResult<()> {
    lua.set_named_registry_value("__instruction_count", 0u32)
}
