//! WeirdPerformance 2.4-B2 GC in-place safe-sweep companion.
//!
//! A/B contract:
//!   A  = validated WeirdPerformance 2.4-A binary, unchanged.
//!   B2 = the exact same 2.4-A binary + this companion DLL.
//!
//! B1 split rootgc by detaching swept survivors from the live rootgc chain
//! between GC calls. B2 keeps the complete rootgc chain linked whenever
//! control returns to WoW. Each chunk is isolated only for the duration of
//! lua_gc_remove_objects(), then reconnected before the detour returns.
//!
//! This preserves B1's bounded 50,000-object rootgc sweep while removing the
//! long-lived fragmented-list state suspected in delayed lua_table_set_value
//! corruption (0x006FA8E2).
//!
//! Target: WoW 1.12.1 build 5875, x86 only.

const std = @import("std");
const hook = @import("zhook");
const WINAPI = std.builtin.CallingConvention.winapi;
const X86_FASTCALL: std.builtin.CallingConvention = .{ .x86_fastcall = .{} };

const LUA_COLLECT_GARBAGE_ADDR: usize = 0x6F7340;
const LUA_CLOSE_ADDR: usize = 0x6F6EF0;
const IS_IN_WORLD_ADDR: u32 = 0x00B4B424;

const GS_ROOTGC: u32 = 0x10;
const GS_ROOTUDATA: u32 = 0x14;
const GS_GCTHRESHOLD: u32 = 0x24;
const GS_TOTALBYTES: u32 = 0x28;
const OBJ_NEXT: u32 = 0;

const CHUNK_SIZE: u32 = 50_000;
const BATCH_HEADROOM: u32 = 128 * 1024;

// Crash logs for the supported client show 0x00400000..0x00D2B000.
const EXPECTED_IMAGE_BASE: usize = 0x00400000;
const EXPECTED_IMAGE_SIZE: u32 = 0x0092B000;

const CollectFn = fn (u32) callconv(X86_FASTCALL) void;
const LuaCloseFn = fn (u32) callconv(X86_FASTCALL) void;
const SweepAllFn = fn (u32, u32) callconv(X86_FASTCALL) void;
const RemoveObjectsFn = fn (u32, u32, u32) callconv(X86_FASTCALL) u32;

const lua_gc_full_collection: *const CollectFn = @ptrFromInt(0x6F73E0);
const lua_gc_shrink_memory: *const CollectFn = @ptrFromInt(0x6F7370);
const luaCallUserDataGC: *const CollectFn = @ptrFromInt(0x6F7080);
const lua_gc_sweep_all_lists: *const SweepAllFn = @ptrFromInt(0x6F72F0);
const lua_gc_remove_objects: *const RemoveObjectsFn = @ptrFromInt(0x6F7210);

var collect_hook: hook.Detour(CollectFn) = .{};
var lua_close_hook: hook.Detour(LuaCloseFn) = .{};

var installed = false;
var in_gc = false;
var closing_lua = false;

// B2 state. Unlike B1 there are no detached rootgc fragments.
// sweep_next is the first object not yet swept in the current native mark cycle.
// New luaC_link objects are prepended before it and are intentionally left for
// the next cycle rather than being exposed to the current sweep without a mark.
var sweeping = false;
var sweep_next: u32 = 0;
var saved_g: u32 = 0;

inline fn readU16(addr: usize) u16 {
    return @as(*volatile const u16, @ptrFromInt(addr)).*;
}

inline fn readU32(addr: u32) u32 {
    return @as(*volatile const u32, @ptrFromInt(addr)).*;
}

inline fn readU32usize(addr: usize) u32 {
    return @as(*volatile const u32, @ptrFromInt(addr)).*;
}

inline fn writeU32(addr: u32, value: u32) void {
    @as(*volatile u32, @ptrFromInt(addr)).* = value;
}

inline fn getGlobalState(L: u32) u32 {
    return readU32(L + 0x10);
}

fn validateClient() bool {
    if (readU16(EXPECTED_IMAGE_BASE) != 0x5A4D) return false; // MZ

    const pe_off = readU32usize(EXPECTED_IMAGE_BASE + 0x3C);
    const pe = EXPECTED_IMAGE_BASE + pe_off;
    if (readU32usize(pe) != 0x00004550) return false; // PE\0\0
    if (readU16(pe + 4) != 0x014C) return false; // IMAGE_FILE_MACHINE_I386
    if (readU16(pe + 24) != 0x010B) return false; // PE32 optional header

    const image_size = readU32usize(pe + 24 + 56);
    return image_size == EXPECTED_IMAGE_SIZE;
}

fn resetSweepState() void {
    sweeping = false;
    sweep_next = 0;
    saved_g = 0;
}

// Locate the pointer field which currently references target. This is required
// because luaC_link prepends newly allocated objects to g->rootgc while a split
// sweep is in progress. We must skip those new objects instead of accidentally
// sweeping them with marks from the previous atomic phase.
fn findLinkTo(g: u32, target: u32) ?u32 {
    if (target == 0) return null;

    var link_addr = g + GS_ROOTGC;
    var obj = readU32(link_addr);
    while (obj != 0) {
        if (obj == target) return link_addr;
        link_addr = obj + OBJ_NEXT;
        obj = readU32(link_addr);
    }
    return null;
}

fn findChunkTail(head: u32, limit: u32) struct { tail: u32, count: u32 } {
    if (head == 0 or limit == 0) return .{ .tail = 0, .count = 0 };

    var obj = head;
    var count: u32 = 1;
    while (count < limit) : (count += 1) {
        const next = readU32(obj + OBJ_NEXT);
        if (next == 0) return .{ .tail = obj, .count = count };
        obj = next;
    }
    return .{ .tail = obj, .count = count };
}

fn findTail(head: u32) u32 {
    var obj = head;
    if (obj == 0) return 0;
    while (true) {
        const next = readU32(obj + OBJ_NEXT);
        if (next == 0) return obj;
        obj = next;
    }
}

const ChunkResult = enum {
    more,
    done,
    invalid,
};

// Sweep exactly one bounded sub-list using WoW's native sweep routine.
// The rootgc chain may be temporarily truncated while the native routine runs,
// but it is always fully reconnected before this function returns.
fn sweepOneChunk(L: u32, g: u32) ChunkResult {
    if (!sweeping or sweep_next == 0) return .done;
    if (g != saved_g) return .invalid;

    const link_addr = findLinkTo(g, sweep_next) orelse return .invalid;
    const head = readU32(link_addr);
    if (head != sweep_next) return .invalid;

    const chunk = findChunkTail(head, CHUNK_SIZE);
    if (chunk.tail == 0) return .done;

    const rest = readU32(chunk.tail + OBJ_NEXT);

    // Final chunk: no temporary split is necessary. Let the native routine
    // update the real list link directly and complete the cycle.
    if (rest == 0) {
        _ = lua_gc_remove_objects(L, link_addr, 0);
        sweep_next = 0;
        return .done;
    }

    // Isolate only the current chunk for the native sweep call.
    writeU32(chunk.tail + OBJ_NEXT, 0);
    _ = lua_gc_remove_objects(L, link_addr, 0);

    // Reconnect the untouched remainder before returning to gameplay.
    // If every object in the chunk died, link_addr itself becomes the bridge.
    const survivors = readU32(link_addr);
    if (survivors == 0) {
        writeU32(link_addr, rest);
    } else {
        const survivor_tail = findTail(survivors);
        if (survivor_tail == 0) return .invalid;
        writeU32(survivor_tail + OBJ_NEXT, rest);
    }

    sweep_next = rest;
    return .more;
}

fn finishCycle(L: u32) void {
    resetSweepState();
    lua_gc_shrink_memory(L);
    luaCallUserDataGC(L);
}

// Used only for transitions such as leaving the world while a B2 cycle is
// active. The list is fully linked, so we can finish remaining chunks without
// any B1-style fragment reconstruction.
fn finishSweepNow(L: u32, g: u32) bool {
    while (sweeping) {
        switch (sweepOneChunk(L, g)) {
            .more => {},
            .done => {
                finishCycle(L);
                return true;
            },
            .invalid => {
                resetSweepState();
                return false;
            },
        }
    }
    return true;
}

fn nativeFallback(L: u32) void {
    resetSweepState();
    collect_hook.callOriginal(.{L});
}

fn collectGarbageDetour(L: u32) callconv(X86_FASTCALL) void {
    if (closing_lua) {
        collect_hook.callOriginal(.{L});
        return;
    }

    if (in_gc) return;
    if (L == 0) return;
    if (readU32(L + 0x60) == 0) return;

    in_gc = true;
    defer in_gc = false;

    const g = getGlobalState(L);
    if (g == 0) {
        collect_hook.callOriginal(.{L});
        return;
    }

    // If gameplay ends mid-cycle, finish the already-marked sweep while the
    // current Lua state is still valid, then return control to native GC.
    if (readU32(IS_IN_WORLD_ADDR) == 0) {
        if (sweeping and g == saved_g) {
            if (!finishSweepNow(L, g)) {
                collect_hook.callOriginal(.{L});
                return;
            }
        } else if (sweeping) {
            resetSweepState();
        }
        collect_hook.callOriginal(.{L});
        return;
    }

    // Never carry a cursor into a different Lua global_State.
    if (sweeping and g != saved_g) {
        nativeFallback(L);
        return;
    }

    if (!sweeping) {
        // Keep WoW's native atomic mark, userdata sweep, and string sweep.
        lua_gc_full_collection(L);
        _ = lua_gc_remove_objects(L, g + GS_ROOTUDATA, 0);
        lua_gc_sweep_all_lists(L, 0);

        sweep_next = readU32(g + GS_ROOTGC);
        saved_g = g;
        sweeping = sweep_next != 0;

        if (!sweeping) {
            finishCycle(L);
            return;
        }
    }

    switch (sweepOneChunk(L, g)) {
        .done => {
            finishCycle(L);
            return;
        },
        .invalid => {
            // The live rootgc chain itself was never fragmented across calls,
            // so falling back to WoW's complete collector is safe here.
            nativeFallback(L);
            return;
        },
        .more => {
            const totalbytes = readU32(g + GS_TOTALBYTES);
            writeU32(g + GS_GCTHRESHOLD, totalbytes + BATCH_HEADROOM);
            return;
        },
    }
}

fn luaCloseDetour(L: u32) callconv(X86_FASTCALL) void {
    // B2 never leaves detached rootgc fragments. lua_close can therefore own
    // teardown normally; just discard the cursor before the state is destroyed.
    closing_lua = true;
    resetSweepState();
    in_gc = false;

    lua_close_hook.callOriginal(.{L});

    closing_lua = false;
}

fn install() void {
    if (installed) return;
    if (!validateClient()) return;

    // Transactional install: lua_close guard first, collector second.
    if (lua_close_hook.attach(LUA_CLOSE_ADDR, &luaCloseDetour) != .ok) return;

    if (collect_hook.attach(LUA_COLLECT_GARBAGE_ADDR, &collectGarbageDetour) != .ok) {
        lua_close_hook.detach();
        return;
    }

    installed = true;
}

const version: [*:0]const u8 = "2.4-B2-gc-inplace-safe-sweep";

pub export fn WeirdPerformanceGC24B2_GetVersion() callconv(.c) [*:0]const u8 {
    return version;
}

pub export fn WeirdPerformanceGC24B2_IsActive() callconv(.c) i32 {
    return if (installed) 1 else 0;
}

pub export fn DllMain(
    _: ?*anyopaque,
    reason: u32,
    _: ?*anyopaque,
) callconv(WINAPI) std.os.windows.BOOL {
    if (reason == 1) install();

    // Process-lifetime A/B companion. We intentionally do not hot-unhook
    // detours during process teardown.
    return @enumFromInt(1);
}