econ/src/lua.zig

const std = @import("std");
const root = @import("root");

pub const c = @import("luajit").c;

const util = root.util;

pub const Func = fn (lua: ?*c.lua_State) callconv(.C) c_int;

const slice_types = .{
    f32,
    usize,
    root.model.Product,
    root.model.Product.Quantity,
};

const map_types = .{
    root.model.Task.ProductMap,
};

var load_alloc: ?std.mem.Allocator = null;

pub fn init(
    alloc: std.mem.Allocator,
    config: root.Config.Start,
    simulation: *root.Simulation,
) !*c.lua_State {
    load_alloc = alloc;

    const lua = c.luaL_newstate().?;
    _ = c.luaopen_base(lua);
    _ = c.luaopen_table(lua);
    _ = c.luaopen_io(lua);
    _ = c.luaopen_string(lua);
    _ = c.luaopen_math(lua);

    const cstr_filename = try alloc.dupeZ(u8, config.script);
    defer alloc.free(cstr_filename);

    if (c.luaL_loadfile(lua, cstr_filename.ptr) != 0) {
        const error_str = load([]const u8, lua);
        defer load_alloc.?.free(error_str);

        std.log.err("lua error: {s}", .{error_str});
        return error.LuaError;
    }

    initTables(simulation, lua);
    initAPI(simulation, lua);

    if (c.lua_pcall(lua, 0, c.LUA_MULTRET, 0) != 0) {
        const error_str = load([]const u8, lua);
        defer load_alloc.?.free(error_str);

        std.log.err("lua error: {s}", .{error_str});
        return error.LuaError;
    }

    return lua;
}

pub fn deinit(_: std.mem.Allocator, lua: *c.lua_State) !void {
    //try call(void, alloc, lua, "printResults", .{});
    load_alloc = null;
    c.lua_close(lua);
}

pub fn update(
    alloc: std.mem.Allocator,
    _: root.Config.Run,
    lua: *c.lua_State,
) !void {
    const api = root.Simulation.interface;
    const tuples = .{
        .{ api.ResourcesOrder, "assignResources", api.addResources },
        .{ api.WageOrder, "assignWages", api.setWages },
        .{ api.PriceOrder, "assignPrices", api.setPrices },
    };
    inline for (tuples) |tuple| {
        const Type = tuple.@"0";
        const funcName = tuple.@"1";
        const apiFunc = tuple.@"2";
        const data = try call([]Type, alloc, lua, funcName, .{});
        defer util.free(alloc, data);

        apiFunc(data);
    }
}

pub fn call(
    comptime Type: type,
    alloc: std.mem.Allocator,
    lua: *c.lua_State,
    name: []const u8,
    args: anytype,
) !Type {
    const cstr = try alloc.dupeZ(u8, name);
    defer alloc.free(cstr);

    c.lua_getglobal(lua, cstr.ptr);

    for (args) |arg| {
        push(lua, arg);
    }

    const err = c.lua_pcall(lua, args.len, 1, 0);
    if (err != 0) {
        const msg = load([]const u8, lua);
        defer load_alloc.?.free(msg);

        std.log.info("error running {s}: {s}", .{ name, msg });
    }

    if (Type != void) {
        return load(Type, lua);
    }
}

fn initTables(_: *root.Simulation, lua: *c.lua_State) void {
    inline for (slice_types) |Type| {
        addSliceMetatable(Type, lua);
    }
    inline for (map_types) |Type| {
        addMapMetatable(Type, lua);
    }
}

fn initAPI(simulation: *root.Simulation, lua: *c.lua_State) void {
    const api = root.Simulation.interface;
    api.simulation = simulation;

    c.lua_newtable(lua);

    setField(lua, "products", simulation.products);
    //setField(lua, "agents", simulation.agents.keys());
    //setField(lua, "tasks", simulation.tasks.keys());

    registerFunc(lua, "getDay", api.getDay);
    registerFunc(lua, "getAvgLabor", api.getAvgLabor);
    registerFunc(lua, "getTasks", api.getTasks);
    registerFunc(lua, "getAgents", api.getAgents);
    registerFunc(lua, "getTaskProduct", api.getTaskProduct);
    registerFunc(lua, "getInventory", api.getInventory);
    registerFunc(lua, "getProduced", api.getProduced);
    registerFunc(lua, "getConsumed", api.getConsumed);
    registerFunc(lua, "getConsumedByAgents", api.getConsumedByAgents);
    registerFunc(lua, "getConsumedByTasks", api.getConsumedByTasks);
    registerFunc(lua, "getAgentTask", api.getAgentTask);
    registerFunc(lua, "getTaskAgents", api.getTaskAgents);
    registerFunc(lua, "getAgentConsumption", api.getAgentConsumption);
    registerFunc(lua, "getAgentProduction", api.getAgentProduction);
    registerFunc(lua, "getAgentSavings", api.getAgentSavings);
    registerFunc(lua, "getTaskInventory", api.getTaskInventory);
    registerFunc(lua, "getTaskUnspent", api.getTaskUnspent);
    registerFunc(lua, "getAverageHealth", api.getAverageHealth);
    registerFunc(lua, "getAverageHappiness", api.getAverageHappiness);

    c.lua_setglobal(lua, "econ");
}

fn registerFunc(
    lua: *c.lua_State,
    comptime name: []const u8,
    comptime func: anytype,
) void {
    const wfunc = wrap(func);
    setFunctionField(lua, name, wfunc);
}

pub fn wrap(comptime func: anytype) Func {
    return struct {
        const Type = @TypeOf(func);
        const info = @typeInfo(Type).Fn;
        const Tuple = std.meta.ArgsTuple(Type);
        const fields = std.meta.fields(Tuple);
        const len = fields.len;
        const return_type = info.return_type;
        pub fn f(opt: ?*c.lua_State) callconv(.C) c_int {
            const lua = opt.?;
            var args: Tuple = undefined;
            inline for (fields, 0..) |field, i| {
                c.lua_pushvalue(lua, -@as(c_int, @intCast(len - i)));
                @field(args, field.name) = load(field.type, lua);
                c.lua_pop(lua, 1);
            }
            if (return_type != null and return_type.? != void) {
                const ret = @call(.auto, func, args);
                push(lua, ret);
                return 1;
            } else {
                @call(.auto, func, args);
                return 0;
            }
        }
    }.f;
}

pub fn load(comptime Type: type, lua: *c.lua_State) Type {
    //@compileLog("load " ++ @typeName(Type));
    //std.log.info("loading {s}", .{@typeName(Type)});
    const alloc = load_alloc.?;
    const value: Type = switch (@typeInfo(Type)) {
        .Bool => c.lua_toboolean(lua, -1) != 0,
        .Int => @as(Type, @intFromFloat(c.lua_tonumber(lua, -1))),
        .Float => @as(Type, @floatCast(c.lua_tonumber(lua, -1))),
        .Optional => |info| if (c.lua_isnil(lua, -1))
            null
        else
            load(info.child, lua),
        .Struct => loadStruct(Type, lua),
        .Pointer => |info| switch (info.size) {
            .Slice => if (info.is_const and info.child == u8) blk: {
                var len: usize = 0;
                const ptr = c.lua_tolstring(lua, -1, &len);
                var slice: [:0]const u8 = undefined;
                slice.ptr = ptr;
                slice.len = len;
                break :blk alloc.dupe(u8, slice) catch
                    unreachable;
            } else loadSlice(info.child, alloc, lua) catch
                unreachable,
            else => @compileError("unsupported type"),
        },
        else => @compileError("unsupported type"),
    };
    if (comptime @typeInfo(Type) != .Optional) {
        c.lua_pop(lua, 1);
    }
    return value;
}

fn loadStruct(comptime Type: type, lua: *c.lua_State) Type {
    var value: Type = undefined; //TODO: init defaults
    inline for (std.meta.fields(Type)) |field| {
        const name = field.name;
        @field(value, name) = getField(field.type, lua, name);
    }
    return value;
}

fn getField(comptime Type: type, lua: *c.lua_State, name: []const u8) Type {
    push(lua, name);
    c.lua_gettable(lua, -2);
    return load(Type, lua);
}

fn loadSlice(
    comptime Type: type,
    alloc: std.mem.Allocator,
    lua: *c.lua_State,
) ![]Type {
    const len = c.lua_objlen(lua, -1);
    const slice = try alloc.alloc(Type, len);
    for (slice, 0..) |*item, i| {
        item.* = getIndex(Type, lua, i + 1);
    }
    return slice;
}

fn getIndex(comptime Type: type, lua: *c.lua_State, index: usize) Type {
    push(lua, index);
    c.lua_gettable(lua, -2);
    return load(Type, lua);
}

pub fn push(lua: *c.lua_State, value: anytype) void {
    const Type = @TypeOf(value);
    if (comptime in(Type, map_types)) {
        pushMap(Type, lua, value);
        return;
    } else {
        switch (@typeInfo(Type)) {
            .Bool => c.lua_pushboolean(lua, @intFromBool(value)),
            .Int => c.lua_pushnumber(lua, @as(f64, @floatFromInt(value))),
            .Float => c.lua_pushnumber(lua, @as(f64, @floatCast(value))),
            .Optional => if (value) |v| push(lua, v) else c.lua_pushnil(lua),
            .Array => |info| pushSliceAsTable(info.child, lua, &value),
            .Struct => pushStruct(Type, lua, value),
            .Pointer => |info| switch (info.size) {
                .Slice => if (info.is_const and info.child == u8) {
                    c.lua_pushlstring(lua, value.ptr, value.len);
                } else {
                    if (in(info.child, slice_types)) {
                        pushSlice(info.child, lua, value);
                    } else {
                        pushSliceAsTable(info.child, lua, value);
                    }
                },
                .One => push(lua, value.*),
                else => @compileError("unsupported type"),
            },
            else => @compileError("unsupported type"),
        }
    }
}

fn pushStruct(comptime Type: type, lua: *c.lua_State, value: Type) void {
    c.lua_newtable(lua);
    inline for (std.meta.fields(Type)) |field| {
        const name = field.name;
        setField(lua, name, @field(value, name));
    }
}

fn pushMap(comptime Type: type, lua: *c.lua_State, map: Type) void {
    const Payload = Type;
    const n = @sizeOf(Payload);
    const ptr = cast(*Payload, c.lua_newuserdata(lua, n));
    c.luaL_getmetatable(lua, mapMetatableName(Type));
    _ = c.lua_setmetatable(lua, -2);
    ptr.* = map;
}

fn pushSlice(comptime Type: type, lua: *c.lua_State, slice: []Type) void {
    const Payload = []Type;
    const n = @sizeOf(Payload);
    const ptr = cast(*Payload, c.lua_newuserdata(lua, n));
    c.luaL_getmetatable(lua, sliceMetatableName(Type));
    _ = c.lua_setmetatable(lua, -2);
    ptr.* = slice;
}

fn pushSliceAsTable(
    comptime Type: type,
    lua: *c.lua_State,
    value: []Type,
) void {
    c.lua_newtable(lua);
    for (value, 0..) |item, i| {
        push(lua, i);
        push(lua, item);
        c.lua_settable(lua, -3);
    }
    //setField(lua, "size", value.len);
}

fn setField(lua: *c.lua_State, index: []const u8, value: anytype) void {
    push(lua, index);
    push(lua, value);
    c.lua_settable(lua, -3);
}

fn setFunctionField(
    lua: *c.lua_State,
    index: []const u8,
    value: *const Func,
) void {
    push(lua, index);
    c.lua_pushcfunction(lua, value);
    c.lua_settable(lua, -3);
}

fn addMapMetatable(comptime Type: type, lua: *c.lua_State) void {
    const Payload = Type;
    const getItem = struct {
        pub fn f(opt_l: ?*c.lua_State) callconv(.C) c_int {
            const l = opt_l.?;
            const opt = cast(?*Payload, c.lua_touserdata(l, 1));
            const index = c.luaL_checkint(l, 2);
            const ptr = opt.?;
            push(l, ptr.get(@as(usize, @intCast(index))));
            return 1;
        }
    }.f;
    const getSize = struct {
        pub fn f(opt_l: ?*c.lua_State) callconv(.C) c_int {
            const l = opt_l.?;
            const opt = cast(?*Payload, c.lua_touserdata(l, 1));
            const ptr = opt.?;
            push(l, ptr.count());
            return 1;
        }
    }.f;
    const index_str: []const u8 = "__index";

    _ = c.luaL_newmetatable(lua, mapMetatableName(Type));
    setFunctionField(lua, "get", getItem);
    setFunctionField(lua, "size", getSize);
    push(lua, index_str);
    c.lua_pushvalue(lua, -2);
    c.lua_settable(lua, -3);
    c.lua_pop(lua, 1);
}

fn mapMetatableName(comptime Type: type) [*c]const u8 {
    return "econ_map" ++ @typeName(Type);
}

fn addSliceMetatable(comptime Type: type, lua: *c.lua_State) void {
    const Payload = []Type;
    const getItem = struct {
        pub fn f(opt_l: ?*c.lua_State) callconv(.C) c_int {
            const l = opt_l.?;
            const opt = cast(?*Payload, c.lua_touserdata(l, 1));
            const index = c.luaL_checkint(l, 2);
            const ptr = opt.?;
            push(l, ptr.*[@as(usize, @intCast(index))]);
            return 1;
        }
    }.f;
    const getSize = struct {
        pub fn f(opt_l: ?*c.lua_State) callconv(.C) c_int {
            const l = opt_l.?;
            const opt = cast(?*Payload, c.lua_touserdata(l, 1));
            const ptr = opt.?;
            push(l, ptr.*.len);
            return 1;
        }
    }.f;
    const index_str: []const u8 = "__index";

    _ = c.luaL_newmetatable(lua, sliceMetatableName(Type));
    setFunctionField(lua, "get", getItem);
    setFunctionField(lua, "size", getSize);
    push(lua, index_str);
    c.lua_pushvalue(lua, -2);
    c.lua_settable(lua, -3);
    c.lua_pop(lua, 1);
}

fn sliceMetatableName(comptime Type: type) [*c]const u8 {
    return "econ_slice" ++ @typeName(Type);
}

fn cast(comptime Type: type, ptr: anytype) Type {
    return @as(Type, @ptrFromInt(@intFromPtr(ptr)));
}

fn in(comptime Type: type, comptime tuple: anytype) bool {
    var found: bool = false;
    inline for (tuple) |Item| {
        if (Item == Type) {
            found = true;
        }
    }
    return found;
}