solidb 1.0.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, 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 `lua.remove_hook()` afterwards.
fn install_deadline_hook(lua: &Lua) {
    static TIMEOUT_SECS: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
    let timeout = *TIMEOUT_SECS.get_or_init(|| {
        std::env::var("SOLIDB_LUA_TIMEOUT_SECS")
            .ok()
            .and_then(|v| v.parse::<u64>().ok())
            .unwrap_or(30)
    });
    if timeout == 0 {
        return;
    }
    let deadline = std::time::Instant::now() + std::time::Duration::from_secs(timeout);
    let _ = lua.set_hook(
        mlua::HookTriggers::new().every_nth_instruction(50_000),
        move |_lua, _debug| {
            if std::time::Instant::now() > deadline {
                Err(mlua::Error::RuntimeError(format!(
                    "script exceeded execution time limit ({}s)",
                    timeout
                )))
            } else {
                Ok(mlua::VmState::Continue)
            }
        },
    );
}

/// Lua scripting engine
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> {
        let pool_guard = pool.acquire();

        // 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)
        };

        // SINGLE lock acquisition for entire operation
        pool_guard.with_lua(|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);

                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 = chunk.eval::<LuaValue>();
            lua.remove_hook();
            let lua_result =
                lua_result.map_err(|e| DbError::InternalError(format!("Lua error: {}", e)))?;

            // 4. Check for RawJson (pre-serialized) or convert to JSON
            if let LuaValue::UserData(ref ud) = lua_result {
                if let Ok(raw) = ud.borrow::<crate::scripting::conversion::RawJson>() {
                    return Ok(ScriptResult {
                        status: 200,
                        body: JsonValue::Null,
                        headers: HashMap::new(),
                        raw_body: Some(raw.0.clone()),
                    });
                }
            }
            let body = self.lua_to_json(lua, lua_result)?;
            Ok(ScriptResult {
                status: 200,
                body,
                headers: HashMap::new(),
                raw_body: None,
            })
        })
    }

    /// 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 = chunk.eval_async::<LuaValue>().await;
        lua.remove_hook();

        match eval_result {
            Ok(result) => {
                // Convert Lua result to JSON
                let json_result = self.lua_to_json(&lua, result)?;
                Ok(ScriptResult {
                    status: 200,
                    body: json_result,
                    headers: HashMap::new(),
                    raw_body: None,
                })
            }
            Err(e) => Err(DbError::InternalError(format!("Lua 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
    pub async fn execute_repl(
        &self,
        code: &str,
        db_name: &str,
        variables: &HashMap<String, JsonValue>,
        history: &[String],
        output_capture: &mut Vec<String>,
    ) -> Result<(JsonValue, HashMap<String, JsonValue>), DbError> {
        if !lua_runtime_enabled() {
            return Err(lua_disabled_error());
        }
        repl::execute_repl(self, code, db_name, variables, history, output_capture).await
    }

    // 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> {
        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("")));
    }
}