path: root/src/lua_bridge.zig

//! Lua C-API bridge for the panto CLI.
//!
//! Exposes a `panto` global table inside any `lua_State` we construct, with
//! a single function:
//!
//!     panto.register_tool {
//!         name = "...",
//!         description = "...",
//!         schema = { ... },          -- JSON Schema as a Lua table
//!         handler = function(input) ... end,
//!     }
//!
//! The single-table-argument form is idiomatic Lua "named arguments". It's
//! also forward-compatible: future optional fields (examples, version, etc.)
//! can be added without breaking existing extensions.
//!
//! Each call records a registration in a Lua-side table at a fixed registry
//! slot. The Zig side then reads that table to decide what to do with it:
//!
//! - **Discovery** (`harvestRegistrations`): runs an extension script once at
//!   startup to learn the *names*, *descriptions*, and *schemas* of every
//!   tool it declares. The handler functions are discarded — that throwaway
//!   state will be closed immediately.
//!
//! - **Invocation** (`fetchHandler` + `runHandler`): per tool call, we open
//!   a fresh `lua_State`, re-run the script, then look up the handler by
//!   name in the same registry table.
//!
//! No `lua_State` pooling, no shared mutable state across calls. Every
//! `LuaTool.invoke` builds and tears down its own state. This is slow per-
//! call (~ms of Lua startup) but mechanically the simplest model: there is
//! nothing that can leak between invocations.

const std = @import("std");
const Allocator = std.mem.Allocator;

pub const c = @cImport({
    @cInclude("lua.h");
    @cInclude("lauxlib.h");
    @cInclude("lualib.h");
});

// Lua type constants are #defines that translate_c surfaces as inline
// functions returning the literal; using them in switch prongs needs
// explicit `c_int` constants. Define our own clean aliases.
pub const T_NIL: c_int = 0;
pub const T_BOOLEAN: c_int = 1;
pub const T_NUMBER: c_int = 3;
pub const T_STRING: c_int = 4;
pub const T_TABLE: c_int = 5;
pub const T_FUNCTION: c_int = 6;

pub const LUA_MULTRET: c_int = -1;
pub const LUA_REGISTRYINDEX: c_int = -1001000; // matches lua.h with LUAI_MAXSTACK=1000000

/// Errors the bridge can produce when talking to a Lua state.
pub const BridgeError = error{
    LuaInitFailed,
    LuaLoadFailed,
    LuaRunFailed,
    LuaHandlerCrashed,
    LuaHandlerNotFound,
    BadRegistration,
    BadHandlerReturn,
    InputNotJsonObject,
    OutOfMemory,
};

/// The key under which we stash the registrations table in
/// `LUA_REGISTRYINDEX`. Any unique pointer works — we use the address of a
/// module-level `u8` so multiple states all use the same key value.
var registrations_key: u8 = 0;

/// A single declared tool, as harvested from a script's top-level call to
/// `panto.register_tool`. All slices reference Lua-owned strings on the
/// state's stack/registry; copy them before closing the state.
pub const Registration = struct {
    name: []const u8,
    description: []const u8,
    /// Serialized JSON Schema for the tool's input.
    schema_json: []const u8,
};

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

/// Install the `panto.register_tool` global into the given state.
///
/// Also creates the registry table that holds harvested registrations and
/// the per-name handler references.
pub fn install(L: *c.lua_State) void {
    // Create the registrations table: an array of records, each shaped
    // { name=, description=, schema_json=, handler= }.
    c.lua_createtable(L, 0, 0);
    // Stash under our registry key.
    c.lua_rawsetp(L, LUA_REGISTRYINDEX, &registrations_key);

    // Build the `panto` global table with `register_tool`.
    c.lua_createtable(L, 0, 1);
    c.lua_pushcclosure(L, registerToolThunk, 0);
    c.lua_setfield(L, -2, "register_tool");
    c.lua_setglobal(L, "panto");
}

/// Load and execute a Lua source file in the given state. The file's
/// top-level code typically calls `panto.register_tool(...)` one or more
/// times, populating the registrations table.
///
/// On Lua error, the error message is left on the stack — callers that
/// want to surface it can read the top with `lua_tolstring`.
pub fn loadFile(L: *c.lua_State, path: [:0]const u8) BridgeError!void {
    if (c.luaL_loadfilex(L, path.ptr, null) != 0) return BridgeError.LuaLoadFailed;
    if (c.lua_pcallk(L, 0, 0, 0, 0, null) != 0) return BridgeError.LuaRunFailed;
}

/// Walk the registrations table and copy each entry's name, description,
/// and schema_json into freshly-allocated bytes owned by `arena`. The
/// returned slice (and every byte slice inside each `Registration`) lives
/// as long as the arena does.
///
/// The handler field is ignored — discovery mode doesn't care about it.
pub fn harvestRegistrations(
    L: *c.lua_State,
    arena: Allocator,
) BridgeError![]Registration {
    // Push the registrations table.
    _ = c.lua_rawgetp(L, LUA_REGISTRYINDEX, &registrations_key);
    defer c.lua_settop(L, c.lua_gettop(L) - 1); // pop the table when done

    const n: usize = @intCast(c.lua_rawlen(L, -1));
    if (n == 0) return arena.alloc(Registration, 0) catch BridgeError.OutOfMemory;

    var out = arena.alloc(Registration, n) catch return BridgeError.OutOfMemory;
    var i: usize = 1;
    while (i <= n) : (i += 1) {
        _ = c.lua_rawgeti(L, -1, @intCast(i)); // record table on top
        defer c.lua_settop(L, c.lua_gettop(L) - 1);

        const name = try readStringField(L, -1, "name", arena);
        const desc = try readStringField(L, -1, "description", arena);
        const schema = try readStringField(L, -1, "schema_json", arena);
        out[i - 1] = .{
            .name = name,
            .description = desc,
            .schema_json = schema,
        };
    }
    return out;
}

/// In an *invocation-mode* state (registrations table populated by re-
/// running the script), push the handler function for `tool_name` onto the
/// stack. Caller is responsible for popping it after use.
///
/// Returns LuaHandlerNotFound if no registration with that name exists.
pub fn pushHandler(L: *c.lua_State, tool_name: []const u8) BridgeError!void {
    _ = c.lua_rawgetp(L, LUA_REGISTRYINDEX, &registrations_key);
    const n: usize = @intCast(c.lua_rawlen(L, -1));
    var i: usize = 1;
    while (i <= n) : (i += 1) {
        _ = c.lua_rawgeti(L, -1, @intCast(i));
        _ = c.lua_getfield(L, -1, "name");
        var len: usize = 0;
        const ptr = c.lua_tolstring(L, -1, &len);
        const matched = ptr != null and std.mem.eql(u8, ptr[0..len], tool_name);
        c.lua_settop(L, c.lua_gettop(L) - 1); // pop name
        if (matched) {
            // Replace the record with its handler field.
            _ = c.lua_getfield(L, -1, "handler");
            // Stack: ..., regs_table, record, handler. Remove record, regs_table.
            c.lua_copy(L, -1, -3);
            c.lua_settop(L, c.lua_gettop(L) - 2);
            return;
        }
        c.lua_settop(L, c.lua_gettop(L) - 1); // pop record
    }
    c.lua_settop(L, c.lua_gettop(L) - 1); // pop regs table
    return BridgeError.LuaHandlerNotFound;
}

/// Convert raw JSON bytes into a Lua value and push it onto the stack.
/// Top-level value must be a JSON object (matches our schema convention
/// that tool input is always an object).
pub fn pushJsonAsLua(
    L: *c.lua_State,
    arena: Allocator,
    input: []const u8,
) BridgeError!void {
    var parsed = std.json.parseFromSlice(std.json.Value, arena, input, .{}) catch {
        return BridgeError.InputNotJsonObject;
    };
    defer parsed.deinit();
    if (parsed.value != .object) return BridgeError.InputNotJsonObject;
    pushJsonValue(L, parsed.value) catch return BridgeError.OutOfMemory;
}

/// Read a Lua value at `idx` and serialize it to a JSON-compatible owned
/// byte string. Used to convert handler return values into ToolResult
/// content. For now we only accept string returns; extending to richer
/// types is straightforward but unnecessary for slice 2.
pub fn readHandlerResult(
    L: *c.lua_State,
    idx: c_int,
    allocator: Allocator,
) BridgeError![]u8 {
    if (c.lua_type(L, idx) != T_STRING) return BridgeError.BadHandlerReturn;
    var len: usize = 0;
    const ptr = c.lua_tolstring(L, idx, &len);
    if (ptr == null) return BridgeError.BadHandlerReturn;
    return allocator.dupe(u8, ptr[0..len]) catch BridgeError.OutOfMemory;
}

// ---------------------------------------------------------------------------
// Lua-callable C functions
// ---------------------------------------------------------------------------

/// Implementation of `panto.register_tool { name=, description=, schema=, handler= }`.
///
/// Expects a single table argument with the four named fields. Validates
/// each field type, serializes `schema` to JSON, and appends a record to
/// the registrations table at `registry[&registrations_key]`. Throws a Lua
/// error via `luaL_error` if anything is malformed — that propagates out
/// of the running script as a Lua exception, which our `loadFile` surfaces
/// as `LuaRunFailed`.
fn registerToolThunk(L_opt: ?*c.lua_State) callconv(.c) c_int {
    const L = L_opt.?;
    c.luaL_checktype(L, 1, T_TABLE);

    // Pull each named field onto the stack and type-check it. After this
    // block, the stack layout is:
    //   1: args table (input)
    //   2: name      (string)
    //   3: description (string)
    //   4: schema      (table)
    //   5: handler     (function)
    expectField(L, 1, "name", T_STRING);
    expectField(L, 1, "description", T_STRING);
    expectField(L, 1, "schema", T_TABLE);
    expectField(L, 1, "handler", T_FUNCTION);

    // Serialize the schema table to JSON, leaving the string on top.
    pushSchemaAsJson(L, 4) catch |err| {
        const msg = switch (err) {
            BridgeError.OutOfMemory => "register_tool: out of memory serializing schema",
            else => "register_tool: schema is not JSON-serializable",
        };
        _ = c.luaL_error(L, msg);
        unreachable;
    };
    // Stack now: args(1), name(2), desc(3), schema(4), handler(5), schema_json(6)

    // Build the record table.
    c.lua_createtable(L, 0, 4);
    c.lua_pushvalue(L, 2);
    c.lua_setfield(L, -2, "name");
    c.lua_pushvalue(L, 3);
    c.lua_setfield(L, -2, "description");
    c.lua_pushvalue(L, 6); // schema_json
    c.lua_setfield(L, -2, "schema_json");
    c.lua_pushvalue(L, 5); // handler
    c.lua_setfield(L, -2, "handler");

    // Append the record to the registrations table.
    _ = c.lua_rawgetp(L, LUA_REGISTRYINDEX, &registrations_key);
    // Stack: ..., record, regs_table
    const n: c_int = @intCast(c.lua_rawlen(L, -1));
    c.lua_pushvalue(L, -2); // copy record above regs_table
    c.lua_rawseti(L, -2, n + 1);
    c.lua_settop(L, c.lua_gettop(L) - 1); // pop regs_table

    return 0;
}

/// Push `args_table[field_name]` onto the stack and assert it has the
/// expected type. Raises a Lua error if missing or wrong type.
fn expectField(
    L: *c.lua_State,
    args_idx: c_int,
    comptime field_name: [:0]const u8,
    expected_type: c_int,
) void {
    _ = c.lua_getfield(L, args_idx, field_name.ptr);
    const got = c.lua_type(L, -1);
    if (got != expected_type) {
        _ = c.luaL_error(
            L,
            "register_tool: field '%s' must be %s (got %s)",
            field_name.ptr,
            c.lua_typename(L, expected_type),
            c.lua_typename(L, got),
        );
        unreachable;
    }
}

// ---------------------------------------------------------------------------
// JSON <-> Lua conversion
// ---------------------------------------------------------------------------

/// Serialize the Lua table at stack index `idx` to a JSON string and push
/// that string onto the stack.
fn pushSchemaAsJson(L: *c.lua_State, idx: c_int) BridgeError!void {
    var aw: std.Io.Writer.Allocating = .init(std.heap.c_allocator);
    defer aw.deinit();

    var s: std.json.Stringify = .{ .writer = &aw.writer };
    writeLuaValueAsJson(L, idx, &s) catch return BridgeError.OutOfMemory;

    const bytes = aw.written();
    _ = c.lua_pushlstring(L, bytes.ptr, bytes.len);
}

/// Serialize the value at stack index `idx` to JSON via `stringifier`. We
/// allow strings, numbers, booleans, nil (→ JSON null), arrays (Lua tables
/// with sequential integer keys 1..n), and objects (any other table).
fn writeLuaValueAsJson(L: *c.lua_State, idx: c_int, w: *std.json.Stringify) anyerror!void {
    switch (c.lua_type(L, idx)) {
        T_NIL => try w.write(null),
        T_BOOLEAN => try w.write(c.lua_toboolean(L, idx) != 0),
        T_NUMBER => {
            // Distinguish integer vs float so we emit clean integers.
            var isnum: c_int = 0;
            const as_int = c.lua_tointegerx(L, idx, &isnum);
            if (isnum != 0) {
                try w.write(as_int);
            } else {
                const as_float = c.lua_tonumberx(L, idx, null);
                try w.write(as_float);
            }
        },
        T_STRING => {
            var len: usize = 0;
            const ptr = c.lua_tolstring(L, idx, &len);
            try w.write(ptr[0..len]);
        },
        T_TABLE => try writeLuaTableAsJson(L, idx, w),
        else => return error.UnsupportedLuaType,
    }
}

/// Decide whether the table at `idx` is JSON-array-shaped or object-shaped
/// and serialize accordingly.
///
/// Heuristic: if `lua_rawlen > 0`, treat it as an array (Lua's standard
/// length operator returns the array-part border, so this catches the
/// common case of `{ "a", "b", "c" }`). Otherwise iterate via `lua_next`
/// and emit a JSON object keyed by stringified keys.
///
/// An empty Lua table is ambiguous (could be either) and we serialize it
/// as `{}`, since JSON Schema usage almost always wants empty-object
/// shape (e.g. `properties = {}`).
fn writeLuaTableAsJson(L: *c.lua_State, idx_in: c_int, w: *std.json.Stringify) !void {
    // Normalize negative indices since we'll be pushing more on the stack.
    const abs_idx = c.lua_absindex(L, idx_in);

    const len = c.lua_rawlen(L, abs_idx);
    if (len > 0) {
        try w.beginArray();
        var i: c.lua_Integer = 1;
        while (i <= @as(c.lua_Integer, @intCast(len))) : (i += 1) {
            _ = c.lua_rawgeti(L, abs_idx, i);
            try writeLuaValueAsJson(L, -1, w);
            c.lua_settop(L, c.lua_gettop(L) - 1);
        }
        try w.endArray();
        return;
    }

    try w.beginObject();
    c.lua_pushnil(L); // first key
    while (c.lua_next(L, abs_idx) != 0) {
        // Stack: ..., key, value. Key must be a string for JSON objects.
        if (c.lua_type(L, -2) != T_STRING) {
            // Pop value, leave key for next lua_next iteration.
            c.lua_settop(L, c.lua_gettop(L) - 1);
            return error.UnsupportedLuaKey;
        }
        var klen: usize = 0;
        // CAREFUL: lua_tolstring on a non-string-key would mutate the key
        // and break lua_next. We already verified it's T_STRING above.
        const kptr = c.lua_tolstring(L, -2, &klen);
        try w.objectField(kptr[0..klen]);
        try writeLuaValueAsJson(L, -1, w);
        c.lua_settop(L, c.lua_gettop(L) - 1); // pop value, keep key
    }
    try w.endObject();
}

/// Push a parsed `std.json.Value` onto the stack.
fn pushJsonValue(L: *c.lua_State, v: std.json.Value) !void {
    switch (v) {
        .null => c.lua_pushnil(L),
        .bool => |b| c.lua_pushboolean(L, if (b) 1 else 0),
        .integer => |i| c.lua_pushinteger(L, @intCast(i)),
        .float => |f| c.lua_pushnumber(L, f),
        .number_string => |s| {
            // Best effort: try integer, else float.
            if (std.fmt.parseInt(c.lua_Integer, s, 10)) |i| {
                c.lua_pushinteger(L, i);
            } else |_| {
                const f = try std.fmt.parseFloat(c.lua_Number, s);
                c.lua_pushnumber(L, f);
            }
        },
        .string => |s| _ = c.lua_pushlstring(L, s.ptr, s.len),
        .array => |arr| {
            c.lua_createtable(L, @intCast(arr.items.len), 0);
            for (arr.items, 0..) |item, i| {
                try pushJsonValue(L, item);
                c.lua_rawseti(L, -2, @intCast(i + 1));
            }
        },
        .object => |obj| {
            c.lua_createtable(L, 0, @intCast(obj.count()));
            var it = obj.iterator();
            while (it.next()) |kv| {
                // Push key as Lua string (length-prefixed; doesn't need NUL).
                const key = kv.key_ptr.*;
                _ = c.lua_pushlstring(L, key.ptr, key.len);
                try pushJsonValue(L, kv.value_ptr.*);
                // Stack: ..., table, key, value -> table[key]=value, leaving table.
                c.lua_rawset(L, -3);
            }
        },
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Read a string-typed field `field_name` from the record at stack index
/// `record_idx`, duplicate its bytes into `arena`, and return the slice.
fn readStringField(
    L: *c.lua_State,
    record_idx: c_int,
    field_name: [:0]const u8,
    arena: Allocator,
) BridgeError![]const u8 {
    _ = c.lua_getfield(L, record_idx, field_name.ptr);
    defer c.lua_settop(L, c.lua_gettop(L) - 1);
    if (c.lua_type(L, -1) != T_STRING) return BridgeError.BadRegistration;
    var len: usize = 0;
    const ptr = c.lua_tolstring(L, -1, &len);
    if (ptr == null) return BridgeError.BadRegistration;
    return arena.dupe(u8, ptr[0..len]) catch BridgeError.OutOfMemory;
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

test "install creates panto.register_tool global" {
    const L = c.luaL_newstate() orelse return error.LuaInitFailed;
    defer c.lua_close(L);
    c.luaL_openlibs(L);
    install(L);

    _ = c.lua_getglobal(L, "panto");
    try std.testing.expectEqual(@as(c_int, T_TABLE), c.lua_type(L, -1));
    _ = c.lua_getfield(L, -1, "register_tool");
    try std.testing.expectEqual(@as(c_int, T_FUNCTION), c.lua_type(L, -1));
}

test "register_tool records name, description, schema_json" {
    const L = c.luaL_newstate() orelse return error.LuaInitFailed;
    defer c.lua_close(L);
    c.luaL_openlibs(L);
    install(L);

    const script =
        \\panto.register_tool {
        \\  name = "echo",
        \\  description = "Echoes its input back.",
        \\  schema = { type = "object", properties = { msg = { type = "string" } } },
        \\  handler = function(input) return input.msg end,
        \\}
    ;
    if (c.luaL_loadstring(L, script) != 0 or c.lua_pcallk(L, 0, 0, 0, 0, null) != 0) {
        var len: usize = 0;
        const msg = c.lua_tolstring(L, -1, &len);
        std.debug.print("lua error: {s}\n", .{msg[0..len]});
        return error.LuaScriptFailed;
    }

    var arena_state = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena_state.deinit();
    const regs = try harvestRegistrations(L, arena_state.allocator());

    try std.testing.expectEqual(@as(usize, 1), regs.len);
    try std.testing.expectEqualStrings("echo", regs[0].name);
    try std.testing.expectEqualStrings("Echoes its input back.", regs[0].description);
    // schema_json should be valid JSON containing "type": "object"
    try std.testing.expect(std.mem.indexOf(u8, regs[0].schema_json, "\"type\"") != null);
    try std.testing.expect(std.mem.indexOf(u8, regs[0].schema_json, "\"object\"") != null);
}

test "handler invocation: input parsed, result captured" {
    const L = c.luaL_newstate() orelse return error.LuaInitFailed;
    defer c.lua_close(L);
    c.luaL_openlibs(L);
    install(L);

    const script =
        \\panto.register_tool {
        \\  name = "echo", description = "echoes",
        \\  schema = { type = "object" },
        \\  handler = function(input) return "got: " .. input.msg end,
        \\}
    ;
    if (c.luaL_loadstring(L, script) != 0 or c.lua_pcallk(L, 0, 0, 0, 0, null) != 0) {
        return error.LuaScriptFailed;
    }

    try pushHandler(L, "echo");
    try std.testing.expectEqual(@as(c_int, T_FUNCTION), c.lua_type(L, -1));

    var arena_state = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena_state.deinit();
    try pushJsonAsLua(L, arena_state.allocator(), "{\"msg\":\"hello\"}");

    // Call: 1 arg, 1 return.
    if (c.lua_pcallk(L, 1, 1, 0, 0, null) != 0) {
        return error.LuaCallFailed;
    }

    const result = try readHandlerResult(L, -1, std.testing.allocator);
    defer std.testing.allocator.free(result);
    try std.testing.expectEqualStrings("got: hello", result);
}

test "handler crash: error message surfaces via xpcall traceback hook" {
    const L = c.luaL_newstate() orelse return error.LuaInitFailed;
    defer c.lua_close(L);
    c.luaL_openlibs(L);
    install(L);

    const script =
        \\panto.register_tool {
        \\  name = "boom", description = "crashes",
        \\  schema = { type = "object" },
        \\  handler = function(input) error("explosion") end,
        \\}
    ;
    if (c.luaL_loadstring(L, script) != 0 or c.lua_pcallk(L, 0, 0, 0, 0, null) != 0) {
        return error.LuaScriptFailed;
    }

    // Push a traceback error handler at the bottom of the call frame.
    _ = c.lua_getglobal(L, "debug");
    _ = c.lua_getfield(L, -1, "traceback");
    c.lua_copy(L, -1, -2);
    c.lua_settop(L, c.lua_gettop(L) - 1);
    const errfunc_idx = c.lua_gettop(L);

    try pushHandler(L, "boom");
    c.lua_pushnil(L); // input arg

    // pcallk with errfunc index = where we put debug.traceback.
    const rc = c.lua_pcallk(L, 1, 1, errfunc_idx, 0, null);
    try std.testing.expect(rc != 0);

    var len: usize = 0;
    const msg = c.lua_tolstring(L, -1, &len);
    try std.testing.expect(msg != null);
    const slice = msg[0..len];
    try std.testing.expect(std.mem.indexOf(u8, slice, "explosion") != null);
    // Should also contain a traceback marker since we used debug.traceback.
    try std.testing.expect(std.mem.indexOf(u8, slice, "stack traceback") != null);
}