solidb 1.2.1

use std::collections::HashMap;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use mlua::{Lua, Value as LuaValue};
use serde_json::Value as JsonValue;
use tokio::sync::broadcast;

use crate::error::DbError;
use crate::scripting::channel_manager::ChannelManager;
use crate::storage::StorageEngine;
use crate::stream::StreamManager;

use super::conversion::lua_to_json_value;
use super::types::{Script, ScriptContext, ScriptDbName, ScriptResult, ScriptStats};

pub mod cache;
pub mod globals;
pub mod pool;
pub mod repl;
pub mod script_index;
pub mod websocket;

pub use cache::ScriptCache;
pub use pool::LuaPool;
pub use script_index::ScriptIndex;

/// `SOLIDB_NO_LUA=1` / `true` / `yes` skips the VM pool and refuses script
/// execution. `--no-lua` sets the same variable at process start.
///
/// Resolved once: `main` writes the variable before the Tokio runtime starts,
/// and every script execution asks this, so re-reading the environment (and
/// allocating a `String`) per call buys nothing.
pub fn lua_runtime_enabled() -> bool {
    static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ENABLED.get_or_init(|| !env_flag_is_set(std::env::var("SOLIDB_NO_LUA").ok().as_deref()))
}

pub(crate) fn env_flag_is_set(value: Option<&str>) -> bool {
    matches!(
        value,
        Some(v) if v == "1" || v.eq_ignore_ascii_case("true") || v.eq_ignore_ascii_case("yes")
    )
}

pub fn lua_disabled_error() -> DbError {
    DbError::OperationNotSupported(
        "Lua is disabled (--no-lua / SOLIDB_NO_LUA). Custom scripts, services, and the Lua REPL are not available.".to_string(),
    )
}

/// Abort a script that runs past its wall-clock deadline (default 30s,
/// override with `SOLIDB_LUA_TIMEOUT_SECS`, 0 disables). Checked every 50k
/// VM instructions, so a `while true do end` cannot pin a pooled state (and
/// its OS thread) forever. The caller must call `remove_deadline_hook` afterwards.
/// Stop a long-lived script that has stopped yielding.
///
/// WebSocket service scripts live as long as the connection and spend most
/// of it awaiting messages, so a wall-clock deadline from start would kill
/// every legitimate one. What must not happen is `while true do end`: the VM
/// never yields, and the tokio worker polling it never returns to the
/// scheduler. The hook fires every 50k instructions and adds up the time
/// between consecutive firings while they are close together — that is
/// time spent executing Lua without a break. A gap longer than a few
/// milliseconds means the script awaited (or ran a host call) and the count
/// starts over. Continuous execution past the script timeout is an error.
fn install_busy_loop_hook(lua: &Lua) {
    let timeout = script_timeout_secs();
    if timeout == 0 {
        return;
    }
    let limit = std::time::Duration::from_secs(timeout);
    let state = std::sync::Mutex::new((std::time::Instant::now(), std::time::Duration::ZERO));
    let _ = lua.set_global_hook(
        mlua::HookTriggers::new().every_nth_instruction(50_000),
        move |_lua, _debug| {
            let now = std::time::Instant::now();
            let mut st = state.lock().unwrap_or_else(|e| e.into_inner());
            let gap = now.duration_since(st.0);
            st.0 = now;
            if gap > std::time::Duration::from_millis(20) {
                st.1 = std::time::Duration::ZERO;
            } else {
                st.1 += gap;
            }
            if st.1 > limit {
                Err(mlua::Error::RuntimeError(format!(
                    "script ran for {}s without yielding (execution time limit)",
                    timeout
                )))
            } else {
                Ok(mlua::VmState::Continue)
            }
        },
    );
}

fn script_timeout_secs() -> u64 {
    static TIMEOUT_SECS: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
    *TIMEOUT_SECS.get_or_init(|| {
        std::env::var("SOLIDB_LUA_TIMEOUT_SECS")
            .ok()
            .and_then(|v| v.parse::<u64>().ok())
            .unwrap_or(30)
    })
}

/// Turn a Lua failure into the error the caller should see.
///
/// A `DbError` raised inside a binding (a refused write to `_scripts`, a
/// missing document) travels out of mlua as `CallbackError { cause:
/// ExternalError(..) }`. It used to be flattened into
/// `InternalError("Lua error: ...")`, so every refusal was a 500 with the
/// real reason buried in the message. `Forbidden` and friends now come back
/// out as themselves.
pub(crate) fn lua_error_to_db_error(e: mlua::Error) -> DbError {
    fn db_error_in(e: &mlua::Error) -> Option<DbError> {
        match e {
            mlua::Error::CallbackError { cause, .. } => db_error_in(cause),
            mlua::Error::WithContext { cause, .. } => db_error_in(cause),
            mlua::Error::ExternalError(arc) => match arc.downcast_ref::<DbError>() {
                Some(DbError::Forbidden(m)) => Some(DbError::Forbidden(m.clone())),
                Some(DbError::Unauthorized(m)) => Some(DbError::Unauthorized(m.clone())),
                _ => None,
            },
            // `solidb.error(msg, code)` raises "ERROR:{code}:{msg}"; Lua's
            // own `error()` on such a message prefixes the location, so
            // search rather than anchor.
            mlua::Error::RuntimeError(m) => script_error_in(m),
            _ => None,
        }
    }
    fn script_error_in(m: &str) -> Option<DbError> {
        let at = m.find("ERROR:")?;
        let rest = &m[at + "ERROR:".len()..];
        let (code, message) = rest.split_once(':')?;
        let status: u16 = code.trim().parse().ok()?;
        if !(100..=599).contains(&status) {
            return None;
        }
        Some(DbError::ScriptError {
            status,
            message: message.trim().to_string(),
        })
    }
    db_error_in(&e).unwrap_or_else(|| DbError::InternalError(format!("Lua error: {}", e)))
}

/// The script wall-clock budget, or `None` when `SOLIDB_LUA_TIMEOUT_SECS=0`
/// disables it.
pub(crate) fn script_timeout() -> Option<std::time::Duration> {
    match script_timeout_secs() {
        0 => None,
        secs => Some(std::time::Duration::from_secs(secs)),
    }
}

fn install_deadline_hook(lua: &Lua) {
    if let Some(limit) = script_timeout() {
        install_deadline_hook_for(lua, limit);
    }
}

/// Abort the script once `limit` of wall-clock time has passed. Checked
/// every 50k VM instructions; time spent awaiting a host call is bounded
/// separately by the caller (see [`eval_with_deadline`]).
pub(crate) fn install_deadline_hook_for(lua: &Lua, limit: std::time::Duration) {
    let deadline = std::time::Instant::now() + limit;
    let hook = move |_lua: &Lua, _debug: &mlua::Debug| {
        if std::time::Instant::now() > deadline {
            Err(mlua::Error::RuntimeError(format!(
                "script exceeded execution time limit ({}ms)",
                limit.as_millis()
            )))
        } else {
            Ok(mlua::VmState::Continue)
        }
    };
    let triggers = mlua::HookTriggers::new().every_nth_instruction(50_000);
    // `set_hook` covers only the main thread. `eval_async` runs the chunk in
    // a new coroutine thread, which takes the *global* hook, so without the
    // second call `while true do end` never hits the deadline there.
    let _ = lua.set_hook(triggers, hook);
    let _ = lua.set_global_hook(triggers, hook);
}

/// Undo [`install_deadline_hook_for`].
pub(crate) fn remove_deadline_hook(lua: &Lua) {
    lua.remove_hook();
    lua.remove_global_hook();
}

/// Evaluate `chunk` asynchronously, giving up after `limit`.
///
/// The instruction hook cannot fire while the script is parked in an async
/// binding (`time.sleep`, `fetch`, `solidb.timeout`), so the await itself is
/// bounded too.
pub(crate) async fn eval_with_deadline(
    chunk: mlua::Chunk<'_>,
    limit: Option<std::time::Duration>,
) -> mlua::Result<LuaValue> {
    match limit {
        Some(limit) => tokio::time::timeout(limit, chunk.eval_async::<LuaValue>())
            .await
            .unwrap_or_else(|_| {
                Err(mlua::Error::RuntimeError(format!(
                    "script exceeded execution time limit ({}ms)",
                    limit.as_millis()
                )))
            }),
        None => chunk.eval_async::<LuaValue>().await,
    }
}

/// Lua scripting engine
///
/// Cheap to clone: every field is shared. The pooled path moves a clone
/// onto a blocking thread (audit A5).
#[derive(Clone)]
pub struct ScriptEngine {
    pub(crate) storage: Arc<StorageEngine>,
    pub(crate) queue_notifier: Option<broadcast::Sender<()>>,
    pub(crate) stream_manager: Option<Arc<StreamManager>>,
    pub(crate) channel_manager: Option<Arc<ChannelManager>>,
    pub(crate) stats: Arc<ScriptStats>,
    /// Optional Lua VM pool for efficient state reuse
    pub(crate) lua_pool: Option<Arc<LuaPool>>,
    /// Optional bytecode cache for avoiding recompilation
    pub(crate) script_cache: Option<Arc<ScriptCache>>,
}

impl ScriptEngine {
    /// Create a new script engine with access to the storage layer
    pub fn new(storage: Arc<StorageEngine>, stats: Arc<ScriptStats>) -> Self {
        Self {
            storage,
            queue_notifier: None,
            stream_manager: None,
            channel_manager: None,
            stats,
            lua_pool: None,
            script_cache: None,
        }
    }

    pub fn with_queue_notifier(mut self, notifier: broadcast::Sender<()>) -> Self {
        self.queue_notifier = Some(notifier);
        self
    }

    pub fn with_stream_manager(mut self, manager: Arc<StreamManager>) -> Self {
        self.stream_manager = Some(manager);
        self
    }

    pub fn with_channel_manager(mut self, manager: Arc<ChannelManager>) -> Self {
        self.channel_manager = Some(manager);
        self
    }

    /// Configure the engine to use a Lua VM pool for efficient state reuse.
    ///
    /// This dramatically reduces per-request overhead by reusing pre-initialized
    /// Lua states instead of creating new ones for each request.
    pub fn with_lua_pool(mut self, pool: Arc<LuaPool>) -> Self {
        self.lua_pool = Some(pool);
        self
    }

    /// Configure the engine to use a bytecode cache.
    ///
    /// This avoids recompiling scripts on every request by caching
    /// the compiled bytecode.
    pub fn with_script_cache(mut self, cache: Arc<ScriptCache>) -> Self {
        self.script_cache = Some(cache);
        self
    }

    /// Execute a Lua script with the given context
    pub async fn execute(
        &self,
        script: &Script,
        db_name: &str,
        context: &ScriptContext,
    ) -> Result<ScriptResult, DbError> {
        if !lua_runtime_enabled() {
            return Err(lua_disabled_error());
        }
        // Use Relaxed ordering for stats - exact counts not critical for performance
        self.stats.active_scripts.fetch_add(1, Ordering::Relaxed);
        self.stats
            .total_scripts_executed
            .fetch_add(1, Ordering::Relaxed);

        // Ensure active counter is decremented even on panic or early return
        // Use a reference to avoid Arc clone overhead
        struct ActiveScriptGuard<'a>(&'a ScriptStats);
        impl Drop for ActiveScriptGuard<'_> {
            fn drop(&mut self) {
                self.0.active_scripts.fetch_sub(1, Ordering::Relaxed);
            }
        }
        let _guard = ActiveScriptGuard(&self.stats);

        // Use pooled Lua state if available, otherwise create a new one
        if let Some(ref pool) = self.lua_pool {
            self.execute_with_pool(pool, script, db_name, context).await
        } else {
            self.execute_without_pool(script, db_name, context).await
        }
    }

    /// Execute using a pooled Lua state (fast path)
    ///
    /// This method uses a two-tier globals system for maximum performance:
    /// - Static globals (crypto, time, json, etc.) are initialized once per pool state
    /// - Per-request globals (db, request, context, etc.) are set up only if referenced
    async fn execute_with_pool(
        &self,
        pool: &Arc<LuaPool>,
        script: &Script,
        db_name: &str,
        context: &ScriptContext,
    ) -> Result<ScriptResult, DbError> {
        // Every state busy means `pool_size` scripts are already running.
        // Wait for one asynchronously — the old fallback spun a tokio worker
        // at 100% until a state freed up, with no exit — and give up after
        // the script timeout rather than queueing forever.
        let acquire_deadline = std::time::Instant::now()
            + std::time::Duration::from_secs(script_timeout_secs().max(1));
        let pool_guard = loop {
            if let Some(guard) = pool.try_acquire() {
                break guard;
            }
            if std::time::Instant::now() >= acquire_deadline {
                return Err(DbError::InternalError(
                    "Script engine busy: every Lua state is in use".to_string(),
                ));
            }
            tokio::time::sleep(std::time::Duration::from_millis(1)).await;
        };

        // Analyze what globals this script needs (cached in script_cache)
        let needs = if let Some(ref cache) = self.script_cache {
            cache.get_or_analyze_needs(&script.key, &script.code)
        } else {
            globals::ScriptNeeds::analyze(&script.code)
        };

        // Run on the blocking pool, not on a tokio worker (audit A5). The
        // pool used to be sized to the worker count, so as many slow scripts
        // — anonymous ones included — pinned every runtime worker for up to
        // the script timeout. On a blocking thread the script is driven by
        // `Handle::block_on(eval_async)`, which also makes the async
        // bindings (`fetch`, `crypto.hash_password`, `time.sleep`,
        // `solidb.timeout`/`retry`/`try`) work here: under the previous
        // synchronous `eval` they could not yield.
        let engine = self.clone();
        let script = script.clone();
        let db_name = db_name.to_string();
        let context = context.clone();
        let handle = tokio::runtime::Handle::current();
        let task = tokio::task::spawn_blocking(move || {
            pool_guard.with_lua(|lua| {
                engine.run_pooled(&handle, lua, &script, &db_name, &context, &needs)
            })
        });
        match task.await {
            Ok(result) => result,
            Err(e) => Err(DbError::InternalError(format!(
                "Script execution task failed: {}",
                e
            ))),
        }
    }

    /// The body of [`Self::execute_with_pool`], on a blocking thread with
    /// the pooled state locked.
    fn run_pooled(
        &self,
        handle: &tokio::runtime::Handle,
        lua: &Lua,
        script: &Script,
        db_name: &str,
        context: &ScriptContext,
        needs: &globals::ScriptNeeds,
    ) -> Result<ScriptResult, DbError> {
        // Identity, database and response overrides are installed for every
        // request, not only when the globals analysis found something to set
        // up: a script it misjudged must run as its own caller, never as the
        // previous one on this state (audit C3). Reset also removes them.
        lua.set_app_data(globals::LuaCaller {
            actor: context.write_actor(),
            principal: context.query_principal(),
        });
        lua.set_app_data(ScriptDbName(db_name.to_string()));
        crate::scripting::response::reset_overrides(lua);

        // Set up only the globals this script actually uses
        if needs.any() {
            // Check if static globals are already initialized (two-tier optimization)
            let has_static_globals = lua
                .globals()
                .get::<bool>("__solidb_static_initialized")
                .unwrap_or(false);

            // The per-request `solidb.*` fields are written into the
            // read-only shared tables; only engine code runs in here.
            pool::with_shared_tables_unlocked(lua, || {
                if has_static_globals {
                    // Fast path: only set up per-request globals the script needs
                    globals::setup_request_globals_selective(
                        self,
                        lua,
                        db_name,
                        context,
                        Some((&script.key, &script.name)),
                        Some(needs),
                    )
                } else {
                    // Fallback: set up all globals (for states without static initialization)
                    self.setup_lua_globals(lua, db_name, context, Some((&script.key, &script.name)))
                }
            })?;
        }

        // 2. Get or compile bytecode
        let bytecode = if let Some(ref cache) = self.script_cache {
            cache
                .get_or_compile(&script.key, &script.code, |code| {
                    let chunk = lua.load(code);
                    let func = chunk.into_function()?;
                    Ok(func.dump(false))
                })
                .map_err(|e| DbError::InternalError(format!("Bytecode compilation error: {}", e)))?
        } else {
            // No cache - compile directly
            let chunk = lua.load(&script.code);
            let func = chunk
                .into_function()
                .map_err(|e| DbError::InternalError(format!("Script compilation error: {}", e)))?;
            func.dump(false)
        };

        // 3. Execute the bytecode under a wall-clock deadline: an
        // infinite loop in a script would otherwise pin a pooled state
        // (and its OS thread) forever.
        install_deadline_hook(lua);
        let chunk = lua.load(&bytecode[..]);
        let lua_result = handle.block_on(eval_with_deadline(chunk, script_timeout()));
        remove_deadline_hook(lua);
        let lua_result = lua_result.map_err(lua_error_to_db_error)?;

        // 4. Turn the return value into the response
        self.finish_result(lua, db_name, lua_result)
    }

    /// Execute without pooling (original behavior, used as fallback)
    async fn execute_without_pool(
        &self,
        script: &Script,
        db_name: &str,
        context: &ScriptContext,
    ) -> Result<ScriptResult, DbError> {
        let lua = Lua::new();
        LuaPool::apply_memory_limit(&lua);

        // Secure environment: Remove unsafe standard libraries and functions
        let globals = lua.globals();
        globals
            .set("os", LuaValue::Nil)
            .map_err(|e| DbError::InternalError(format!("Failed to secure os: {}", e)))?;
        globals
            .set("io", LuaValue::Nil)
            .map_err(|e| DbError::InternalError(format!("Failed to secure io: {}", e)))?;
        globals
            .set("debug", LuaValue::Nil)
            .map_err(|e| DbError::InternalError(format!("Failed to secure debug: {}", e)))?;
        globals
            .set("package", LuaValue::Nil)
            .map_err(|e| DbError::InternalError(format!("Failed to secure package: {}", e)))?;
        globals
            .set("dofile", LuaValue::Nil)
            .map_err(|e| DbError::InternalError(format!("Failed to secure dofile: {}", e)))?;
        globals
            .set("load", LuaValue::Nil)
            .map_err(|e| DbError::InternalError(format!("Failed to secure load: {}", e)))?;
        globals
            .set("loadfile", LuaValue::Nil)
            .map_err(|e| DbError::InternalError(format!("Failed to secure loadfile: {}", e)))?;
        globals
            .set("require", LuaValue::Nil)
            .map_err(|e| DbError::InternalError(format!("Failed to secure require: {}", e)))?;

        // Set up the Lua environment
        self.setup_lua_globals(&lua, db_name, context, Some((&script.key, &script.name)))?;

        // Execute the script under the same deadline as the pooled path
        install_deadline_hook(&lua);
        let chunk = lua.load(&script.code);
        let eval_result = eval_with_deadline(chunk, script_timeout()).await;
        remove_deadline_hook(&lua);

        match eval_result {
            Ok(result) => self.finish_result(&lua, db_name, result),
            Err(e) => Err(lua_error_to_db_error(e)),
        }
    }

    /// Execute a Lua script as a WebSocket handler
    pub async fn execute_ws(
        &self,
        script: &Script,
        db_name: &str,
        context: &ScriptContext,
        ws: axum::extract::ws::WebSocket,
    ) -> Result<(), DbError> {
        if !lua_runtime_enabled() {
            return Err(lua_disabled_error());
        }
        websocket::execute_ws(self, script, db_name, context, ws).await
    }

    /// Execute Lua code in REPL mode with variable persistence
    #[allow(clippy::too_many_arguments)]
    pub async fn execute_repl(
        &self,
        code: &str,
        db_name: &str,
        user: crate::scripting::auth::ScriptUser,
        variables: &HashMap<String, JsonValue>,
        history: &[String],
        output_capture: &mut Vec<String>,
        timeout_ms: u64,
    ) -> Result<(JsonValue, HashMap<String, JsonValue>), DbError> {
        if !lua_runtime_enabled() {
            return Err(lua_disabled_error());
        }
        repl::execute_repl(
            self,
            code,
            db_name,
            user,
            variables,
            history,
            output_capture,
            timeout_ms,
        )
        .await
    }

    /// Build the HTTP-facing result from a script's return value.
    ///
    /// Three shapes: an explicit `ScriptResponse` from `response.*`, a
    /// `RawJson` fast path, or any other Lua value serialised as JSON. The
    /// last two honour `solidb.status` / `solidb.header` overrides.
    fn finish_result(
        &self,
        lua: &Lua,
        db_name: &str,
        value: LuaValue,
    ) -> Result<ScriptResult, DbError> {
        use crate::scripting::response::{take_overrides, ResponseBody, ScriptResponse};
        let overrides = take_overrides(lua);

        if let LuaValue::UserData(ref ud) = value {
            if let Ok(resp) = ud.borrow::<ScriptResponse>() {
                let mut headers: HashMap<String, String> = resp.headers.iter().cloned().collect();
                let (body, raw_body) = match &resp.body {
                    ResponseBody::Json(v) => (v.clone(), None),
                    ResponseBody::Raw {
                        content_type,
                        bytes,
                    } => {
                        headers.insert("content-type".to_string(), content_type.clone());
                        (JsonValue::Null, Some(bytes.clone()))
                    }
                    ResponseBody::File { key, filename } => {
                        let (mime, bytes) = self.read_stored_file(db_name, key)?;
                        headers.insert("content-type".to_string(), mime);
                        if let Some(name) = filename {
                            headers.insert(
                                "content-disposition".to_string(),
                                format!("attachment; filename=\"{}\"", name.replace('"', "")),
                            );
                        }
                        (JsonValue::Null, Some(bytes))
                    }
                };
                return Ok(ScriptResult {
                    status: resp.status,
                    body,
                    headers,
                    raw_body,
                });
            }
            if let Ok(raw) = ud.borrow::<crate::scripting::conversion::RawJson>() {
                return Ok(ScriptResult {
                    status: overrides.status.unwrap_or(200),
                    body: JsonValue::Null,
                    headers: overrides.headers.into_iter().collect(),
                    raw_body: Some(raw.0.clone().into_bytes()),
                });
            }
        }
        let body = self.lua_to_json(lua, value)?;
        Ok(ScriptResult {
            status: overrides.status.unwrap_or(200),
            body,
            headers: overrides.headers.into_iter().collect(),
            raw_body: None,
        })
    }

    /// A file stored through `solidb.upload`: its MIME type and bytes.
    fn read_stored_file(&self, db_name: &str, key: &str) -> Result<(String, Vec<u8>), DbError> {
        let database = self.storage.get_database(db_name)?;
        let collection = database
            .get_collection(crate::scripting::file_handling::FILES_COLLECTION)
            .map_err(|_| DbError::DocumentNotFound(key.to_string()))?;
        let doc = collection
            .get(key)
            .map_err(|_| DbError::DocumentNotFound(key.to_string()))?;
        let meta = doc.to_value();
        let mime = meta
            .get("mime_type")
            .and_then(|v| v.as_str())
            .unwrap_or("application/octet-stream")
            .to_string();
        let chunk_count = meta.get("chunks").and_then(|v| v.as_u64()).unwrap_or(1) as u32;
        let mut data = Vec::new();
        for i in 0..chunk_count {
            if let Ok(Some(chunk)) = collection.get_blob_chunk(key, i) {
                data.extend(chunk);
            }
        }
        Ok((mime, data))
    }

    // Helper exposed for submodules
    pub(crate) fn setup_lua_globals(
        &self,
        lua: &Lua,
        db_name: &str,
        context: &ScriptContext,
        script_info: Option<(&str, &str)>,
    ) -> Result<(), DbError> {
        // Every path that builds a script environment goes through here, so
        // this is where the database namespace for cache / rate-limit /
        // channel keys is installed (audit H6).
        lua.set_app_data(ScriptDbName(db_name.to_string()));
        globals::setup_lua_globals(self, lua, db_name, context, script_info)
    }

    /// Convert Lua value to JSON
    pub(crate) fn lua_to_json(&self, lua: &Lua, value: LuaValue) -> Result<JsonValue, DbError> {
        lua_to_json_value(lua, value)
            .map_err(|e| DbError::InternalError(format!("Failed to convert Lua to JSON: {}", e)))
    }
}

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

    #[test]
    fn env_flag_accepts_common_truthy_values() {
        assert!(env_flag_is_set(Some("1")));
        assert!(env_flag_is_set(Some("true")));
        assert!(env_flag_is_set(Some("YES")));
        assert!(!env_flag_is_set(Some("0")));
        assert!(!env_flag_is_set(Some("false")));
        assert!(!env_flag_is_set(None));
        assert!(!env_flag_is_set(Some("")));
    }
}