// SPDX-FileCopyrightText: © 2026 Vladimir Zorin <vladimir@deviant.guru>
// SPDX-License-Identifier: LicenseRef-OWL-1.0-or-later
// Licensed under OWL v1.0+. See LICENSE.

#include "mneme.h"

#include <lauxlib.h>
#include <lua.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>

#define MNEME_DB_MT  "mneme.db"
#define MNEME_TXN_MT "mneme.txn"
#define MNEME_CUR_MT "mneme.cursor"
#define MNEME_DEK_MT "mneme.dek"

/* ── DB userdata ───────────────────────────────────────────────────── */

typedef struct {
    mneme_db_t *db;
} lua_mneme_db_t;

static lua_mneme_db_t *check_db(lua_State *L) {
    lua_mneme_db_t *ud = (lua_mneme_db_t *)luaL_checkudata(L, 1, MNEME_DB_MT);
    if (!ud->db || ud->db->closed)
        luaL_error(L, "db is closed");
    return ud;
}

/* ── TXN userdata ──────────────────────────────────────────────────── */

typedef struct {
    mneme_txn_t *txn;
} lua_mneme_txn_t;

static lua_mneme_txn_t *check_txn(lua_State *L) {
    lua_mneme_txn_t *ud = (lua_mneme_txn_t *)luaL_checkudata(L, 1, MNEME_TXN_MT);
    if (!ud->txn || ud->txn->committed)
        luaL_error(L, "transaction is finished");
    return ud;
}

/* For bindings that mutate the database: also reject readonly txns,
   whose write-side state (dirty array etc.) was never allocated. */
static lua_mneme_txn_t *check_txn_rw(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn(L);
    if (ud->txn->readonly)
        luaL_error(L, "transaction is read-only");
    return ud;
}

/* Anchor the userdata at parent_idx in the environment table of the
   userdata at the stack top, so the parent (and the C struct it owns)
   cannot be collected while the child is alive. */
static void anchor_parent(lua_State *L, int parent_idx) {
    lua_createtable(L, 1, 0);
    lua_pushvalue(L, parent_idx);
    lua_rawseti(L, -2, 1);
    lua_setfenv(L, -2);
}

/* ── db_open(path, flags_table) ────────────────────────────────────── */

static int l_db_open(lua_State *L) {
    const char *path = luaL_checkstring(L, 1);
    uint32_t flags   = MNEME_OPEN_CREATE; /* default: create if missing */

    if (lua_istable(L, 2)) {
        lua_getfield(L, 2, "readonly");
        if (lua_toboolean(L, -1))
            flags |= MNEME_OPEN_READONLY;
        lua_pop(L, 1);

        lua_getfield(L, 2, "create");
        if (!lua_isnil(L, -1) && !lua_toboolean(L, -1))
            flags &= ~MNEME_OPEN_CREATE;
        lua_pop(L, 1);

        lua_getfield(L, 2, "sync");
        if (lua_isstring(L, -1)) {
            const char *s = lua_tostring(L, -1);
            if (strcmp(s, "none") == 0)
                flags |= MNEME_OPEN_NOSYNC;
        }
        lua_pop(L, 1);
    }

    /* Create the userdata before acquiring the C resource: lua_newuserdata
       can longjmp on OOM, which would otherwise leak the open db. */
    lua_mneme_db_t *ud = (lua_mneme_db_t *)lua_newuserdata(L, sizeof(*ud));
    ud->db             = NULL;
    luaL_getmetatable(L, MNEME_DB_MT);
    lua_setmetatable(L, -2);

    mneme_db_t *db = NULL;
    int rc         = mneme_db_open(path, flags, &db);
    if (rc < 0) {
        lua_pushnil(L);
        lua_pushstring(L, rc == -2 ? "unsupported database format version" : "failed to open database");
        return 2;
    }
    ud->db = db;
    return 1;
}

/* ── db:close() ────────────────────────────────────────────────────── */

static int l_db_close(lua_State *L) {
    lua_mneme_db_t *ud = (lua_mneme_db_t *)luaL_checkudata(L, 1, MNEME_DB_MT);
    if (ud->db && !ud->db->closed && ud->db->write_txn_active) {
        /* Closing now would free the struct out from under the live
           txn; its later commit/abort would be a use-after-free. */
        lua_pushnil(L);
        lua_pushstring(L, "write transaction active");
        return 2;
    }
    if (ud->db) {
        /* The handle may already be closed (poisoned by a failed mmap
           refresh after commit); the struct still has to be freed. */
        if (!ud->db->closed)
            mneme_db_close(ud->db);
        free(ud->db);
        ud->db = NULL;
    }
    lua_pushboolean(L, 1);
    return 1;
}

/* ── db:sync() ─────────────────────────────────────────────────────── */

static int l_db_sync(lua_State *L) {
    lua_mneme_db_t *ud = check_db(L);
    if (mneme_db_sync(ud->db) < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "sync failed");
        return 2;
    }
    lua_pushboolean(L, 1);
    return 1;
}

/* ── db:info() ─────────────────────────────────────────────────────── */

static int l_db_info(lua_State *L) {
    lua_mneme_db_t *ud = check_db(L);
    mneme_db_t *db     = ud->db;

    if (mneme_meta_refresh(db) < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "failed to refresh database metadata");
        return 2;
    }

    lua_newtable(L);

    lua_pushnumber(L, (lua_Number)db->meta.page_count);
    lua_setfield(L, -2, "page_count");

    lua_pushnumber(L, (lua_Number)db->meta.txn_id);
    lua_setfield(L, -2, "txn_id");

    lua_pushnumber(L, (lua_Number)db->meta.retired_root_pgno);
    lua_setfield(L, -2, "retired_root_pgno");

    lua_pushnumber(L, (lua_Number)db->meta.retired_page_count);
    lua_setfield(L, -2, "retired_page_count");

    lua_pushnumber(L, (lua_Number)db->meta.keyspace_root_pgno);
    lua_setfield(L, -2, "keyspace_root_pgno");

    lua_pushnumber(L, (lua_Number)db->meta.page_count * MNEME_PAGE_SIZE);
    lua_setfield(L, -2, "file_size");

    return 1;
}

/* ── txn_begin(db, readonly) ───────────────────────────────────────── */

static int l_txn_begin(lua_State *L) {
    lua_mneme_db_t *ud = check_db(L);
    int readonly       = lua_toboolean(L, 2);

    if (!readonly && ud->db->write_txn_active) {
        lua_pushnil(L);
        lua_pushstring(L, "write transaction already active");
        return 2;
    }

    /* All Lua allocations happen before mneme_txn_begin: a longjmp on
       OOM after the begin would leak the txn with the write flock held. */
    lua_mneme_txn_t *tud = (lua_mneme_txn_t *)lua_newuserdata(L, sizeof(*tud));
    tud->txn             = NULL;
    luaL_getmetatable(L, MNEME_TXN_MT);
    lua_setmetatable(L, -2);
    anchor_parent(L, 1); /* keep the db userdata alive while the txn exists */

    mneme_txn_t *txn = NULL;
    if (mneme_txn_begin(ud->db, readonly, &txn) < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "failed to begin transaction");
        return 2;
    }
    tud->txn = txn;
    return 1;
}

/* ── txn:commit() ──────────────────────────────────────────────────── */

static int l_txn_commit(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn(L);
    int rc              = mneme_txn_commit(ud->txn);
    if (rc == -2) {
        /* Commit is durable but mmap refresh failed; handle is now closed. */
        lua_pushnil(L);
        lua_pushstring(L, "commit succeeded but mmap refresh failed; db handle invalidated");
        return 2;
    }
    if (rc < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "commit failed");
        return 2;
    }
    lua_pushboolean(L, 1);
    return 1;
}

/* ── txn:abort() ───────────────────────────────────────────────────── */

static int l_txn_abort(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn(L);
    mneme_txn_abort(ud->txn);
    return 0;
}

/* ── txn:page_alloc() ──────────────────────────────────────────────── */

static int l_txn_page_alloc(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn_rw(L);
    uint32_t pgno       = mneme_page_alloc(ud->txn);
    if (pgno == 0) {
        lua_pushnil(L);
        lua_pushstring(L, "page allocation failed");
        return 2;
    }
    lua_pushnumber(L, (lua_Number)pgno);
    return 1;
}

/* ── txn:page_free(pgno) ──────────────────────────────────────────── */

static int l_txn_page_free(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn_rw(L);
    uint32_t pgno       = (uint32_t)luaL_checknumber(L, 2);
    if (mneme_page_free(ud->txn, pgno) < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "page free failed");
        return 2;
    }
    lua_pushboolean(L, 1);
    return 1;
}

/* ── txn:meta() -- return current txn meta state ───────────────────── */

static int l_txn_meta(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn(L);
    mneme_txn_t *txn    = ud->txn;

    lua_newtable(L);

    lua_pushnumber(L, (lua_Number)txn->meta.page_count);
    lua_setfield(L, -2, "page_count");

    lua_pushnumber(L, (lua_Number)txn->meta.txn_id);
    lua_setfield(L, -2, "txn_id");

    lua_pushnumber(L, (lua_Number)txn->meta.retired_root_pgno);
    lua_setfield(L, -2, "retired_root_pgno");

    lua_pushnumber(L, (lua_Number)txn->meta.retired_page_count);
    lua_setfield(L, -2, "retired_page_count");

    lua_pushnumber(L, (lua_Number)txn->meta.keyspace_root_pgno);
    lua_setfield(L, -2, "keyspace_root_pgno");

    return 1;
}

/* Resolve an overflow or inline value.
   On inline values (no MNEME_OVERFLOW_FLAG), sets *out_ptr = val,
   *out_len = val_len, *out_owned = NULL.
   On overflow values, allocates *out_owned (caller must free) with the
   resolved content, and sets *out_ptr = *out_owned, *out_len = bytes_read.
   Returns 0 on success, -1 on malloc or read failure. */
static int resolve_value(mneme_txn_t *txn, uint8_t type_tag, const uint8_t *val, uint32_t val_len,
                         const uint8_t **out_ptr, uint32_t *out_len, uint8_t **out_owned) {
    if (!(type_tag & MNEME_OVERFLOW_FLAG)) {
        *out_ptr   = val;
        *out_len   = val_len;
        *out_owned = NULL;
        return 0;
    }
    uint32_t ovfl_pgno;
    memcpy(&ovfl_pgno, val, 4);
    uint8_t *buf = (uint8_t *)malloc(val_len);
    if (!buf)
        return -1;
    uint32_t bytes_read;
    if (mneme_overflow_read(txn, ovfl_pgno, buf, val_len, &bytes_read) < 0) {
        free(buf);
        return -1;
    }
    *out_ptr   = buf;
    *out_len   = bytes_read;
    *out_owned = buf;
    return 0;
}

/* ── txn:btree_get(root_pgno, key) ─────────────────────────────────── */

static int l_txn_btree_get(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn(L);
    uint32_t root       = (uint32_t)luaL_checknumber(L, 2);
    size_t key_len;
    const char *key = luaL_checklstring(L, 3, &key_len);
    if (key_len == 0 || key_len > MNEME_MAX_KEY_LEN) {
        lua_pushnil(L);
        lua_pushstring(L, "invalid key length");
        return 2;
    }

    const uint8_t *val;
    uint32_t val_len;
    uint8_t type_tag;
    int64_t ttl;

    int rc = mneme_btree_get(ud->txn, root, (const uint8_t *)key, (uint16_t)key_len, &val, &val_len, &type_tag, &ttl);
    if (rc == 1) {
        lua_pushnil(L);
        lua_pushstring(L, "not found");
        return 2;
    }
    if (rc < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "btree get failed");
        return 2;
    }

    const uint8_t *vp;
    uint32_t vl;
    uint8_t *owned;
    if (resolve_value(ud->txn, type_tag, val, val_len, &vp, &vl, &owned) < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "overflow read failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)vp, vl);
    free(owned);

    lua_pushnumber(L, (lua_Number)(type_tag & 0x7F));
    lua_pushnumber(L, (lua_Number)ttl);
    return 3;
}

/* ── txn:btree_put(root_pgno, key, val, type_tag, ttl) ────────────── */

static int l_txn_btree_put(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn_rw(L);
    uint32_t root       = (uint32_t)luaL_checknumber(L, 2);
    size_t key_len;
    const char *key = luaL_checklstring(L, 3, &key_len);
    if (key_len == 0 || key_len > MNEME_MAX_KEY_LEN) {
        lua_pushnil(L);
        lua_pushstring(L, "invalid key length");
        return 2;
    }
    size_t val_len;
    const char *val = luaL_checklstring(L, 4, &val_len);
    double tag_d    = luaL_optnumber(L, 5, 1); /* default: Bytes */
    if (tag_d < 1 || tag_d > 0x7F) {
        /* Tags >= 0x80 would forge MNEME_OVERFLOW_FLAG, making readers
           treat the first 4 value bytes as an overflow page number. */
        lua_pushnil(L);
        lua_pushstring(L, "invalid type_tag");
        return 2;
    }
    uint8_t type_tag = (uint8_t)tag_d;
    int64_t ttl      = (int64_t)luaL_optnumber(L, 6, 0);

    uint32_t new_root = mneme_btree_put(ud->txn, root, (const uint8_t *)key, (uint16_t)key_len, (const uint8_t *)val,
                                        (uint32_t)val_len, type_tag, ttl);
    if (new_root == 0) {
        lua_pushnil(L);
        lua_pushstring(L, "btree put failed");
        return 2;
    }
    lua_pushnumber(L, (lua_Number)new_root);
    return 1;
}

/* ── txn:btree_del(root_pgno, key) ────────────────────────────────── */

static int l_txn_btree_del(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn_rw(L);
    uint32_t root       = (uint32_t)luaL_checknumber(L, 2);
    size_t key_len;
    const char *key = luaL_checklstring(L, 3, &key_len);
    if (key_len == 0 || key_len > MNEME_MAX_KEY_LEN) {
        lua_pushnil(L);
        lua_pushstring(L, "invalid key length");
        return 2;
    }

    uint32_t new_root = mneme_btree_del(ud->txn, root, (const uint8_t *)key, (uint16_t)key_len);
    if (new_root == 0) {
        lua_pushnil(L);
        lua_pushstring(L, "btree del failed");
        return 2;
    }
    lua_pushnumber(L, (lua_Number)new_root);
    return 1;
}

/* ── txn:btree_free(root_pgno) -- recursively free a whole subtree ── */

static int l_txn_btree_free(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn_rw(L);
    uint32_t root       = (uint32_t)luaL_checknumber(L, 2);
    if (root == 0) {
        lua_pushboolean(L, 1);
        return 1;
    }
    if (mneme_btree_free(ud->txn, root) < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "btree free failed");
        return 2;
    }
    lua_pushboolean(L, 1);
    return 1;
}

/* ── txn:set_ks_root(pgno) -- update meta keyspace_root_pgno ───────── */

static int l_txn_set_ks_root(lua_State *L) {
    lua_mneme_txn_t *ud              = check_txn_rw(L);
    uint32_t pgno                    = (uint32_t)luaL_checknumber(L, 2);
    ud->txn->meta.keyspace_root_pgno = pgno;
    return 0;
}

/* ── txn:alloc_leaf() -- allocate and initialize an empty leaf page ── */

static int l_txn_alloc_leaf(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn_rw(L);
    uint32_t pgno       = mneme_page_alloc(ud->txn);
    if (pgno == 0) {
        lua_pushnil(L);
        lua_pushstring(L, "allocation failed");
        return 2;
    }
    /* Initialize as empty leaf */
    uint8_t *page = mneme_dirty_find(ud->txn, pgno);
    if (page) {
        mneme_page_hdr_t *hdr = (mneme_page_hdr_t *)page;
        hdr->page_type        = MNEME_PAGE_LEAF;
        hdr->key_count        = 0;
        hdr->free_offset      = MNEME_PAGE_HDR_SIZE;
        hdr->cell_area_end    = MNEME_PAGE_SIZE;
        hdr->right_ptr        = 0;
    }
    lua_pushnumber(L, (lua_Number)pgno);
    return 1;
}

/* ── txn:btree_stats(root_pgno) → { keys, depth, leaf_pages, branch_pages, overflow_pages } */

static int seen_contains(uint32_t *seen, int seen_count, uint32_t pgno) {
    for (int i = 0; i < seen_count; i++) {
        if (seen[i] == pgno)
            return 1;
    }
    return 0;
}

static int seen_add(uint32_t **seen, int *seen_count, int *seen_cap, uint32_t pgno) {
    if (seen_contains(*seen, *seen_count, pgno))
        return 0;
    if (*seen_count >= *seen_cap) {
        int new_cap       = (*seen_cap == 0) ? 64 : (*seen_cap * 2);
        uint32_t *new_arr = (uint32_t *)realloc(*seen, (size_t)new_cap * sizeof(uint32_t));
        if (!new_arr)
            return -1;
        *seen     = new_arr;
        *seen_cap = new_cap;
    }
    (*seen)[(*seen_count)++] = pgno;
    return 1;
}

/* Recursion/iteration guard for the stats walks: a corrupted or crafted
   file can contain page cycles, so bound tree depth and chain length. */
#define STATS_MAX_DEPTH 64

static void btree_stats_recurse(mneme_txn_t *txn, uint32_t pgno, int depth, int *max_depth, int *keys, int *leaves,
                                int *branches, int *overflows) {
    if (depth > STATS_MAX_DEPTH)
        return;
    const uint8_t *page = mneme_page_resolve(txn, pgno);
    if (!page)
        return;
    const mneme_page_hdr_t *hdr = (const mneme_page_hdr_t *)page;

    if (depth > *max_depth)
        *max_depth = depth;

    if (hdr->page_type == MNEME_PAGE_LEAF) {
        (*leaves)++;
        *keys += hdr->key_count;
        /* Count overflow pages for each cell */
        for (int i = 0; i < hdr->key_count; i++) {
            uint16_t off;
            memcpy(&off, page + MNEME_PAGE_HDR_SIZE + i * 2, 2);
            uint8_t tag = page[off + MNEME_LCELL_HDR - 1]; /* type_tag is last byte of cell header */
            if (tag & MNEME_OVERFLOW_FLAG) {
                /* Overflow: count the chain */
                uint16_t kl;
                memcpy(&kl, page + off, 2);
                uint32_t ovfl_pgno;
                memcpy(&ovfl_pgno, page + off + MNEME_LCELL_HDR + kl, 4);
                /* A chain can't legitimately exceed the page count; a
                   cyclic chain in a corrupt file would spin forever. */
                uint64_t chain_budget = txn->meta.page_count;
                while (ovfl_pgno != 0 && chain_budget-- > 0) {
                    (*overflows)++;
                    const uint8_t *op = mneme_page_resolve(txn, ovfl_pgno);
                    if (!op)
                        break;
                    memcpy(&ovfl_pgno, op + 4, 4);
                }
            }
        }
    } else if (hdr->page_type == MNEME_PAGE_BRANCH) {
        (*branches)++;
        for (int i = 0; i < hdr->key_count; i++) {
            uint16_t off;
            memcpy(&off, page + MNEME_PAGE_HDR_SIZE + i * 2, 2);
            uint32_t child;
            memcpy(&child, page + off, 4);
            btree_stats_recurse(txn, child, depth + 1, max_depth, keys, leaves, branches, overflows);
        }
        /* right_ptr child */
        if (hdr->right_ptr != 0)
            btree_stats_recurse(txn, hdr->right_ptr, depth + 1, max_depth, keys, leaves, branches, overflows);
    }
}

static int l_txn_btree_stats(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn(L);
    uint32_t root       = (uint32_t)luaL_checknumber(L, 2);

    int max_depth = 0, keys = 0, leaves = 0, branches = 0, overflows = 0;
    btree_stats_recurse(ud->txn, root, 1, &max_depth, &keys, &leaves, &branches, &overflows);

    lua_newtable(L);
    lua_pushnumber(L, keys);
    lua_setfield(L, -2, "keys");
    lua_pushnumber(L, max_depth);
    lua_setfield(L, -2, "depth");
    lua_pushnumber(L, leaves);
    lua_setfield(L, -2, "leaf_pages");
    lua_pushnumber(L, branches);
    lua_setfield(L, -2, "branch_pages");
    lua_pushnumber(L, overflows);
    lua_setfield(L, -2, "overflow_pages");
    return 1;
}

static void live_pages_recurse(mneme_txn_t *txn, uint32_t pgno, int depth, uint32_t **seen, int *seen_count,
                               int *seen_cap) {
    if (depth > STATS_MAX_DEPTH)
        return;
    int added = seen_add(seen, seen_count, seen_cap, pgno);
    if (added <= 0)
        return;

    const uint8_t *page = mneme_page_resolve(txn, pgno);
    if (!page)
        return;
    const mneme_page_hdr_t *hdr = (const mneme_page_hdr_t *)page;

    if (hdr->page_type == MNEME_PAGE_LEAF) {
        for (int i = 0; i < hdr->key_count; i++) {
            uint16_t off;
            memcpy(&off, page + MNEME_PAGE_HDR_SIZE + i * 2, 2);
            uint8_t tag = page[off + MNEME_LCELL_HDR - 1];
            if (tag & MNEME_OVERFLOW_FLAG) {
                uint16_t kl;
                memcpy(&kl, page + off, 2);
                uint32_t ovfl_pgno;
                memcpy(&ovfl_pgno, page + off + MNEME_LCELL_HDR + kl, 4);
                while (ovfl_pgno != 0) {
                    int ovfl_added    = seen_add(seen, seen_count, seen_cap, ovfl_pgno);
                    const uint8_t *op = mneme_page_resolve(txn, ovfl_pgno);
                    if (ovfl_added < 0 || !op)
                        break;
                    memcpy(&ovfl_pgno, op + 4, 4);
                    if (ovfl_added == 0)
                        break;
                }
            }
        }
        return;
    }

    if (hdr->page_type == MNEME_PAGE_BRANCH) {
        for (int i = 0; i < hdr->key_count; i++) {
            uint16_t off;
            memcpy(&off, page + MNEME_PAGE_HDR_SIZE + i * 2, 2);
            uint32_t child;
            memcpy(&child, page + off, 4);
            live_pages_recurse(txn, child, depth + 1, seen, seen_count, seen_cap);
        }
        if (hdr->right_ptr != 0)
            live_pages_recurse(txn, hdr->right_ptr, depth + 1, seen, seen_count, seen_cap);
    }
}

/* ── txn:live_count() → number of pages reachable from current meta ─── */

static int l_txn_live_count(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn(L);
    uint32_t *seen      = NULL;
    int seen_count = 0, seen_cap = 0;

    seen_add(&seen, &seen_count, &seen_cap, 0);
    seen_add(&seen, &seen_count, &seen_cap, 1);
    live_pages_recurse(ud->txn, ud->txn->meta.keyspace_root_pgno, 1, &seen, &seen_count, &seen_cap);
    if (ud->txn->meta.retired_root_pgno != 0)
        live_pages_recurse(ud->txn, ud->txn->meta.retired_root_pgno, 1, &seen, &seen_count, &seen_cap);

    mneme_cursor_t *cur = NULL;
    if (mneme_cursor_open(ud->txn, ud->txn->meta.keyspace_root_pgno, &cur) == 0) {
        mneme_cursor_first(cur);
        while (cur->valid) {
            const uint8_t *key, *val;
            uint16_t key_len;
            uint32_t val_len;
            uint8_t type_tag;
            int64_t ttl;
            if (mneme_cursor_entry(cur, &key, &key_len, &val, &val_len, &type_tag, &ttl) == 0) {
                if (key_len > 0 && key[0] != '\0' && val_len == 4) {
                    uint32_t root_pgno;
                    memcpy(&root_pgno, val, 4);
                    live_pages_recurse(ud->txn, root_pgno, 1, &seen, &seen_count, &seen_cap);
                }
            }
            mneme_cursor_next(cur);
        }
        mneme_cursor_close(cur);
    }

    free(seen);
    lua_pushnumber(L, seen_count);
    return 1;
}

/* ── Cursor bindings ───────────────────────────────────────────────── */

typedef struct {
    mneme_cursor_t *cur;
} lua_mneme_cur_t;

/* txn:cursor_open(root_pgno) */
static int l_txn_cursor_open(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn(L);
    uint32_t root       = (uint32_t)luaL_checknumber(L, 2);

    /* Userdata before the C allocation, so an OOM longjmp can't leak it. */
    lua_mneme_cur_t *cud = (lua_mneme_cur_t *)lua_newuserdata(L, sizeof(*cud));
    cud->cur             = NULL;
    luaL_getmetatable(L, MNEME_CUR_MT);
    lua_setmetatable(L, -2);
    anchor_parent(L, 1); /* keep the txn userdata alive while the cursor exists */

    mneme_cursor_t *cur = NULL;
    if (mneme_cursor_open(ud->txn, root, &cur) < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "cursor open failed");
        return 2;
    }
    cud->cur = cur;
    return 1;
}

static lua_mneme_cur_t *check_cur(lua_State *L) {
    lua_mneme_cur_t *ud = (lua_mneme_cur_t *)luaL_checkudata(L, 1, MNEME_CUR_MT);
    if (!ud->cur)
        luaL_error(L, "cursor is closed");
    /* The cursor anchors its txn userdata, so cur->txn is always valid
       memory here -- but the txn may have committed/aborted, freeing the
       snapshot/dirty state the cursor walks. */
    if (ud->cur->txn->committed)
        luaL_error(L, "transaction is finished");
    return ud;
}

/* Push key, value, type_tag, ttl from current cursor position.
   Handles overflow values. Returns 4 on success, or 1 (nil) on failure. */
static int push_cursor_entry(lua_State *L, mneme_cursor_t *cur) {
    const uint8_t *key, *val;
    uint16_t key_len;
    uint32_t val_len;
    uint8_t type_tag;
    int64_t ttl;

    if (mneme_cursor_entry(cur, &key, &key_len, &val, &val_len, &type_tag, &ttl) < 0) {
        lua_pushnil(L);
        return 1;
    }

    const uint8_t *vp;
    uint32_t vl;
    uint8_t *owned;
    if (resolve_value(cur->txn, type_tag, val, val_len, &vp, &vl, &owned) < 0) {
        lua_pushnil(L);
        return 1;
    }

    lua_pushlstring(L, (const char *)key, key_len);
    lua_pushlstring(L, (const char *)vp, vl);
    free(owned);
    lua_pushnumber(L, (lua_Number)(type_tag & 0x7F));
    lua_pushnumber(L, (lua_Number)ttl);
    return 4;
}

/* cur:first() -> key, val, type, ttl OR nil */
static int l_cur_first(lua_State *L) {
    lua_mneme_cur_t *ud = check_cur(L);
    mneme_cursor_first(ud->cur);
    if (!ud->cur->valid) {
        lua_pushnil(L);
        return 1;
    }
    return push_cursor_entry(L, ud->cur);
}

/* cur:last() -> key, val, type, ttl OR nil */
static int l_cur_last(lua_State *L) {
    lua_mneme_cur_t *ud = check_cur(L);
    mneme_cursor_last(ud->cur);
    if (!ud->cur->valid) {
        lua_pushnil(L);
        return 1;
    }
    return push_cursor_entry(L, ud->cur);
}

/* cur:seek(key) -> key, val, type, ttl OR nil */
static int l_cur_seek(lua_State *L) {
    lua_mneme_cur_t *ud = check_cur(L);
    size_t key_len;
    const char *key = luaL_checklstring(L, 2, &key_len);
    if (key_len > MNEME_MAX_KEY_LEN) {
        lua_pushnil(L);
        lua_pushstring(L, "key too long");
        return 2;
    }
    mneme_cursor_seek(ud->cur, (const uint8_t *)key, (uint16_t)key_len);
    if (!ud->cur->valid) {
        lua_pushnil(L);
        return 1;
    }
    return push_cursor_entry(L, ud->cur);
}

/* cur:next() -> key, val, type, ttl OR nil */
static int l_cur_next(lua_State *L) {
    lua_mneme_cur_t *ud = check_cur(L);
    mneme_cursor_next(ud->cur);
    if (!ud->cur->valid) {
        lua_pushnil(L);
        return 1;
    }
    return push_cursor_entry(L, ud->cur);
}

/* cur:prev() -> key, val, type, ttl OR nil */
static int l_cur_prev(lua_State *L) {
    lua_mneme_cur_t *ud = check_cur(L);
    mneme_cursor_prev(ud->cur);
    if (!ud->cur->valid) {
        lua_pushnil(L);
        return 1;
    }
    return push_cursor_entry(L, ud->cur);
}

/* cur:current() -> key, val, type, ttl OR nil (read without moving) */
static int l_cur_current(lua_State *L) {
    lua_mneme_cur_t *ud = check_cur(L);
    if (!ud->cur->valid) {
        lua_pushnil(L);
        return 1;
    }
    return push_cursor_entry(L, ud->cur);
}

/* cur:valid() -> boolean */
static int l_cur_valid(lua_State *L) {
    lua_mneme_cur_t *ud = check_cur(L);
    lua_pushboolean(L, ud->cur->valid);
    return 1;
}

/* cur:next_batch(batch_size [, end_key [, prefix]]) -> table, count
   Reads up to batch_size entries from the current cursor position into
   a flat Lua table: {k1, v1, tag1, ttl1, k2, v2, tag2, ttl2, ...}.
   Stops early if end_key or prefix boundary is reached.
   Advances the cursor past the last returned entry. */
static int l_cur_next_batch(lua_State *L) {
    lua_mneme_cur_t *ud = check_cur(L);
    int batch_size      = (int)luaL_checkinteger(L, 2);

    size_t end_key_len  = 0;
    const char *end_key = NULL;
    if (lua_isstring(L, 3))
        end_key = lua_tolstring(L, 3, &end_key_len);
    if (end_key && end_key_len > MNEME_MAX_KEY_LEN) {
        lua_pushnil(L);
        lua_pushstring(L, "end key too long");
        return 2;
    }

    size_t prefix_len  = 0;
    const char *prefix = NULL;
    if (lua_isstring(L, 4))
        prefix = lua_tolstring(L, 4, &prefix_len);

    if (batch_size < 1)
        batch_size = 1;
    if (batch_size > 1024)
        batch_size = 1024;

    lua_createtable(L, batch_size * 4, 0);
    int count   = 0;
    int tab_idx = 1;

    for (int i = 0; i < batch_size; i++) {
        if (!ud->cur->valid)
            break;

        const uint8_t *key, *val;
        uint16_t key_len;
        uint32_t val_len;
        uint8_t type_tag;
        int64_t ttl;

        if (mneme_cursor_entry(ud->cur, &key, &key_len, &val, &val_len, &type_tag, &ttl) < 0)
            break;

        /* Range boundary check */
        if (end_key && mneme_keycmp(key, key_len, (const uint8_t *)end_key, (uint16_t)end_key_len) >= 0)
            break;

        /* Prefix boundary check (compare lengths in size_t: a truncated
           uint16 prefix_len would let memcmp over-read past the key) */
        if (prefix && ((size_t)key_len < prefix_len || memcmp(key, prefix, prefix_len) != 0))
            break;

        /* Key */
        lua_pushlstring(L, (const char *)key, key_len);
        lua_rawseti(L, -2, tab_idx++);

        /* Value (handle overflow) */
        {
            const uint8_t *vp;
            uint32_t vl;
            uint8_t *owned = NULL;
            if (resolve_value(ud->cur->txn, type_tag, val, val_len, &vp, &vl, &owned) < 0) {
                free(owned);
                lua_pushnil(L);
                lua_pushstring(L, "overflow read failed");
                return 2;
            }
            lua_pushlstring(L, (const char *)vp, vl);
            free(owned);
            lua_rawseti(L, -2, tab_idx++);
        }

        /* Type tag (clear overflow flag) */
        lua_pushnumber(L, (lua_Number)(type_tag & 0x7F));
        lua_rawseti(L, -2, tab_idx++);

        /* TTL */
        lua_pushnumber(L, (lua_Number)ttl);
        lua_rawseti(L, -2, tab_idx++);

        count++;
        mneme_cursor_next(ud->cur);
    }

    lua_pushinteger(L, count);
    return 2;
}

/* cur:close() */
static int l_cur_close(lua_State *L) {
    lua_mneme_cur_t *ud = (lua_mneme_cur_t *)luaL_checkudata(L, 1, MNEME_CUR_MT);
    if (ud->cur) {
        mneme_cursor_close(ud->cur);
        ud->cur = NULL;
    }
    return 0;
}

/* ── GC for txn (abort if not committed) ───────────────────────────── */

static int l_txn_gc(lua_State *L) {
    lua_mneme_txn_t *ud = (lua_mneme_txn_t *)luaL_checkudata(L, 1, MNEME_TXN_MT);
    if (ud->txn && !ud->txn->committed) {
        mneme_txn_abort(ud->txn);
    }
    if (ud->txn) {
        free(ud->txn);
        ud->txn = NULL;
    }
    return 0;
}

/* ── Registration ──────────────────────────────────────────────────── */

static const luaL_Reg db_methods[] = {
    {"close",     l_db_close },
    {"sync",      l_db_sync  },
    {"info",      l_db_info  },
    {"txn_begin", l_txn_begin},
    {NULL,        NULL       }
};

static const luaL_Reg txn_methods[] = {
    {"commit",      l_txn_commit     },
    {"abort",       l_txn_abort      },
    {"page_alloc",  l_txn_page_alloc },
    {"page_free",   l_txn_page_free  },
    {"meta",        l_txn_meta       },
    {"btree_get",   l_txn_btree_get  },
    {"btree_put",   l_txn_btree_put  },
    {"btree_del",   l_txn_btree_del  },
    {"btree_free",  l_txn_btree_free },
    {"cursor_open", l_txn_cursor_open},
    {"set_ks_root", l_txn_set_ks_root},
    {"alloc_leaf",  l_txn_alloc_leaf },
    {"btree_stats", l_txn_btree_stats},
    {"live_count",  l_txn_live_count },
    {NULL,          NULL             }
};

static const luaL_Reg cursor_methods[] = {
    {"first",      l_cur_first     },
    {"last",       l_cur_last      },
    {"seek",       l_cur_seek      },
    {"next",       l_cur_next      },
    {"prev",       l_cur_prev      },
    {"current",    l_cur_current   },
    {"valid",      l_cur_valid     },
    {"next_batch", l_cur_next_batch},
    {"close",      l_cur_close     },
    {NULL,         NULL            }
};

/* ── Pack/unpack helpers for typed values ───────────────────────────── */

/* pack_int64(number) -> 8-byte native-endian string */
static int l_pack_int64(lua_State *L) {
    double d  = luaL_checknumber(L, 1);
    int64_t v = (int64_t)d;
    char buf[8];
    memcpy(buf, &v, 8);
    lua_pushlstring(L, buf, 8);
    return 1;
}

/* unpack_int64(string) -> number */
static int l_unpack_int64(lua_State *L) {
    size_t len;
    const char *s = luaL_checklstring(L, 1, &len);
    if (len < 8) {
        lua_pushnil(L);
        lua_pushstring(L, "need 8 bytes");
        return 2;
    }
    int64_t v;
    memcpy(&v, s, 8);
    lua_pushnumber(L, (lua_Number)v);
    return 1;
}

/* pack_double(number) -> 8-byte native-endian string */
static int l_pack_double(lua_State *L) {
    double v = luaL_checknumber(L, 1);
    char buf[8];
    memcpy(buf, &v, 8);
    lua_pushlstring(L, buf, 8);
    return 1;
}

/* unpack_double(string) -> number */
static int l_unpack_double(lua_State *L) {
    size_t len;
    const char *s = luaL_checklstring(L, 1, &len);
    if (len < 8) {
        lua_pushnil(L);
        lua_pushstring(L, "need 8 bytes");
        return 2;
    }
    double v;
    memcpy(&v, s, 8);
    lua_pushnumber(L, v);
    return 1;
}

/* ── Score encoding for sorted sets (IEEE 754 → lexicographic) ─────── */

/* Encode a double so that memcmp gives numeric ordering.
   Flip sign bit; if original was negative, flip all remaining bits. */
static int l_encode_score(lua_State *L) {
    double d = luaL_checknumber(L, 1);
    uint64_t bits;
    memcpy(&bits, &d, 8);

    /* Convert to big-endian for memcmp ordering */
    if (bits & ((uint64_t)1 << 63)) {
        /* Negative: flip all bits */
        bits = ~bits;
    } else {
        /* Positive (or zero): flip sign bit only */
        bits ^= ((uint64_t)1 << 63);
    }

    /* Store as big-endian */
    uint8_t buf[8];
    for (int i = 7; i >= 0; i--) {
        buf[7 - i] = (uint8_t)(bits >> (i * 8));
    }
    lua_pushlstring(L, (const char *)buf, 8);
    return 1;
}

/* Decode score bytes back to double */
static int l_decode_score(lua_State *L) {
    size_t len;
    const char *s = luaL_checklstring(L, 1, &len);
    if (len < 8) {
        lua_pushnil(L);
        lua_pushstring(L, "need 8 bytes");
        return 2;
    }

    /* Read big-endian */
    const uint8_t *b = (const uint8_t *)s;
    uint64_t bits    = 0;
    for (int i = 0; i < 8; i++) {
        bits = (bits << 8) | b[i];
    }

    /* Reverse the encoding */
    if (bits & ((uint64_t)1 << 63)) {
        /* Was positive: flip sign bit back */
        bits ^= ((uint64_t)1 << 63);
    } else {
        /* Was negative: flip all bits back */
        bits = ~bits;
    }

    double d;
    memcpy(&d, &bits, 8);
    lua_pushnumber(L, d);
    return 1;
}

/* ── Big-endian int64 for list sequence numbers ────────────────────── */

static int l_pack_int64_be(lua_State *L) {
    double d  = luaL_checknumber(L, 1);
    int64_t v = (int64_t)d;
    uint8_t buf[8];
    for (int i = 7; i >= 0; i--) {
        buf[7 - i] = (uint8_t)((uint64_t)v >> (i * 8));
    }
    lua_pushlstring(L, (const char *)buf, 8);
    return 1;
}

static int l_unpack_int64_be(lua_State *L) {
    size_t len;
    const char *s = luaL_checklstring(L, 1, &len);
    if (len < 8) {
        lua_pushnil(L);
        lua_pushstring(L, "need 8 bytes");
        return 2;
    }
    const uint8_t *b = (const uint8_t *)s;
    int64_t v        = 0;
    for (int i = 0; i < 8; i++) {
        v = (v << 8) | b[i];
    }
    lua_pushnumber(L, (lua_Number)v);
    return 1;
}

/* ── Vector helpers ─────────────────────────────────────────────────── */

#include <math.h>

/* pack_floats(table_of_numbers) → binary string of float32 values */
static int l_pack_floats(lua_State *L) {
    luaL_checktype(L, 1, LUA_TTABLE);
    int dim = (int)lua_objlen(L, 1);
    if (dim <= 0) {
        lua_pushnil(L);
        lua_pushstring(L, "empty vector");
        return 2;
    }

    /* Encode: [uint16 dim] [float32 × dim] */
    size_t total = 2 + (size_t)dim * 4;
    uint8_t *buf = (uint8_t *)malloc(total);
    if (!buf) {
        lua_pushnil(L);
        lua_pushstring(L, "out of memory");
        return 2;
    }

    uint16_t d16 = (uint16_t)dim;
    memcpy(buf, &d16, 2);

    float *fp = (float *)(buf + 2);
    for (int i = 0; i < dim; i++) {
        lua_rawgeti(L, 1, i + 1);
        fp[i] = (float)lua_tonumber(L, -1);
        lua_pop(L, 1);
    }

    lua_pushlstring(L, (const char *)buf, total);
    free(buf);
    return 1;
}

/* unpack_floats(binary_string) → table_of_numbers */
static int l_unpack_floats(lua_State *L) {
    size_t len;
    const char *s = luaL_checklstring(L, 1, &len);
    if (len < 2) {
        lua_pushnil(L);
        lua_pushstring(L, "too short");
        return 2;
    }

    uint16_t dim;
    memcpy(&dim, s, 2);
    if (len < 2 + (size_t)dim * 4) {
        lua_pushnil(L);
        lua_pushstring(L, "truncated");
        return 2;
    }

    const float *fp = (const float *)(s + 2);
    lua_createtable(L, dim, 0);
    for (int i = 0; i < dim; i++) {
        lua_pushnumber(L, (lua_Number)fp[i]);
        lua_rawseti(L, -2, i + 1);
    }
    return 1;
}

/* vec_dim(binary_string) → dimension (uint16 from first 2 bytes) */
static int l_vec_dim(lua_State *L) {
    size_t len;
    const char *s = luaL_checklstring(L, 1, &len);
    if (len < 2) {
        lua_pushnil(L);
        lua_pushstring(L, "too short");
        return 2;
    }
    uint16_t dim;
    memcpy(&dim, s, 2);
    lua_pushnumber(L, (lua_Number)dim);
    return 1;
}

typedef struct {
    float score;
    int idx;
} heap_entry_t;

static int heap_entry_cmp(const void *a, const void *b) {
    float as = ((const heap_entry_t *)a)->score;
    float bs = ((const heap_entry_t *)b)->score;
    return (as > bs) - (as < bs);
}

/* vsearch_scan(txn, root_pgno, query_packed, metric, top_k, prefix_filter)
   Scans all entries with type_tag=VECTOR, computes distance, returns top_k.
   metric: 1=cosine, 2=dot, 3=l2
   Returns a table of { key=string, score=number } sorted by score. */
static int l_vsearch_scan(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn(L);

    uint32_t root = (uint32_t)luaL_checknumber(L, 2);
    size_t q_len;
    const char *q_packed = luaL_checklstring(L, 3, &q_len);
    int metric           = (int)luaL_checknumber(L, 4); /* 1=cosine, 2=dot, 3=l2 */
    int top_k            = (int)luaL_checknumber(L, 5);
    if (top_k < 1 || top_k > 10000) {
        lua_pushnil(L);
        lua_pushstring(L, "invalid top_k");
        return 2;
    }
    size_t filter_len  = 0;
    const char *filter = NULL;
    if (lua_isstring(L, 6)) {
        filter = lua_tolstring(L, 6, &filter_len);
    }

    /* Parse query vector */
    if (q_len < 2) {
        lua_pushnil(L);
        lua_pushstring(L, "bad query");
        return 2;
    }
    uint16_t q_dim;
    memcpy(&q_dim, q_packed, 2);
    if (q_len < 2 + (size_t)q_dim * 4) {
        lua_pushnil(L);
        lua_pushstring(L, "truncated query");
        return 2;
    }
    const float *q_vec = (const float *)(q_packed + 2);

    /* Max-heap for top_k results (score, key_idx_in_results table).
       For cosine/l2: lower is better → max-heap evicts largest.
       For dot: higher is better → max-heap evicts smallest (negate scores). */
    heap_entry_t *heap = (heap_entry_t *)calloc((size_t)top_k, sizeof(heap_entry_t));
    if (!heap) {
        lua_pushnil(L);
        lua_pushstring(L, "out of memory");
        return 2;
    }
    int heap_size = 0;

    /* We also need to store keys. Use a Lua table for that. */
    lua_newtable(L); /* results accumulator at stack top */
    int results_idx  = lua_gettop(L);
    int result_count = 0;

    /* Scan via cursor */
    mneme_cursor_t *cur = NULL;
    if (mneme_cursor_open(ud->txn, root, &cur) < 0) {
        free(heap);
        lua_pushnil(L);
        lua_pushstring(L, "cursor failed");
        return 2;
    }

    mneme_cursor_first(cur);
    while (cur->valid) {
        const uint8_t *key, *val;
        uint16_t key_len;
        uint32_t val_len;
        uint8_t type_tag;
        int64_t ttl;

        if (mneme_cursor_key(cur, &key, &key_len) < 0 || mneme_cursor_val(cur, &val, &val_len, &type_tag, &ttl) < 0)
            break;

        /* Only process vector entries */
        uint8_t base_tag = type_tag & 0x7F;
        if (base_tag != MNEME_TYPE_VECTOR) {
            mneme_cursor_next(cur);
            continue;
        }

        /* Optional prefix filter on key */
        if (filter && filter_len > 0) {
            if (key_len < filter_len || memcmp(key, filter, filter_len) != 0) {
                mneme_cursor_next(cur);
                continue;
            }
        }

        /* Parse stored vector: [uint16 dim][float32 × dim] */
        /* For overflow vectors, val points to the overflow pgno — resolve it. */
        const uint8_t *vp;
        uint32_t vl;
        uint8_t *ovfl_buf;
        if (resolve_value(ud->txn, type_tag, val, val_len, &vp, &vl, &ovfl_buf) < 0) {
            mneme_cursor_next(cur);
            continue;
        }
        val     = vp;
        val_len = vl;

        if (val_len < 2) {
            free(ovfl_buf);
            mneme_cursor_next(cur);
            continue;
        }
        uint16_t v_dim;
        memcpy(&v_dim, val, 2);
        if (v_dim != q_dim || val_len < 2 + (size_t)v_dim * 4) {
            free(ovfl_buf);
            mneme_cursor_next(cur);
            continue;
        }
        const float *v_vec = (const float *)(val + 2);

        /* Compute distance */
        float score;
        switch (metric) {
        case 1:
            score = mneme_vec_cosine(q_vec, v_vec, q_dim);
            break;
        case 2:
            score = mneme_vec_dot(q_vec, v_vec, q_dim);
            break;
        case 3:
            score = mneme_vec_l2(q_vec, v_vec, q_dim);
            break;
        default:
            score = mneme_vec_cosine(q_vec, v_vec, q_dim);
            break;
        }

        /* For dot product, negate for max-heap (we want highest dot) */
        float heap_score = (metric == 2) ? -score : score;

        /* NaN compares false against everything, so a NaN entry admitted
           during the fill phase would wedge the heap ordering. */
        if (isnan(heap_score)) {
            free(ovfl_buf);
            mneme_cursor_next(cur);
            continue;
        }

        /* Insert into heap */
        if (heap_size < top_k) {
            /* Add to results table */
            result_count++;
            lua_pushlstring(L, (const char *)key, key_len);
            lua_rawseti(L, results_idx, result_count);

            heap[heap_size].score = heap_score;
            heap[heap_size].idx   = result_count;
            heap_size++;

            /* Sift up */
            int c = heap_size - 1;
            while (c > 0) {
                int p = (c - 1) / 2;
                if (heap[p].score < heap[c].score) {
                    heap_entry_t tmp = heap[p];
                    heap[p]          = heap[c];
                    heap[c]          = tmp;
                    c                = p;
                } else
                    break;
            }
        } else if (heap_score < heap[0].score) {
            /* Replace root of max-heap */
            /* Overwrite old key in results table */
            lua_pushlstring(L, (const char *)key, key_len);
            lua_rawseti(L, results_idx, heap[0].idx);

            heap[0].score = heap_score;

            /* Sift down */
            int p = 0;
            for (;;) {
                int l = 2 * p + 1, r = 2 * p + 2, largest = p;
                if (l < heap_size && heap[l].score > heap[largest].score)
                    largest = l;
                if (r < heap_size && heap[r].score > heap[largest].score)
                    largest = r;
                if (largest == p)
                    break;
                heap_entry_t tmp = heap[p];
                heap[p]          = heap[largest];
                heap[largest]    = tmp;
                p                = largest;
            }
        }

        free(ovfl_buf);
        mneme_cursor_next(cur);
    }
    mneme_cursor_close(cur);

    /* Build sorted result table */
    /* Collect heap entries, sort by score (ascending for cosine/l2, descending for dot) */
    lua_createtable(L, heap_size, 0);
    int out_idx = lua_gettop(L);

    /* Ascending by heap score (top_k can be as large as 10000) */
    qsort(heap, (size_t)heap_size, sizeof(heap_entry_t), heap_entry_cmp);

    for (int i = 0; i < heap_size; i++) {
        lua_createtable(L, 0, 2);

        lua_rawgeti(L, results_idx, heap[i].idx);
        lua_setfield(L, -2, "key");

        /* Restore original score for dot product */
        float orig = (metric == 2) ? -heap[i].score : heap[i].score;
        lua_pushnumber(L, (lua_Number)orig);
        lua_setfield(L, -2, "score");

        lua_rawseti(L, out_idx, i + 1);
    }

    free(heap);
    /* Remove the intermediate results table, keep the output table */
    lua_remove(L, results_idx);
    return 1;
}

/* ── u32 packing helpers ───────────────────────────────────────────── */

static int l_pack_u32(lua_State *L) {
    uint32_t n = (uint32_t)luaL_checknumber(L, 1);
    uint8_t buf[4];
    memcpy(buf, &n, 4); /* native-endian */
    lua_pushlstring(L, (const char *)buf, 4);
    return 1;
}

static int l_unpack_u32(lua_State *L) {
    size_t len;
    const char *s = luaL_checklstring(L, 1, &len);
    if (len < 4) {
        lua_pushnil(L);
        lua_pushstring(L, "too short");
        return 2;
    }
    uint32_t n;
    memcpy(&n, s, 4);
    lua_pushnumber(L, (lua_Number)n);
    return 1;
}

/* ── Full-text search helpers ──────────────────────────────────────── */

/* ft_analyze(text) → table of term strings */
static int l_ft_analyze(lua_State *L) {
    size_t text_len;
    const char *text = luaL_checklstring(L, 1, &text_len);

    mneme_ft_terms_t terms;
    if (mneme_ft_analyze((const uint8_t *)text, (uint32_t)text_len, &terms) < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "out of memory");
        return 2;
    }

    lua_createtable(L, terms.count, 0);
    for (int i = 0; i < terms.count; i++) {
        lua_pushlstring(L, (const char *)terms.terms[i].term, terms.terms[i].len);
        lua_rawseti(L, -2, i + 1);
    }

    mneme_ft_terms_free(&terms);
    return 1;
}

/* ft_search(txn, ks_root, name, query_terms, top_k, k1, b)
   Performs BM25 search over the inverted index.
   query_terms: Lua table of deduplicated term strings.
   Returns table of { key=string, score=number } sorted by descending score. */
static int l_ft_search(lua_State *L) {
    lua_mneme_txn_t *ud = check_txn(L);

    uint32_t ks_root = (uint32_t)luaL_checknumber(L, 2);
    size_t name_len;
    const char *name = luaL_checklstring(L, 3, &name_len);
    luaL_checktype(L, 4, LUA_TTABLE);
    int top_k = (int)luaL_checknumber(L, 5);
    double k1 = luaL_checknumber(L, 6);
    double b  = luaL_checknumber(L, 7);

    if (top_k <= 0)
        top_k = 10;
    if (top_k > 10000)
        top_k = 10000;

    /* Read corpus stats from key: name \x00 \x00 */
    uint8_t meta_key[4096];
    if (name_len + 2 > sizeof(meta_key)) {
        lua_createtable(L, 0, 0);
        return 1;
    }
    memcpy(meta_key, name, name_len);
    meta_key[name_len]     = 0x00;
    meta_key[name_len + 1] = 0x00;

    const uint8_t *meta_val;
    uint32_t meta_val_len;
    uint8_t meta_tag;
    int64_t meta_ttl;
    int rc = mneme_btree_get(ud->txn, ks_root, meta_key, (uint16_t)(name_len + 2), &meta_val, &meta_val_len, &meta_tag,
                             &meta_ttl);
    if (rc != 0 || meta_val_len < 12) {
        /* No corpus stats → empty corpus */
        lua_createtable(L, 0, 0);
        return 1;
    }

    uint32_t doc_count;
    uint64_t total_doc_len;
    memcpy(&doc_count, meta_val, 4);
    memcpy(&total_doc_len, meta_val + 4, 8);

    if (doc_count == 0) {
        lua_createtable(L, 0, 0);
        return 1;
    }

    double avgdl = (double)total_doc_len / (double)doc_count;
    if (avgdl <= 0.0)
        avgdl = 1.0; /* drifted stats would yield 0/0 = NaN scores below */
    double N = (double)doc_count;

    /* Count query terms */
    int q_count = (int)lua_objlen(L, 4);
    if (q_count == 0) {
        lua_createtable(L, 0, 0);
        return 1;
    }

    /* ── Hash map for score accumulation ──────────────────────────── */

    typedef struct {
        const uint8_t *doc_id;
        uint16_t doc_id_len;
        double score;
        uint32_t doc_len; /* cached, 0 = not yet loaded */
        int doc_len_loaded;
        int next; /* index into arena, -1 = end of chain */
    } score_entry_t;

    int map_cap = 256;
    while (map_cap < top_k * 2)
        map_cap *= 2;
    /* Bucket array stores indices into arena (-1 = empty) */
    int *map_buckets = (int *)malloc((size_t)map_cap * sizeof(int));
    if (!map_buckets) {
        lua_pushnil(L);
        lua_pushstring(L, "out of memory");
        return 2;
    }
    for (int i = 0; i < map_cap; i++)
        map_buckets[i] = -1;

    /* Arena for score entries */
    int arena_cap        = 1024;
    int arena_count      = 0;
    score_entry_t *arena = (score_entry_t *)malloc((size_t)arena_cap * sizeof(score_entry_t));
    if (!arena) {
        free(map_buckets);
        lua_pushnil(L);
        lua_pushstring(L, "out of memory");
        return 2;
    }

    /* We need to store doc_id copies because cursor pointers are transient.
       Use a Lua table to hold doc_id strings (keeps them GC-rooted). */
    lua_newtable(L); /* doc_id string cache at stack top */
    int docid_cache_idx = lua_gettop(L);
    int docid_count     = 0;

    /* ── Process each query term ─────────────────────────────────── */

    /* Build prefix buffers:
       df_key:      name \x03 <term> \x00 \x00
       posting_pfx: name \x03 <term> \x00 \x01
       doc_meta_pfx: name \x02  */
    uint8_t key_buf[4096 + 256 + 10];

    /* doc_meta prefix: name \x02 */
    uint8_t dm_prefix[4096 + 2];
    memcpy(dm_prefix, name, name_len);
    dm_prefix[name_len]    = 0x02;
    uint16_t dm_prefix_len = (uint16_t)(name_len + 1);

    for (int qi = 1; qi <= q_count; qi++) {
        /* The term stays on the stack for the whole iteration: for a
           number element, lua_tolstring's result is anchored only by
           this slot and could be collected by the pushes below. */
        lua_rawgeti(L, 4, qi);
        size_t term_len;
        const char *term = lua_tolstring(L, -1, &term_len);
        /* Indexed terms are capped at MNEME_FT_MAX_TERM_LEN, so longer
           query terms can't match; this also keeps the key sizes below
           bounded (checked in size_t -- casting to uint16 first can wrap
           past the bounds check and overflow key_buf in the memcpy). */
        if (!term || term_len == 0 || term_len > MNEME_FT_MAX_TERM_LEN ||
            name_len + 1 + term_len + 2 > sizeof(key_buf)) {
            lua_pop(L, 1);
            continue;
        }

        /* df key: name \x03 <term> \x00 \x00 */
        uint16_t df_key_len = (uint16_t)(name_len + 1 + term_len + 2);
        memcpy(key_buf, name, name_len);
        key_buf[name_len] = 0x03;
        memcpy(key_buf + name_len + 1, term, term_len);
        key_buf[name_len + 1 + term_len]     = 0x00;
        key_buf[name_len + 1 + term_len + 1] = 0x00;

        const uint8_t *df_val;
        uint32_t df_val_len;
        uint8_t df_tag;
        int64_t df_ttl;
        rc = mneme_btree_get(ud->txn, ks_root, key_buf, df_key_len, &df_val, &df_val_len, &df_tag, &df_ttl);
        if (rc != 0 || df_val_len < 4) {
            lua_pop(L, 1);
            continue; /* term not in corpus */
        }

        uint32_t df;
        memcpy(&df, df_val, 4);
        if (df == 0) {
            lua_pop(L, 1);
            continue;
        }

        /* IDF = ln((N - df + 0.5) / (df + 0.5) + 1) */
        double idf = log(((N - (double)df + 0.5) / ((double)df + 0.5)) + 1.0);

        /* Posting prefix: name \x03 <term> \x00 \x01 */
        uint16_t post_pfx_len                = df_key_len;
        key_buf[name_len + 1 + term_len + 1] = 0x01;

        /* Scan postings via cursor */
        mneme_cursor_t *cur = NULL;
        if (mneme_cursor_open(ud->txn, ks_root, &cur) < 0) {
            lua_pop(L, 1);
            continue;
        }

        if (mneme_cursor_seek(cur, key_buf, post_pfx_len) < 0 || !cur->valid) {
            mneme_cursor_close(cur);
            lua_pop(L, 1);
            continue;
        }

        while (cur->valid) {
            const uint8_t *ck;
            uint16_t ck_len;
            if (mneme_cursor_key(cur, &ck, &ck_len) < 0)
                break;

            /* Check prefix match */
            if (ck_len < post_pfx_len || memcmp(ck, key_buf, post_pfx_len) != 0)
                break;

            /* Extract doc_id: everything after the posting prefix */
            const uint8_t *doc_id = ck + post_pfx_len;
            uint16_t doc_id_len   = ck_len - post_pfx_len;

            /* Read tf from value */
            const uint8_t *tv;
            uint32_t tv_len;
            uint8_t tv_tag;
            int64_t tv_ttl;
            if (mneme_cursor_val(cur, &tv, &tv_len, &tv_tag, &tv_ttl) < 0 || tv_len < 4) {
                mneme_cursor_next(cur);
                continue;
            }
            uint32_t tf;
            memcpy(&tf, tv, 4);

            /* Find or create score entry (index-based chaining) */
            score_entry_t *ent = NULL;
            {
                uint32_t _h = mneme_checksum(doc_id, doc_id_len) & ((uint32_t)map_cap - 1);
                int _idx    = map_buckets[_h];
                while (_idx >= 0) {
                    if (arena[_idx].doc_id_len == doc_id_len && memcmp(arena[_idx].doc_id, doc_id, doc_id_len) == 0) {
                        ent = &arena[_idx];
                        break;
                    }
                    _idx = arena[_idx].next;
                }
                if (!ent) {
                    if (arena_count >= arena_cap) {
                        int _nc            = arena_cap * 2;
                        score_entry_t *_na = (score_entry_t *)realloc(arena, (size_t)_nc * sizeof(score_entry_t));
                        if (_na) {
                            arena     = _na;
                            arena_cap = _nc;
                        }
                    }
                    if (arena_count < arena_cap) {
                        int _ni = arena_count++;
                        docid_count++;
                        lua_pushlstring(L, (const char *)doc_id, doc_id_len);
                        lua_rawseti(L, docid_cache_idx, docid_count);
                        lua_rawgeti(L, docid_cache_idx, docid_count);
                        arena[_ni].doc_id = (const uint8_t *)lua_tostring(L, -1);
                        lua_pop(L, 1);
                        arena[_ni].doc_id_len     = doc_id_len;
                        arena[_ni].score          = 0.0;
                        arena[_ni].doc_len        = 0;
                        arena[_ni].doc_len_loaded = 0;
                        arena[_ni].next           = map_buckets[_h];
                        map_buckets[_h]           = _ni;
                        ent                       = &arena[_ni];
                    }
                }
            }
            if (!ent) {
                mneme_cursor_next(cur);
                continue;
            }

            /* Load doc_len on first encounter */
            if (!ent->doc_len_loaded) {
                /* doc_meta key: name \x02 <doc_id> */
                uint8_t dm_key[4096 + 256 + 4];
                uint16_t dm_key_len = (uint16_t)(dm_prefix_len + doc_id_len);
                if (dm_key_len <= sizeof(dm_key)) {
                    memcpy(dm_key, dm_prefix, dm_prefix_len);
                    memcpy(dm_key + dm_prefix_len, doc_id, doc_id_len);

                    const uint8_t *dv;
                    uint32_t dv_len;
                    uint8_t dv_tag;
                    int64_t dv_ttl;
                    if (mneme_btree_get(ud->txn, ks_root, dm_key, dm_key_len, &dv, &dv_len, &dv_tag, &dv_ttl) == 0 &&
                        dv_len >= 4) {
                        memcpy(&ent->doc_len, dv, 4);
                    }
                }
                ent->doc_len_loaded = 1;
            }

            /* BM25 term contribution */
            double dl    = (double)ent->doc_len;
            double tf_d  = (double)tf;
            double num   = tf_d * (k1 + 1.0);
            double denom = tf_d + k1 * (1.0 - b + b * dl / avgdl);
            ent->score += idf * num / denom;

            mneme_cursor_next(cur);
        }
        mneme_cursor_close(cur);
        lua_pop(L, 1); /* the term pushed by lua_rawgeti */
    }

    /* ── Collect results and top-k sort ───────────────────────────── */

    /* Simple approach: collect all scored docs, partial sort for top_k.
       For expected scale (thousands of docs), this is fine. */
    typedef struct {
        double score;
        int arena_idx;
    } result_entry_t;

    int result_count        = arena_count;
    result_entry_t *results = NULL;
    if (result_count > 0) {
        results = (result_entry_t *)malloc((size_t)result_count * sizeof(result_entry_t));
        if (results) {
            for (int i = 0; i < result_count; i++) {
                results[i].score     = arena[i].score;
                results[i].arena_idx = i;
            }

            /* Partial sort: find top_k by descending score.
               Simple insertion sort since result sets are small. */
            int limit = result_count < top_k ? result_count : top_k;
            for (int i = 0; i < limit; i++) {
                int best = i;
                for (int j = i + 1; j < result_count; j++) {
                    if (results[j].score > results[best].score)
                        best = j;
                }
                if (best != i) {
                    result_entry_t tmp = results[i];
                    results[i]         = results[best];
                    results[best]      = tmp;
                }
            }

            /* Build Lua result table */
            lua_createtable(L, limit, 0);
            for (int i = 0; i < limit; i++) {
                if (results[i].score <= 0.0) {
                    /* Truncate at zero/negative scores */
                    break;
                }
                lua_createtable(L, 0, 2);

                score_entry_t *e = &arena[results[i].arena_idx];
                lua_pushlstring(L, (const char *)e->doc_id, e->doc_id_len);
                lua_setfield(L, -2, "key");

                lua_pushnumber(L, (lua_Number)results[i].score);
                lua_setfield(L, -2, "score");

                lua_rawseti(L, -2, i + 1);
            }
            free(results);
        } else {
            lua_createtable(L, 0, 0);
        }
    } else {
        lua_createtable(L, 0, 0);
    }

    free(arena);
    free(map_buckets);
    /* Remove docid_cache table */
    lua_remove(L, docid_cache_idx);
    return 1;
}

/* ── DEK (Data Encryption Key) secure userdata ─────────────────────── */

#define MNEME_DEK_SIZE 32

typedef struct {
    uint8_t key[MNEME_DEK_SIZE];
    int locked; /* 1 if mlock'd */
} mneme_dek_t;

/* dek_new(key_string_32) -> userdata */
static int l_dek_new(lua_State *L) {
    size_t key_len;
    const char *key = luaL_checklstring(L, 1, &key_len);
    if (key_len != MNEME_DEK_SIZE) {
        lua_pushnil(L);
        lua_pushstring(L, "DEK must be 32 bytes");
        return 2;
    }

    mneme_dek_t *dek = (mneme_dek_t *)lua_newuserdata(L, sizeof(*dek));
    memcpy(dek->key, key, MNEME_DEK_SIZE);
    dek->locked = (mlock(dek->key, MNEME_DEK_SIZE) == 0) ? 1 : 0;

    luaL_getmetatable(L, MNEME_DEK_MT);
    lua_setmetatable(L, -2);
    return 1;
}

/* dek:encrypt(nonce12, aad, plaintext) -> ciphertext_with_tag
   Calls ChaCha20-Poly1305 using the DEK directly from locked memory. */
static int l_dek_encrypt(lua_State *L) {
    mneme_dek_t *dek = (mneme_dek_t *)luaL_checkudata(L, 1, MNEME_DEK_MT);
    size_t nonce_len, aad_len, pt_len;
    const char *nonce = luaL_checklstring(L, 2, &nonce_len);
    const char *aad   = luaL_checklstring(L, 3, &aad_len);
    const char *pt    = luaL_checklstring(L, 4, &pt_len);
    if (nonce_len != 12) {
        lua_pushnil(L);
        lua_pushstring(L, "nonce must be 12 bytes");
        return 2;
    }

    extern int litls_chacha20_poly1305_encrypt(const uint8_t key[32], const uint8_t nonce[12], const uint8_t *aad,
                                               size_t aad_len, const uint8_t *plaintext, size_t pt_len,
                                               uint8_t *ciphertext, uint8_t tag[16]);

    /* Allocate buffer: ciphertext + 16-byte tag */
    size_t out_len = pt_len + 16;
    uint8_t *buf   = (uint8_t *)malloc(out_len);
    if (!buf)
        return luaL_error(L, "out of memory");
    uint8_t tag[16];

    extern void litls_secure_zero(void *buf, size_t len);
    if (litls_chacha20_poly1305_encrypt(dek->key, (const uint8_t *)nonce, (const uint8_t *)aad, aad_len,
                                        (const uint8_t *)pt, pt_len, buf, tag) != 0) {
        litls_secure_zero(buf, out_len);
        litls_secure_zero(tag, 16);
        free(buf);
        return luaL_error(L, "encryption failed");
    }
    memcpy(buf + pt_len, tag, 16);

    lua_pushlstring(L, (const char *)buf, out_len);
    litls_secure_zero(buf, out_len);
    free(buf);
    return 1;
}

/* dek:decrypt(nonce12, aad, ciphertext_with_tag) -> plaintext or nil, err */
static int l_dek_decrypt(lua_State *L) {
    mneme_dek_t *dek = (mneme_dek_t *)luaL_checkudata(L, 1, MNEME_DEK_MT);
    size_t nonce_len, aad_len, ct_len;
    const char *nonce = luaL_checklstring(L, 2, &nonce_len);
    const char *aad   = luaL_checklstring(L, 3, &aad_len);
    const char *ct    = luaL_checklstring(L, 4, &ct_len);
    if (nonce_len != 12) {
        lua_pushnil(L);
        lua_pushstring(L, "nonce must be 12 bytes");
        return 2;
    }
    if (ct_len < 16) {
        lua_pushnil(L);
        lua_pushstring(L, "ciphertext too short");
        return 2;
    }

    extern int litls_chacha20_poly1305_decrypt(const uint8_t key[32], const uint8_t nonce[12], const uint8_t *aad,
                                               size_t aad_len, const uint8_t *ciphertext, size_t ct_len,
                                               const uint8_t tag[16], uint8_t *plaintext);

    size_t pt_len      = ct_len - 16;
    const uint8_t *tag = (const uint8_t *)ct + pt_len;
    uint8_t *buf       = (uint8_t *)malloc(pt_len > 0 ? pt_len : 1);
    if (!buf)
        return luaL_error(L, "out of memory");

    extern void litls_secure_zero(void *buf, size_t len);
    int rc = litls_chacha20_poly1305_decrypt(dek->key, (const uint8_t *)nonce, (const uint8_t *)aad, aad_len,
                                             (const uint8_t *)ct, pt_len, tag, buf);
    if (rc != 0) {
        litls_secure_zero(buf, pt_len);
        free(buf);
        lua_pushnil(L);
        lua_pushstring(L, "decryption failed (authentication error)");
        return 2;
    }

    lua_pushlstring(L, (const char *)buf, pt_len);
    litls_secure_zero(buf, pt_len);
    free(buf);
    return 1;
}

/* dek:derive(info_string) -> new dek userdata
   HKDF-Expand using the DEK as PRK and info as context. */
static int l_dek_derive(lua_State *L) {
    mneme_dek_t *dek = (mneme_dek_t *)luaL_checkudata(L, 1, MNEME_DEK_MT);
    size_t info_len;
    const char *info = luaL_checklstring(L, 2, &info_len);

    uint8_t subkey[MNEME_DEK_SIZE] = {0};

    extern int litls_hkdf_expand_sha256(const uint8_t prk[32], const uint8_t *info, size_t info_len, uint8_t *okm,
                                        size_t okm_len);

    if (litls_hkdf_expand_sha256(dek->key, (const uint8_t *)info, info_len, subkey, MNEME_DEK_SIZE) < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "hkdf expand failed");
        return 2;
    }

    mneme_dek_t *sub = (mneme_dek_t *)lua_newuserdata(L, sizeof(*sub));
    memcpy(sub->key, subkey, MNEME_DEK_SIZE);
    sub->locked = (mlock(sub->key, MNEME_DEK_SIZE) == 0) ? 1 : 0;

    extern void litls_secure_zero(void *buf, size_t len);
    litls_secure_zero(subkey, MNEME_DEK_SIZE);

    luaL_getmetatable(L, MNEME_DEK_MT);
    lua_setmetatable(L, -2);
    return 1;
}

static int l_dek_gc(lua_State *L) {
    mneme_dek_t *dek = (mneme_dek_t *)luaL_checkudata(L, 1, MNEME_DEK_MT);

    extern void litls_secure_zero(void *buf, size_t len);
    litls_secure_zero(dek->key, MNEME_DEK_SIZE);

    /* Deliberately no munlock: it operates on whole pages and the kernel
       does not refcount lock state, so unlocking here would drop the lock
       for any other DEK sharing this page. The kernel releases the locks
       when the process exits. */
    dek->locked = 0;
    return 0;
}

static const luaL_Reg dek_methods[] = {
    {"encrypt", l_dek_encrypt},
    {"decrypt", l_dek_decrypt},
    {"derive",  l_dek_derive },
    {NULL,      NULL         }
};

static const luaL_Reg module_functions[] = {
    {"db_open",         l_db_open        },
    {"dek_new",         l_dek_new        },
    {"pack_int64",      l_pack_int64     },
    {"unpack_int64",    l_unpack_int64   },
    {"pack_double",     l_pack_double    },
    {"unpack_double",   l_unpack_double  },
    {"encode_score",    l_encode_score   },
    {"decode_score",    l_decode_score   },
    {"pack_int64_be",   l_pack_int64_be  },
    {"unpack_int64_be", l_unpack_int64_be},
    {"pack_floats",     l_pack_floats    },
    {"unpack_floats",   l_unpack_floats  },
    {"vec_dim",         l_vec_dim        },
    {"vsearch_scan",    l_vsearch_scan   },
    {"pack_u32",        l_pack_u32       },
    {"unpack_u32",      l_unpack_u32     },
    {"ft_analyze",      l_ft_analyze     },
    {"ft_search",       l_ft_search      },
    {NULL,              NULL             }
};

int luaopen_mneme_core(lua_State *L) {
    /* DB metatable */
    luaL_newmetatable(L, MNEME_DB_MT);

    lua_newtable(L);
    luaL_register(L, NULL, db_methods);
    lua_setfield(L, -2, "__index");

    lua_pushcfunction(L, l_db_close);
    lua_setfield(L, -2, "__gc");

    lua_pop(L, 1);

    /* TXN metatable */
    luaL_newmetatable(L, MNEME_TXN_MT);

    lua_newtable(L);
    luaL_register(L, NULL, txn_methods);
    lua_setfield(L, -2, "__index");

    lua_pushcfunction(L, l_txn_gc);
    lua_setfield(L, -2, "__gc");

    lua_pop(L, 1);

    /* Cursor metatable */
    luaL_newmetatable(L, MNEME_CUR_MT);

    lua_newtable(L);
    luaL_register(L, NULL, cursor_methods);
    lua_setfield(L, -2, "__index");

    lua_pushcfunction(L, l_cur_close);
    lua_setfield(L, -2, "__gc");

    lua_pop(L, 1);

    /* DEK metatable */
    luaL_newmetatable(L, MNEME_DEK_MT);

    lua_newtable(L);
    luaL_register(L, NULL, dek_methods);
    lua_setfield(L, -2, "__index");

    lua_pushcfunction(L, l_dek_gc);
    lua_setfield(L, -2, "__gc");

    lua_pop(L, 1);

    /* Module table */
    luaL_register(L, "mneme.core", module_functions);
    return 1;
}