// 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 <stdlib.h>
#include <string.h>

/* ── Constants ─────────────────────────────────────────────────────── */

#define LCELL_HDR MNEME_LCELL_HDR
#define BCELL_HDR MNEME_BCELL_HDR

#define OVERFLOW_FLAG MNEME_OVERFLOW_FLAG

/* Inline size of an overflow reference */
#define OVFL_PTR_LEN 4

/* Max cell body (key + inline val) to guarantee 2 cells per page */
#define MAX_CELL_BODY (((MNEME_PAGE_SIZE - MNEME_PAGE_HDR_SIZE) / 2) - 2 - LCELL_HDR)

/* Overflow page layout */
#define OVFL_HDR_SIZE  8
#define OVFL_DATA_SIZE (MNEME_PAGE_SIZE - OVFL_HDR_SIZE)

/* Max path depth */
#define MAX_PATH 32

/* Path entry for tree traversal */
typedef struct {
    uint32_t pgno;
    int child_idx;
} path_entry_t;

/* Use shared inline helpers from mneme.h */
#define keycmp   mneme_keycmp
#define cptr_get mneme_cptr_get

/* ── Cell pointer access (write) ──────────────────────────────────── */

static void cptr_set(uint8_t *page, int idx, uint16_t off) {
    memcpy(page + MNEME_PAGE_HDR_SIZE + idx * 2, &off, 2);
}

/* ── Page init ─────────────────────────────────────────────────────── */

static void page_init(uint8_t *page, uint8_t type) {
    memset(page, 0, MNEME_PAGE_SIZE);
    mneme_page_hdr_t *hdr = (mneme_page_hdr_t *)page;
    hdr->page_type        = type;
    hdr->key_count        = 0;
    hdr->free_offset      = MNEME_PAGE_HDR_SIZE;
    hdr->cell_area_end    = MNEME_PAGE_SIZE;
    hdr->right_ptr        = 0;
}

/* ── Leaf cell helpers ─────────────────────────────────────────────── */

static void lcell_read(const uint8_t *page, int idx, const uint8_t **key, uint16_t *key_len, const uint8_t **val,
                       uint32_t *val_len, uint8_t *type_tag, int64_t *ttl) {
    uint16_t off     = cptr_get(page, idx);
    const uint8_t *c = page + off;
    memcpy(key_len, c, 2);
    c += 2;
    memcpy(val_len, c, 4);
    c += 4;
    memcpy(ttl, c, 8);
    c += 8;
    *type_tag = *c;
    c += 1;
    *key = c;
    c += *key_len;
    *val = c;
}

static int lcell_stored_size(const uint8_t *page, uint16_t off) {
    uint16_t kl;
    uint32_t vl;
    memcpy(&kl, page + off, 2);
    memcpy(&vl, page + off + 2, 4);
    uint8_t tag = page[off + 14];
    int ivl     = (tag & OVERFLOW_FLAG) ? OVFL_PTR_LEN : (int)vl;
    return LCELL_HDR + kl + ivl;
}

/* Use shared page search helpers from mneme.h */
#define leaf_search   mneme_leaf_search
#define branch_search mneme_branch_search
#define branch_child  mneme_branch_child

/* ── Available space in a page ─────────────────────────────────────── */

static int page_avail(const uint8_t *page) {
    const mneme_page_hdr_t *hdr = (const mneme_page_hdr_t *)page;
    return (int)hdr->cell_area_end - (int)hdr->free_offset;
}

/* ── Generic cell insert into page ─────────────────────────────────── */

/* Insert raw cell data into a slotted page at pointer index idx.
   Returns 0 on success, -1 if no space. */
static int cell_insert(uint8_t *page, int idx, const uint8_t *cell_data, int cell_size) {
    mneme_page_hdr_t *hdr = (mneme_page_hdr_t *)page;
    if (cell_size + 2 > page_avail(page))
        return -1;

    uint16_t cell_off = hdr->cell_area_end - (uint16_t)cell_size;
    memcpy(page + cell_off, cell_data, cell_size);

    for (int i = hdr->key_count; i > idx; i--)
        cptr_set(page, i, cptr_get(page, i - 1));
    cptr_set(page, idx, cell_off);

    hdr->key_count++;
    hdr->cell_area_end = cell_off;
    hdr->free_offset   = MNEME_PAGE_HDR_SIZE + hdr->key_count * 2;
    return 0;
}

/* Remove cell pointer at index idx (does not reclaim cell data space) */
static void cell_remove(uint8_t *page, int idx) {
    mneme_page_hdr_t *hdr = (mneme_page_hdr_t *)page;
    for (int i = idx; i < hdr->key_count - 1; i++)
        cptr_set(page, i, cptr_get(page, i + 1));
    hdr->key_count--;
    hdr->free_offset = MNEME_PAGE_HDR_SIZE + hdr->key_count * 2;
}

/* ── Build leaf cell data ──────────────────────────────────────────── */

/* Write a leaf cell into buf. Returns cell size. */
static int build_lcell(uint8_t *buf, const uint8_t *key, uint16_t key_len, const uint8_t *inline_val,
                       int inline_val_len, uint32_t real_val_len, uint8_t type_tag, int64_t ttl) {
    uint8_t *p = buf;
    memcpy(p, &key_len, 2);
    p += 2;
    memcpy(p, &real_val_len, 4);
    p += 4;
    memcpy(p, &ttl, 8);
    p += 8;
    *p = type_tag;
    p += 1;
    memcpy(p, key, key_len);
    p += key_len;
    memcpy(p, inline_val, inline_val_len);
    return LCELL_HDR + key_len + inline_val_len;
}

/* ── Build branch cell data ────────────────────────────────────────── */

static int build_bcell(uint8_t *buf, uint32_t child_pgno, const uint8_t *key, uint16_t key_len) {
    memcpy(buf, &child_pgno, 4);
    memcpy(buf + 4, &key_len, 2);
    memcpy(buf + BCELL_HDR, key, key_len);
    return BCELL_HDR + key_len;
}

/* ── Page compact ──────────────────────────────────────────────────── */

/* Generic compact: rebuild cell area without holes. Works for both
   leaf and branch pages. */
static void page_compact(uint8_t *page) {
    mneme_page_hdr_t *hdr = (mneme_page_hdr_t *)page;
    uint8_t tmp[MNEME_PAGE_SIZE];
    uint16_t off = MNEME_PAGE_SIZE;

    for (int i = 0; i < hdr->key_count; i++) {
        uint16_t co = cptr_get(page, i);
        int sz;
        if (hdr->page_type == MNEME_PAGE_LEAF) {
            sz = lcell_stored_size(page, co);
        } else {
            uint16_t kl;
            memcpy(&kl, page + co + 4, 2);
            sz = BCELL_HDR + kl;
        }
        off -= (uint16_t)sz;
        memcpy(tmp + off, page + co, sz);
        cptr_set(page, i, off);
    }

    memcpy(page + off, tmp + off, MNEME_PAGE_SIZE - off);
    hdr->cell_area_end = off;
    hdr->free_offset   = MNEME_PAGE_HDR_SIZE + hdr->key_count * 2;
}

/* ── COW a page ────────────────────────────────────────────────────── */

static uint32_t cow_page(mneme_txn_t *txn, uint32_t old_pgno, uint8_t **out) {
    uint8_t *existing = mneme_dirty_find(txn, old_pgno);
    if (existing) {
        *out = existing;
        return old_pgno;
    }

    uint32_t new_pgno = mneme_page_alloc(txn);
    if (new_pgno == 0)
        return 0;

    uint8_t *np       = mneme_dirty_find(txn, new_pgno);
    const uint8_t *op = mneme_page_read(txn->db, old_pgno);
    if (op)
        memcpy(np, op, MNEME_PAGE_SIZE);

    mneme_page_free(txn, old_pgno);

    /* Re-lookup: page_free may have realloc'd the dirty array */
    np   = mneme_dirty_find(txn, new_pgno);
    *out = np;
    return new_pgno;
}

/* ── Overflow operations ───────────────────────────────────────────── */

uint32_t mneme_overflow_write(mneme_txn_t *txn, const uint8_t *data, uint32_t len) {
    uint32_t first = 0, prev = 0, written = 0;
    while (written < len) {
        uint32_t pgno = mneme_page_alloc(txn);
        if (pgno == 0)
            return 0;
        uint8_t *page = mneme_dirty_find(txn, pgno);
        memset(page, 0, MNEME_PAGE_SIZE);
        page[0]        = MNEME_PAGE_OVERFLOW;
        uint32_t chunk = len - written;
        if (chunk > OVFL_DATA_SIZE)
            chunk = OVFL_DATA_SIZE;
        memcpy(page + OVFL_HDR_SIZE, data + written, chunk);
        written += chunk;
        if (first == 0)
            first = pgno;
        if (prev != 0) {
            uint8_t *pp = mneme_dirty_find(txn, prev);
            memcpy(pp + 4, &pgno, 4);
        }
        prev = pgno;
    }
    return first;
}

int mneme_overflow_read(mneme_txn_t *txn, uint32_t pgno, uint8_t *buf, uint32_t buf_len, uint32_t *bytes_read) {
    uint32_t total = 0;
    while (pgno != 0 && total < buf_len) {
        const uint8_t *page = mneme_page_resolve(txn, pgno);
        if (!page)
            return -1;
        uint32_t chunk = buf_len - total;
        if (chunk > OVFL_DATA_SIZE)
            chunk = OVFL_DATA_SIZE;
        memcpy(buf + total, page + OVFL_HDR_SIZE, chunk);
        total += chunk;
        memcpy(&pgno, page + 4, 4);
    }
    if (bytes_read)
        *bytes_read = total;
    return 0;
}

int mneme_overflow_free(mneme_txn_t *txn, uint32_t pgno) {
    while (pgno != 0) {
        const uint8_t *page = mneme_page_resolve(txn, pgno);
        if (!page)
            return -1;
        uint32_t next;
        memcpy(&next, page + 4, 4);
        mneme_page_free(txn, pgno);
        pgno = next;
    }
    return 0;
}

/* ── Propagate COW from level `from` up to root ───────────────────── */

static uint32_t propagate_cow(mneme_txn_t *txn, path_entry_t *path, int from, uint32_t child_pgno) {
    for (int i = from; i >= 0; i--) {
        uint8_t *parent;
        uint32_t ppgno = cow_page(txn, path[i].pgno, &parent);
        if (ppgno == 0)
            return 0;
        int cidx               = path[i].child_idx;
        mneme_page_hdr_t *phdr = (mneme_page_hdr_t *)parent;
        if (cidx < (int)phdr->key_count) {
            uint16_t coff = cptr_get(parent, cidx);
            memcpy(parent + coff, &child_pgno, 4);
        } else {
            phdr->right_ptr = child_pgno;
        }
        child_pgno   = ppgno;
        path[i].pgno = ppgno;
    }
    return child_pgno;
}

/* ── Leaf split ────────────────────────────────────────────────────── */

/* Split a leaf page, distributing old cells + one new cell between
   left and right pages. Cells are variable-size, so the split point is
   chosen by ACCUMULATED BYTES, not cell count — a count split can pile
   more bytes on one side than a page holds and cell_insert would fail,
   silently dropping a key. The left page takes cells until it reaches
   half of the total payload (or the next cell would not fit), the rest
   go right; every insert is still checked so an unfittable layout
   (pathological near-page-size keys) fails the put instead of losing
   data. A rightmost insert keeps the fast path: the full old page stays
   left, the new cell starts the right page (sequential-append density).
   Returns separator key info for the parent. */
static int leaf_split(mneme_txn_t *txn, uint8_t *left_page, uint32_t left_pgno, int ins_idx, const uint8_t *new_cell,
                      int new_cell_size, uint32_t *right_pgno_out, uint8_t *sep_key_buf, uint16_t *sep_key_len_out) {
    mneme_page_hdr_t *hdr = (mneme_page_hdr_t *)left_page;
    int old_count         = hdr->key_count;
    int total             = old_count + 1;

    /* Allocate right page */
    uint32_t rpgno = mneme_page_alloc(txn);
    if (rpgno == 0)
        return -1;
    /* Re-resolve left_page: page_alloc may have realloc'd the dirty array */
    left_page      = mneme_dirty_find(txn, left_pgno);
    uint8_t *rpage = mneme_dirty_find(txn, rpgno);
    page_init(rpage, MNEME_PAGE_LEAF);

    if (ins_idx >= old_count) {
        /* Rightmost insert: old page untouched, new cell alone on the right */
        if (cell_insert(rpage, 0, new_cell, new_cell_size) < 0)
            return -1;
    } else {
        /* Save old cells */
        uint8_t save[MNEME_PAGE_SIZE];
        memcpy(save, left_page, MNEME_PAGE_SIZE);
        page_init(left_page, MNEME_PAGE_LEAF);

        /* Total stored bytes (each cell + its 2-byte pointer slot) */
        int total_bytes = new_cell_size + 2;
        for (int i = 0; i < old_count; i++)
            total_bytes += lcell_stored_size(save, cptr_get(save, i)) + 2;
        int half     = total_bytes / 2;
        int capacity = MNEME_PAGE_SIZE - MNEME_PAGE_HDR_SIZE;

        int cell_i = 0, left_bytes = 0, going_right = 0;
        for (int i = 0; i < total; i++) {
            const uint8_t *cd;
            int sz;
            if (i == ins_idx) {
                cd = new_cell;
                sz = new_cell_size;
            } else {
                uint16_t off = cptr_get(save, cell_i);
                sz           = lcell_stored_size(save, off);
                cd           = save + off;
                cell_i++;
            }
            /* key order: once a cell has gone right, everything after follows */
            if (!going_right && (left_bytes >= half || left_bytes + sz + 2 > capacity || i == total - 1))
                going_right = 1;
            uint8_t *target      = going_right ? rpage : left_page;
            mneme_page_hdr_t *th = (mneme_page_hdr_t *)target;
            if (cell_insert(target, th->key_count, cd, sz) < 0)
                return -1;
            if (!going_right)
                left_bytes += sz + 2;
        }
    }

    /* Separator = first key of right page */
    mneme_page_hdr_t *rh = (mneme_page_hdr_t *)rpage;
    if (rh->key_count > 0) {
        uint16_t off = cptr_get(rpage, 0);
        uint16_t kl;
        memcpy(&kl, rpage + off, 2);
        memcpy(sep_key_buf, rpage + off + LCELL_HDR, kl);
        *sep_key_len_out = kl;
    }

    *right_pgno_out = rpgno;
    return 0;
}

/* ── Branch: insert separator from child split ─────────────────────── */

/* Insert a separator into a branch after a child split.
   Updates child pointer at cidx to left_pgno.
   Inserts sep_key with right_pgno as its child.

   Branch cell layout: cell[i] = (child_i, key_i)
   child_i handles keys < key_i. right_ptr handles keys >= last key.

   When child at cidx splits into left/right with separator sep:
     cidx < key_count: replace cell[cidx] with (left, sep), insert (right, old_key) at cidx+1
     cidx == key_count: insert (left, sep) at end, right_ptr = right

   Returns 0 on success, -1 if branch needs splitting. */
static int branch_insert_sep(uint8_t *parent, int cidx, uint32_t left_pgno, uint32_t right_pgno, const uint8_t *sep_key,
                             uint16_t sep_key_len) {
    mneme_page_hdr_t *phdr = (mneme_page_hdr_t *)parent;

    if (cidx == (int)phdr->key_count) {
        /* Followed right_ptr → insert (left, sep), update right_ptr */
        uint8_t cell[BCELL_HDR + MNEME_MAX_KEY_LEN];
        int sz = build_bcell(cell, left_pgno, sep_key, sep_key_len);
        if (cell_insert(parent, phdr->key_count, cell, sz) < 0)
            return -1;
        phdr->right_ptr = right_pgno;
        return 0;
    }

    /* cidx < key_count: save old key, remove old cell, insert two new cells */
    uint16_t coff = cptr_get(parent, cidx);
    uint16_t old_kl;
    memcpy(&old_kl, parent + coff + 4, 2);
    uint8_t old_key[MNEME_MAX_KEY_LEN];
    memcpy(old_key, parent + coff + BCELL_HDR, old_kl);

    /* Check if both new cells fit (accounting for removed cell space) */
    int new_sz1 = BCELL_HDR + sep_key_len;
    int new_sz2 = BCELL_HDR + old_kl;

    /* Remove old cell and compact to reclaim space */
    cell_remove(parent, cidx);
    page_compact(parent);

    if (page_avail(parent) < new_sz1 + 2 + new_sz2 + 2)
        return -1;

    uint8_t cell1[BCELL_HDR + MNEME_MAX_KEY_LEN];
    int sz1 = build_bcell(cell1, left_pgno, sep_key, sep_key_len);
    cell_insert(parent, cidx, cell1, sz1);

    uint8_t cell2[BCELL_HDR + MNEME_MAX_KEY_LEN];
    int sz2 = build_bcell(cell2, right_pgno, old_key, old_kl);
    cell_insert(parent, cidx + 1, cell2, sz2);

    return 0;
}

/* ── Branch split ──────────────────────────────────────────────────── */

/* Split a branch page. The caller has already tried branch_insert_sep
   and it failed (no space). We need to:
   1. Collect all cells from the old branch + the new separator
   2. Split around the midpoint, promoting the middle key
   3. Return the promoted key and two child pgnos

   parent_save: saved copy of the branch BEFORE the failed insert attempt
   (since branch_insert_sep may have partially modified the page). */
static int branch_split(mneme_txn_t *txn, const uint8_t *parent_save, uint32_t parent_pgno, int cidx,
                        uint32_t left_pgno, uint32_t right_pgno, const uint8_t *sep_key, uint16_t sep_key_len,
                        uint32_t *new_left_out, uint32_t *new_right_out, uint8_t *promoted_key,
                        uint16_t *promoted_key_len) {
    const mneme_page_hdr_t *phdr = (const mneme_page_hdr_t *)parent_save;
    int old_count                = phdr->key_count;

    /* Build virtual cell list: old cells with modifications + new cell */
    /* For cidx < old_count: replace cell[cidx] → (left,sep) + (right,old_key) */
    /* For cidx == old_count: append (left,sep), right_ptr=right */

    /* Count total virtual cells */
    int total = old_count + 1;
    int split = total / 2;

    /* Allocate left and right pages */
    uint32_t lpgno = mneme_page_alloc(txn);
    if (lpgno == 0)
        return -1;
    uint8_t *lpage = mneme_dirty_find(txn, lpgno);
    page_init(lpage, MNEME_PAGE_BRANCH);

    uint32_t rpgno = mneme_page_alloc(txn);
    if (rpgno == 0)
        return -1;
    /* Re-lookup lpage: rpgno alloc may have realloc'd the dirty array */
    lpage          = mneme_dirty_find(txn, lpgno);
    uint8_t *rpage = mneme_dirty_find(txn, rpgno);
    page_init(rpage, MNEME_PAGE_BRANCH);

    /* Iterate through virtual cells */
    int vi            = 0; /* virtual index */
    int oi            = 0; /* old cell index */
    int promoted_done = 0;

    uint32_t old_right_ptr = phdr->right_ptr;

    /* For cidx < old_count: save old key at cidx */
    uint8_t old_key_at_cidx[MNEME_MAX_KEY_LEN];
    uint16_t old_kl_at_cidx = 0;
    if (cidx < old_count) {
        uint16_t coff = cptr_get(parent_save, cidx);
        memcpy(&old_kl_at_cidx, parent_save + coff + 4, 2);
        memcpy(old_key_at_cidx, parent_save + coff + BCELL_HDR, old_kl_at_cidx);
    }

    int first_repl_done = 0;

    while (vi < total) {
        uint32_t cpgno;
        const uint8_t *ck;
        uint16_t ckl;

        if (cidx < old_count) {
            /* Replace cell[cidx] with two virtual cells */
            if (oi == cidx && !first_repl_done) {
                /* First replacement: (left, sep) */
                cpgno           = left_pgno;
                ck              = sep_key;
                ckl             = sep_key_len;
                first_repl_done = 1;
                /* Don't advance oi yet */
            } else if (oi == cidx && first_repl_done) {
                /* Second replacement: (right, old_key) */
                cpgno = right_pgno;
                ck    = old_key_at_cidx;
                ckl   = old_kl_at_cidx;
                oi++; /* now advance past old cell[cidx] */
            } else {
                uint16_t off = cptr_get(parent_save, oi);
                memcpy(&cpgno, parent_save + off, 4);
                memcpy(&ckl, parent_save + off + 4, 2);
                ck = parent_save + off + BCELL_HDR;
                oi++;
            }
        } else {
            /* cidx == old_count: append at end */
            if (vi == old_count) {
                cpgno = left_pgno;
                ck    = sep_key;
                ckl   = sep_key_len;
            } else {
                uint16_t off = cptr_get(parent_save, oi);
                memcpy(&cpgno, parent_save + off, 4);
                memcpy(&ckl, parent_save + off + 4, 2);
                ck = parent_save + off + BCELL_HDR;
                oi++;
            }
        }

        if (vi == split) {
            /* Promoted key */
            memcpy(promoted_key, ck, ckl);
            *promoted_key_len                      = ckl;
            ((mneme_page_hdr_t *)lpage)->right_ptr = cpgno;
            promoted_done                          = 1;
        } else if (!promoted_done) {
            uint8_t cell[BCELL_HDR + MNEME_MAX_KEY_LEN];
            int sz               = build_bcell(cell, cpgno, ck, ckl);
            mneme_page_hdr_t *lh = (mneme_page_hdr_t *)lpage;
            if (cell_insert(lpage, lh->key_count, cell, sz) < 0)
                return -1; /* unfittable layout: fail the put, never drop a cell */
        } else {
            uint8_t cell[BCELL_HDR + MNEME_MAX_KEY_LEN];
            int sz               = build_bcell(cell, cpgno, ck, ckl);
            mneme_page_hdr_t *rh = (mneme_page_hdr_t *)rpage;
            if (cell_insert(rpage, rh->key_count, cell, sz) < 0)
                return -1; /* unfittable layout: fail the put, never drop a cell */
        }

        vi++;
    }

    /* Right page's right_ptr */
    if (cidx == old_count) {
        ((mneme_page_hdr_t *)rpage)->right_ptr = right_pgno;
    } else {
        ((mneme_page_hdr_t *)rpage)->right_ptr = old_right_ptr;
    }

    /* Free old parent page */
    mneme_page_free(txn, parent_pgno);

    *new_left_out  = lpgno;
    *new_right_out = rpgno;
    return 0;
}

/* ── B-tree GET ────────────────────────────────────────────────────── */

int mneme_btree_get(mneme_txn_t *txn, uint32_t root_pgno, const uint8_t *key, uint16_t key_len, const uint8_t **val_out,
                    uint32_t *val_len_out, uint8_t *type_tag_out, int64_t *ttl_out) {
    if (key_len == 0 || key_len > MNEME_MAX_KEY_LEN)
        return -1;
    uint32_t pgno = root_pgno;
    for (;;) {
        const uint8_t *page = mneme_page_resolve(txn, pgno);
        if (!page)
            return -1;
        const mneme_page_hdr_t *hdr = (const mneme_page_hdr_t *)page;

        if (hdr->page_type == MNEME_PAGE_LEAF) {
            int found;
            int idx = leaf_search(page, key, key_len, &found);
            if (!found)
                return 1; /* not found */
            const uint8_t *k, *v;
            uint16_t kl;
            uint32_t vl;
            uint8_t tag;
            int64_t ttl;
            lcell_read(page, idx, &k, &kl, &v, &vl, &tag, &ttl);
            *val_out      = v;
            *val_len_out  = vl;
            *type_tag_out = tag;
            *ttl_out      = ttl;
            return 0;
        }

        int cidx = branch_search(page, key, key_len);
        pgno     = branch_child(page, cidx);
    }
}

/* ── B-tree PUT ────────────────────────────────────────────────────── */

uint32_t mneme_btree_put(mneme_txn_t *txn, uint32_t root_pgno, const uint8_t *key, uint16_t key_len, const uint8_t *val,
                         uint32_t val_len, uint8_t type_tag, int64_t ttl_expires) {
    if (key_len == 0 || key_len > MNEME_MAX_KEY_LEN)
        return 0;
    path_entry_t path[MAX_PATH];
    int path_len = 0;

    /* Walk to leaf */
    uint32_t pgno = root_pgno;
    for (;;) {
        const uint8_t *page = mneme_page_resolve(txn, pgno);
        if (!page)
            return 0;
        if (((const mneme_page_hdr_t *)page)->page_type == MNEME_PAGE_LEAF)
            break;
        int cidx = branch_search(page, key, key_len);
        if (path_len >= MAX_PATH)
            return 0;
        path[path_len].pgno      = pgno;
        path[path_len].child_idx = cidx;
        path_len++;
        pgno = branch_child(page, cidx);
    }

    /* COW the leaf */
    uint8_t *leaf;
    uint32_t leaf_pgno = cow_page(txn, pgno, &leaf);
    if (leaf_pgno == 0)
        return 0;

    /* Prepare inline value (possibly overflow) */
    int need_overflow = ((int)key_len + (int)val_len > (int)MAX_CELL_BODY);
    uint8_t tag       = type_tag;
    const uint8_t *ival;
    int ival_len;
    uint32_t ovfl_pgno = 0;

    if (need_overflow) {
        ovfl_pgno = mneme_overflow_write(txn, val, val_len);
        if (ovfl_pgno == 0)
            return 0;
        /* Re-resolve leaf: overflow_write may have realloc'd dirty array */
        leaf = mneme_dirty_find(txn, leaf_pgno);
        tag |= OVERFLOW_FLAG;
        ival     = (const uint8_t *)&ovfl_pgno;
        ival_len = OVFL_PTR_LEN;
    } else {
        ival     = val;
        ival_len = (int)val_len;
    }

    /* Search leaf */
    int found;
    int idx = leaf_search(leaf, key, key_len, &found);

    if (found) {
        /* Update: free old overflow, remove old cell */
        uint16_t off    = cptr_get(leaf, idx);
        uint8_t old_tag = leaf[off + 14];
        if (old_tag & OVERFLOW_FLAG) {
            uint16_t okl;
            memcpy(&okl, leaf + off, 2);
            uint32_t old_ovfl;
            memcpy(&old_ovfl, leaf + off + LCELL_HDR + okl, 4);
            mneme_overflow_free(txn, old_ovfl);
            /* Re-resolve leaf: overflow_free may have realloc'd dirty array */
            leaf = mneme_dirty_find(txn, leaf_pgno);
        }
        cell_remove(leaf, idx);
        page_compact(leaf);
    }

    /* Build new cell */
    uint8_t cell_buf[LCELL_HDR + MNEME_MAX_KEY_LEN + MNEME_PAGE_SIZE];
    int cell_size = build_lcell(cell_buf, key, key_len, ival, ival_len, val_len, tag, ttl_expires);

    /* Try to insert into leaf */
    if (cell_insert(leaf, idx, cell_buf, cell_size) == 0) {
        /* Fits. Propagate COW up. */
        uint32_t new_root = propagate_cow(txn, path, path_len - 1, leaf_pgno);
        return (path_len > 0) ? new_root : leaf_pgno;
    }

    /* Leaf split */
    uint32_t right_pgno;
    uint8_t sep_buf[MNEME_MAX_KEY_LEN];
    uint16_t sep_len;

    if (leaf_split(txn, leaf, leaf_pgno, idx, cell_buf, cell_size, &right_pgno, sep_buf, &sep_len) < 0)
        return 0;

    /* Propagate split up */
    uint32_t left_up  = leaf_pgno;
    uint32_t right_up = right_pgno;
    uint8_t cur_sep[MNEME_MAX_KEY_LEN];
    memcpy(cur_sep, sep_buf, sep_len);
    uint16_t cur_sep_len = sep_len;

    for (int i = path_len - 1; i >= 0; i--) {
        /* Save branch state before modification (for branch_split fallback) */
        uint8_t branch_save[MNEME_PAGE_SIZE];
        const uint8_t *orig = mneme_page_resolve(txn, path[i].pgno);
        if (!orig)
            return 0;
        memcpy(branch_save, orig, MNEME_PAGE_SIZE);

        uint8_t *parent;
        uint32_t ppgno = cow_page(txn, path[i].pgno, &parent);
        if (ppgno == 0)
            return 0;
        path[i].pgno = ppgno;

        if (branch_insert_sep(parent, path[i].child_idx, left_up, right_up, cur_sep, cur_sep_len) == 0) {
            /* Fits. Propagate remaining COW. */
            uint32_t new_root = propagate_cow(txn, path, i - 1, ppgno);
            return (i > 0) ? new_root : ppgno;
        }

        /* Branch split needed */
        uint32_t new_left, new_right;
        uint8_t promoted[MNEME_MAX_KEY_LEN];
        uint16_t promoted_len;

        if (branch_split(txn, branch_save, ppgno, path[i].child_idx, left_up, right_up, cur_sep, cur_sep_len, &new_left,
                         &new_right, promoted, &promoted_len) < 0)
            return 0;

        left_up  = new_left;
        right_up = new_right;
        memcpy(cur_sep, promoted, promoted_len);
        cur_sep_len = promoted_len;
    }

    /* Split reached the root. Create new root. */
    uint32_t new_root = mneme_page_alloc(txn);
    if (new_root == 0)
        return 0;
    uint8_t *rp = mneme_dirty_find(txn, new_root);
    page_init(rp, MNEME_PAGE_BRANCH);
    uint8_t rcell[BCELL_HDR + MNEME_MAX_KEY_LEN];
    int rsz = build_bcell(rcell, left_up, cur_sep, cur_sep_len);
    if (cell_insert(rp, 0, rcell, rsz) < 0)
        return 0;
    ((mneme_page_hdr_t *)rp)->right_ptr = right_up;
    return new_root;
}

/* ── B-tree DELETE ─────────────────────────────────────────────────── */

uint32_t mneme_btree_del(mneme_txn_t *txn, uint32_t root_pgno, const uint8_t *key, uint16_t key_len) {
    if (key_len == 0 || key_len > MNEME_MAX_KEY_LEN)
        return 0;
    path_entry_t path[MAX_PATH];
    int path_len = 0;

    /* Walk to leaf */
    uint32_t pgno = root_pgno;
    for (;;) {
        const uint8_t *page = mneme_page_resolve(txn, pgno);
        if (!page)
            return 0;
        if (((const mneme_page_hdr_t *)page)->page_type == MNEME_PAGE_LEAF)
            break;
        int cidx = branch_search(page, key, key_len);
        if (path_len >= MAX_PATH)
            return 0;
        path[path_len].pgno      = pgno;
        path[path_len].child_idx = cidx;
        path_len++;
        pgno = branch_child(page, cidx);
    }

    /* Search leaf */
    const uint8_t *lpage = mneme_page_resolve(txn, pgno);
    int found;
    int idx = leaf_search(lpage, key, key_len, &found);
    if (!found)
        return root_pgno;

    /* COW the leaf */
    uint8_t *leaf;
    uint32_t leaf_pgno = cow_page(txn, pgno, &leaf);
    if (leaf_pgno == 0)
        return 0;

    /* Free overflow if any */
    uint16_t off = cptr_get(leaf, idx);
    uint8_t tag  = leaf[off + 14];
    if (tag & OVERFLOW_FLAG) {
        uint16_t kl;
        memcpy(&kl, leaf + off, 2);
        uint32_t ovfl;
        memcpy(&ovfl, leaf + off + LCELL_HDR + kl, 4);
        if (mneme_overflow_free(txn, ovfl) < 0)
            return 0;
        /* Re-resolve leaf: overflow_free may have realloc'd dirty array */
        leaf = mneme_dirty_find(txn, leaf_pgno);
        if (!leaf)
            return 0;
    }

    cell_remove(leaf, idx);
    page_compact(leaf);

    /* Propagate COW up */
    uint32_t new_root = propagate_cow(txn, path, path_len - 1, leaf_pgno);
    return (path_len > 0) ? new_root : leaf_pgno;
}

/* ── B-tree FREE (drop a whole subtree) ────────────────────────────── */

int mneme_btree_free(mneme_txn_t *txn, uint32_t root_pgno) {
    if (root_pgno == 0)
        return 0;

    const uint8_t *page = mneme_page_resolve(txn, root_pgno);
    if (!page)
        return -1;
    const mneme_page_hdr_t *hdr = (const mneme_page_hdr_t *)page;
    uint8_t ptype               = hdr->page_type;

    if (ptype == MNEME_PAGE_BRANCH) {
        /* Snapshot all child pgnos before freeing this page (recursion can
           change dirty contents otherwise). */
        int n          = (int)hdr->key_count;
        uint32_t *kids = (uint32_t *)calloc((size_t)(n + 1), sizeof(uint32_t));
        if (!kids)
            return -1;
        for (int i = 0; i < n; i++) {
            uint16_t off = cptr_get(page, i);
            memcpy(&kids[i], page + off, 4);
        }
        kids[n] = hdr->right_ptr;
        for (int i = 0; i <= n; i++) {
            if (mneme_btree_free(txn, kids[i]) < 0) {
                free(kids);
                return -1;
            }
        }
        free(kids);
        mneme_page_free(txn, root_pgno);
        return 0;
    }

    if (ptype == MNEME_PAGE_LEAF) {
        int n = (int)hdr->key_count;
        for (int i = 0; i < n; i++) {
            uint16_t off = cptr_get(page, i);
            uint8_t tag  = page[off + 14];
            if (tag & OVERFLOW_FLAG) {
                uint16_t kl;
                memcpy(&kl, page + off, 2);
                uint32_t ovfl;
                memcpy(&ovfl, page + off + LCELL_HDR + kl, 4);
                mneme_overflow_free(txn, ovfl);
            }
        }
        mneme_page_free(txn, root_pgno);
        return 0;
    }

    /* Unexpected page type — refuse to free. */
    return -1;
}