from pathlib import Path

root = Path(__file__).resolve().parent
main_path = root / "main.zig"
build_path = root / "build.zig"

main = main_path.read_text(encoding="utf-8")
build = build_path.read_text(encoding="utf-8")

replacements = [
    (
        "const BATCH_HEADROOM: u32 = 128 * 1024;\n",
        "const BATCH_HEADROOM: u32 = 128 * 1024;\n"
        "const NATIVE_GRACE_COLLECTIONS: u32 = 3;\n",
    ),
    (
        "var installed = false;\nvar in_gc = false;\nvar closing_lua = false;\n",
        "var installed = false;\nvar in_gc = false;\nvar closing_lua = false;\n"
        "var active_g: u32 = 0;\n"
        "var native_grace_collections: u32 = 0;\n",
    ),
    (
        "    // Never carry private list state into another Lua global_State.\n"
        "    if (sweeping and g != saved_g) {\n"
        "        resetSweepState();\n"
        "        collect_hook.callOriginal(.{L});\n"
        "        return;\n"
        "    }\n\n"
        "    // If another module changed the birth byte during our split cycle,\n",
        "    // Never carry private list state into another Lua global_State.\n"
        "    if (sweeping and g != saved_g) {\n"
        "        resetSweepState();\n"
        "        collect_hook.callOriginal(.{L});\n"
        "        return;\n"
        "    }\n\n"
        "    // B2: when a fresh Lua global_State appears (initial world entry,\n"
        "    // /reload, logout/relog), keep the first three collections fully\n"
        "    // native. This gives WoW ownership of the fragile state-rebuild\n"
        "    // window before incremental rootgc sweeping resumes.\n"
        "    if (active_g != g) {\n"
        "        active_g = g;\n"
        "        native_grace_collections = NATIVE_GRACE_COLLECTIONS;\n"
        "    }\n"
        "    if (native_grace_collections != 0) {\n"
        "        native_grace_collections -= 1;\n"
        "        collect_hook.callOriginal(.{L});\n"
        "        return;\n"
        "    }\n\n"
        "    // If another module changed the birth byte during our split cycle,\n",
    ),
    (
        "    reconnectSweep();\n    in_gc = false;\n\n    lua_close_hook.callOriginal(.{L});\n",
        "    reconnectSweep();\n    in_gc = false;\n"
        "    active_g = 0;\n"
        "    native_grace_collections = 0;\n\n"
        "    lua_close_hook.callOriginal(.{L});\n",
    ),
    (
        'const version: [*:0]const u8 = "2.4-B1-gc-safe-sweep-abi";\n',
        'const version: [*:0]const u8 = "2.4-B2-gc-safe-sweep-reload-guard";\n',
    ),
]

for old, new in replacements:
    count = main.count(old)
    if count != 1:
        raise SystemExit(f"main.zig patch anchor mismatch: expected 1, found {count}: {old[:80]!r}")
    main = main.replace(old, new, 1)

old_name = '.name = "weirdperformance_gc24b1",'
new_name = '.name = "weirdperformance_gc24b2",'
if build.count(old_name) != 1:
    raise SystemExit("build.zig output-name anchor mismatch")
build = build.replace(old_name, new_name, 1)

main_path.write_text(main, encoding="utf-8")
build_path.write_text(build, encoding="utf-8")

print("B2 patch applied: new-global_State native grace = 3 GC cycles")