lua_udf.hpp

/**
 * @file lua_udf.hpp
 * @brief Register Lua-authored functions as DuckDB SQL UDFs
 *
 * What this buys you
 * ------------------
 * New SQL functions can be written in Lua and used without recompiling
 * anything. The script declares an `exports` table:
 *
 *     exports = { robust_cv = "scalar", summary = "text" }
 *
 * and at load time this file walks that table and registers `lua_stat_<name>`
 * for each entry, bound to the Lua global `lua_<name>`. Nothing here contains a
 * list of function names.
 *
 * That is a deliberate departure from the duckdb-lua-statcpp prototype, where
 * the SQL-name-to-Lua-name mapping was a hardcoded C++ list. In the prototype
 * you could edit the *body* of the seven known functions without rebuilding,
 * but adding an eighth still meant touching C++. Now it does not.
 *
 * Script resolution (first hit wins)
 * ----------------------------------
 *   1. $STATCPP_LUA_SCRIPT
 *   2. $HOME/<STATCPP_LUA_DEFAULT_RELATIVE_PATH>   (~/.duckdb/statcpp/stats.lua)
 *   3. the copy embedded at build time
 *
 * The embedded copy is a fallback, not the primary path: it only guarantees the
 * extension does something sensible out of the box. An external script always
 * takes precedence, which is what keeps the edit-and-reload workflow intact for
 * a single self-contained binary.
 *
 * Threading
 * ---------
 * A lua_State is not thread-safe, and DuckDB executes scalar functions on many
 * threads at once. The prototype shared one state across every UDF, which was a
 * latent data race that only stayed quiet because the demo was single-threaded.
 * Here a LuaRuntime hands out states from a small pool: a thread borrows one for
 * the duration of a chunk and returns it afterwards. The mutex is held only
 * while handing a pointer over, never while Lua is running, so states scale with
 * actual concurrency instead of serialising it.
 *
 * Conventions at the boundary (identical to the C++ UDFs)
 * ------------------------------------------------------
 *  - A NULL element of the input LIST reaches Lua as NaN, so that `ipairs` does
 *    not stop early the way it would on a nil hole.
 *  - A NaN or nil coming back from Lua becomes SQL NULL.
 *  - A Lua error is not a query error: the row yields SQL NULL, matching how
 *    statcpp failures are absorbed elsewhere.
 */

#pragma once

#include <cmath>
#include <cstdlib>
#include <fstream>
#include <limits>
#include <memory>
#include <mutex>
#include <sstream>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

extern "C" {
#include <lauxlib.h>
#include <lua.h>
#include <lualib.h>
}

#include "capi_util.hpp"
#include "embedded_lua.hpp"
#include "lua_statcpp_bindings.hpp"

namespace lua_duckdb {

using statcpp_duckdb::DoubleWriter;
using statcpp_duckdb::ListReader;
using statcpp_duckdb::ListWriter;
using statcpp_duckdb::LogicalType;
using statcpp_duckdb::StringWriter;
using statcpp_duckdb::Vec;

// ---------------------------------------------------------------------------
// Script resolution
// ---------------------------------------------------------------------------

/// Where a script came from, reported through lua_stat_script_origin().
struct LuaScript {
    std::string source;  ///< The script text.
    std::string origin;  ///< Human-readable provenance, for diagnostics.
};

/// Read a whole file. Returns false if it cannot be opened.
inline bool ReadFile(const std::string& path, std::string& out) {
    std::ifstream in(path, std::ios::binary);
    if (!in) {
        return false;
    }
    std::ostringstream buffer;
    buffer << in.rdbuf();
    out = buffer.str();
    return true;
}

/**
 * @brief Resolve the Lua script to load.
 *
 * Never fails: the embedded copy is always available as a last resort.
 */
inline LuaScript ResolveScript() {
    LuaScript script;

    if (const char* env = std::getenv("STATCPP_LUA_SCRIPT")) {
        if (env[0] != '\0' && ReadFile(env, script.source)) {
            script.origin = std::string("STATCPP_LUA_SCRIPT=") + env;
            return script;
        }
    }

    if (const char* home = std::getenv("HOME")) {
        const std::string path = std::string(home) + "/" + STATCPP_LUA_DEFAULT_RELATIVE_PATH;
        if (ReadFile(path, script.source)) {
            script.origin = path;
            return script;
        }
    }

    script.source = kEmbeddedLuaScript;
    script.origin = "embedded (built into the extension)";
    return script;
}

// ---------------------------------------------------------------------------
// Lua runtime
// ---------------------------------------------------------------------------

/**
 * @brief Owns the script text and a pool of interpreter states.
 *
 * Every state is created from the same source, so which one a thread happens to
 * borrow makes no observable difference. Scripts that keep mutable global state
 * between calls are the one exception, and are outside what this layer promises.
 */
class LuaRuntime {
public:
    explicit LuaRuntime(LuaScript script) : script_(std::move(script)) {}

    ~LuaRuntime() {
        for (lua_State* state : all_) {
            lua_close(state);
        }
    }

    LuaRuntime(const LuaRuntime&) = delete;
    LuaRuntime& operator=(const LuaRuntime&) = delete;

    [[nodiscard]] const std::string& origin() const {
        return script_.origin;
    }

    /**
     * @brief Borrow a state, creating one if none is free.
     * @throws std::runtime_error if the script fails to load.
     */
    lua_State* Acquire() {
        {
            const std::lock_guard<std::mutex> lock(mutex_);
            if (!free_.empty()) {
                lua_State* state = free_.back();
                free_.pop_back();
                return state;
            }
        }
        // Built outside the lock: loading a script is slow and does not need it.
        lua_State* state = NewState();
        const std::lock_guard<std::mutex> lock(mutex_);
        all_.push_back(state);
        return state;
    }

    void Release(lua_State* state) {
        if (state == nullptr) {
            return;
        }
        const std::lock_guard<std::mutex> lock(mutex_);
        free_.push_back(state);
    }

private:
    lua_State* NewState() const {
        lua_State* state = lua_statcpp::CreateLuaState();
        if (state == nullptr) {
            throw std::runtime_error("statcpp: could not create a Lua state");
        }
        const int status = luaL_loadbuffer(state, script_.source.data(), script_.source.size(),
                                           script_.origin.c_str()) ||
                           lua_pcall(state, 0, 0, 0);
        if (status != LUA_OK) {
            const std::string error = lua_tostring(state, -1) ? lua_tostring(state, -1) : "unknown";
            lua_close(state);
            throw std::runtime_error("statcpp: failed to load Lua script (" + script_.origin +
                                     "): " + error);
        }
        return state;
    }

    LuaScript script_;
    std::mutex mutex_;
    std::vector<lua_State*> free_;
    std::vector<lua_State*> all_;
};

/// RAII borrow of a state from a LuaRuntime.
class LuaLease {
public:
    explicit LuaLease(LuaRuntime& runtime) : runtime_(runtime), state_(runtime.Acquire()) {}

    ~LuaLease() {
        runtime_.Release(state_);
    }

    LuaLease(const LuaLease&) = delete;
    LuaLease& operator=(const LuaLease&) = delete;

    [[nodiscard]] lua_State* get() const {
        return state_;
    }

private:
    LuaRuntime& runtime_;
    lua_State* state_;
};

// ---------------------------------------------------------------------------
// DuckDB <-> Lua conversion
// ---------------------------------------------------------------------------

/**
 * @brief Push a sample onto the Lua stack as a 1-indexed table.
 *
 * Missing values are pushed as NaN rather than nil, because a nil hole would
 * make `ipairs` stop at the first missing element and silently truncate the
 * sample a script sees.
 */
inline void PushSampleAsTable(lua_State* L, const Vec& values) {
    lua_createtable(L, static_cast<int>(values.size()), 0);
    for (std::size_t i = 0; i < values.size(); ++i) {
        lua_pushnumber(L, values[i]);
        lua_rawseti(L, -2, static_cast<lua_Integer>(i + 1));
    }
}

/// Convert a Lua table at @p index into a sample. NaN/nil elements stay NaN.
inline Vec TableToSample(lua_State* L, int index) {
    Vec out;
    if (!lua_istable(L, index)) {
        return out;
    }
    const int n = static_cast<int>(lua_rawlen(L, index));
    out.reserve(static_cast<std::size_t>(n));
    for (int i = 1; i <= n; ++i) {
        lua_rawgeti(L, index, static_cast<lua_Integer>(i));
        out.push_back(lua_isnil(L, -1) ? statcpp_duckdb::NaN() : lua_tonumber(L, -1));
        lua_pop(L, 1);
    }
    return out;
}

/**
 * @brief Call a global Lua function with one table argument.
 * @return true if the call succeeded and left exactly one value on the stack
 *         (which the caller must pop). On failure nothing is left behind.
 */
inline bool CallLuaWithSample(lua_State* L, const std::string& lua_name, const Vec& values) {
    lua_getglobal(L, lua_name.c_str());
    if (!lua_isfunction(L, -1)) {
        lua_pop(L, 1);
        return false;
    }
    PushSampleAsTable(L, values);
    if (lua_pcall(L, 1, 1, 0) != LUA_OK) {
        lua_pop(L, 1);  // error message
        return false;
    }
    return true;
}

// ---------------------------------------------------------------------------
// Registration helpers
// ---------------------------------------------------------------------------

/// Register LIST<DOUBLE> -> DOUBLE backed by a Lua function.
inline bool RegisterLuaListToScalar(duckdb_connection connection, std::shared_ptr<LuaRuntime> runtime,
                                    const std::string& sql_name, const std::string& lua_name) {
    const LogicalType list_type = statcpp_duckdb::ListOfDoubleType();
    const LogicalType double_type = statcpp_duckdb::DoubleType();
    return statcpp_duckdb::RegisterScalarUdf(
        connection, sql_name, {list_type.get()}, double_type.get(),
        [runtime, lua_name](duckdb_data_chunk input, duckdb_vector output, idx_t rows) {
            const ListReader reader(duckdb_data_chunk_get_vector(input, 0));
            DoubleWriter writer(output);
            LuaLease lease(*runtime);
            lua_State* L = lease.get();
            Vec values;
            for (idx_t row = 0; row < rows; ++row) {
                reader.Read(row, values);
                if (!CallLuaWithSample(L, lua_name, values)) {
                    writer.SetNull(row);
                    continue;
                }
                if (lua_isnil(L, -1)) {
                    writer.SetNull(row);
                } else {
                    writer.Set(row, lua_tonumber(L, -1));
                }
                lua_pop(L, 1);
            }
        });
}

/// Register LIST<DOUBLE> -> LIST<DOUBLE> backed by a Lua function.
inline bool RegisterLuaListToList(duckdb_connection connection, std::shared_ptr<LuaRuntime> runtime,
                                  const std::string& sql_name, const std::string& lua_name) {
    const LogicalType list_type = statcpp_duckdb::ListOfDoubleType();
    return statcpp_duckdb::RegisterScalarUdf(
        connection, sql_name, {list_type.get()}, list_type.get(),
        [runtime, lua_name](duckdb_data_chunk input, duckdb_vector output, idx_t rows) {
            const ListReader reader(duckdb_data_chunk_get_vector(input, 0));
            ListWriter writer(output, rows);
            LuaLease lease(*runtime);
            lua_State* L = lease.get();
            Vec values;
            for (idx_t row = 0; row < rows; ++row) {
                reader.Read(row, values);
                if (!CallLuaWithSample(L, lua_name, values)) {
                    writer.SetNull(row);
                    continue;
                }
                if (lua_istable(L, -1)) {
                    writer.Set(row, TableToSample(L, -1));
                } else {
                    writer.SetNull(row);
                }
                lua_pop(L, 1);
            }
            writer.Finish();
        });
}

/// Register LIST<DOUBLE> -> VARCHAR backed by a Lua function.
inline bool RegisterLuaListToString(duckdb_connection connection,
                                    std::shared_ptr<LuaRuntime> runtime, const std::string& sql_name,
                                    const std::string& lua_name) {
    const LogicalType list_type = statcpp_duckdb::ListOfDoubleType();
    const LogicalType varchar_type = statcpp_duckdb::VarcharType();
    return statcpp_duckdb::RegisterScalarUdf(
        connection, sql_name, {list_type.get()}, varchar_type.get(),
        [runtime, lua_name](duckdb_data_chunk input, duckdb_vector output, idx_t rows) {
            const ListReader reader(duckdb_data_chunk_get_vector(input, 0));
            StringWriter writer(output);
            LuaLease lease(*runtime);
            lua_State* L = lease.get();
            Vec values;
            for (idx_t row = 0; row < rows; ++row) {
                reader.Read(row, values);
                if (!CallLuaWithSample(L, lua_name, values)) {
                    writer.SetNull(row);
                    continue;
                }
                const char* text = lua_tostring(L, -1);
                if (text == nullptr) {
                    writer.SetNull(row);
                } else {
                    writer.Set(row, std::string(text));
                }
                lua_pop(L, 1);
            }
        });
}

// ---------------------------------------------------------------------------
// exports discovery
// ---------------------------------------------------------------------------

/// One entry of the script's `exports` table.
struct LuaExport {
    std::string name;  ///< Base name; SQL sees lua_stat_<name>, Lua defines lua_<name>.
    std::string kind;  ///< "scalar", "list" or "text".
};

/**
 * @brief Read the global `exports` table from a loaded state.
 *
 * A script with no `exports` table registers nothing, which is not an error:
 * it just means that copy of the script exposes no SQL functions.
 */
inline std::vector<LuaExport> ReadExports(lua_State* L) {
    std::vector<LuaExport> exports;
    lua_getglobal(L, "exports");
    if (!lua_istable(L, -1)) {
        lua_pop(L, 1);
        return exports;
    }
    lua_pushnil(L);
    while (lua_next(L, -2) != 0) {
        // key at -2, value at -1. lua_tostring on a key would rewrite it in
        // place and confuse lua_next, so only strings are accepted.
        if (lua_type(L, -2) == LUA_TSTRING && lua_type(L, -1) == LUA_TSTRING) {
            exports.push_back({lua_tostring(L, -2), lua_tostring(L, -1)});
        }
        lua_pop(L, 1);
    }
    lua_pop(L, 1);
    return exports;
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/**
 * @brief Load the Lua script and register everything it exports.
 *
 * @param connection Connection handed to the extension at load time.
 * @param error      Set to a description when false is returned.
 * @return false only if the script itself cannot be loaded. An unknown `kind`
 *         is skipped rather than fatal, so one bad line in a user's script
 *         cannot make the whole extension unloadable.
 *
 * Also registers lua_stat_script_origin(), which reports where the script was
 * read from. With three possible sources, "which stats.lua is actually running"
 * is the first question to ask when a change appears not to take effect.
 */
inline bool RegisterLuaStatcppFunctions(duckdb_connection connection, std::string& error) {
    std::shared_ptr<LuaRuntime> runtime;
    std::vector<LuaExport> exports;
    std::string origin;

    try {
        runtime = std::make_shared<LuaRuntime>(ResolveScript());
        origin = runtime->origin();
        const LuaLease lease(*runtime);
        exports = ReadExports(lease.get());
    } catch (const std::exception& e) {
        error = e.what();
        return false;
    }

    bool ok = true;
    for (const LuaExport& entry : exports) {
        const std::string sql_name = "lua_stat_" + entry.name;
        const std::string lua_name = "lua_" + entry.name;
        if (entry.kind == "scalar") {
            ok = RegisterLuaListToScalar(connection, runtime, sql_name, lua_name) && ok;
        } else if (entry.kind == "list") {
            ok = RegisterLuaListToList(connection, runtime, sql_name, lua_name) && ok;
        } else if (entry.kind == "text") {
            ok = RegisterLuaListToString(connection, runtime, sql_name, lua_name) && ok;
        }
        // Unknown kinds are ignored on purpose; see the doc comment.
    }

    // Diagnostic: which script is in effect. Constant, so no argument is needed.
    const LogicalType varchar_type = statcpp_duckdb::VarcharType();
    ok = statcpp_duckdb::RegisterScalarUdf(
             connection, "lua_stat_script_origin", {}, varchar_type.get(),
             [origin](duckdb_data_chunk, duckdb_vector output, idx_t rows) {
                 StringWriter writer(output);
                 for (idx_t row = 0; row < rows; ++row) {
                     writer.Set(row, origin);
                 }
             }) &&
         ok;

    if (!ok) {
        error = "failed to register one or more Lua-backed functions";
    }
    return ok;
}

}  // namespace lua_duckdb