// 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>

typedef struct {
    uint32_t *pgnos;
    int count;
} retire_val_t;

static int cmp_u32(const void *a, const void *b) {
    uint32_t av = *(const uint32_t *)a;
    uint32_t bv = *(const uint32_t *)b;
    return (av > bv) - (av < bv);
}

static int u32_contains_sorted(const uint32_t *pages, int count, uint32_t pgno) {
    int lo = 0;
    int hi = count;
    while (lo < hi) {
        int mid      = lo + (hi - lo) / 2;
        uint32_t cur = pages[mid];
        if (cur == pgno)
            return 1;
        if (cur < pgno)
            lo = mid + 1;
        else
            hi = mid;
    }
    return 0;
}

void mneme_retire_key_pack(uint8_t out[12], uint64_t txn_id, uint32_t chunk_idx) {
    mneme_be64enc(out, txn_id);
    mneme_be32enc(out + 8, chunk_idx);
}

void mneme_retire_key_unpack(const uint8_t in[12], uint64_t *txn_id, uint32_t *chunk_idx) {
    if (txn_id)
        *txn_id = mneme_be64dec(in);
    if (chunk_idx)
        *chunk_idx = mneme_be32dec(in + 8);
}

static int retire_val_pack(const uint32_t *pgnos, int count, uint8_t **out, uint32_t *out_len) {
    if (count < 0 || count > MNEME_RETIRE_BATCH_PGNOS)
        return -1;
    uint32_t count32 = (uint32_t)count;
    uint32_t len     = 4 + count32 * 4;
    uint8_t *buf     = (uint8_t *)malloc(len);
    if (!buf)
        return -1;
    memcpy(buf, &count32, 4);
    for (int i = 0; i < count; i++)
        memcpy(buf + 4 + i * 4, &pgnos[i], 4);
    *out     = buf;
    *out_len = len;
    return 0;
}

static int retire_val_unpack(const uint8_t *val, uint32_t val_len, retire_val_t *out) {
    out->pgnos = NULL;
    out->count = 0;
    if (val_len < 4)
        return -1;
    uint32_t count;
    memcpy(&count, val, 4);
    if (count == 0 || count > MNEME_RETIRE_BATCH_PGNOS)
        return -1;
    uint32_t expected = 4 + count * 4;
    if (val_len != expected)
        return -1;
    out->pgnos = (uint32_t *)malloc((size_t)count * sizeof(uint32_t));
    if (!out->pgnos)
        return -1;
    out->count = (int)count;
    for (uint32_t i = 0; i < count; i++)
        memcpy(&out->pgnos[i], val + 4 + i * 4, 4);
    return 0;
}

static void retire_val_free(retire_val_t *v) {
    free(v->pgnos);
    v->pgnos = NULL;
    v->count = 0;
}

static uint32_t retire_root_ensure(mneme_txn_t *txn) {
    if (txn->meta.retired_root_pgno != 0)
        return txn->meta.retired_root_pgno;
    uint32_t pgno = mneme_page_alloc(txn);
    if (pgno == 0)
        return 0;
    uint8_t *page = mneme_dirty_find(txn, pgno);
    if (!page)
        return 0;
    memset(page, 0, MNEME_PAGE_SIZE);
    ((mneme_page_hdr_t *)page)->page_type     = MNEME_PAGE_LEAF;
    ((mneme_page_hdr_t *)page)->free_offset   = MNEME_PAGE_HDR_SIZE;
    ((mneme_page_hdr_t *)page)->cell_area_end = MNEME_PAGE_SIZE;
    txn->meta.retired_root_pgno               = pgno;
    return pgno;
}

static int reuse_append(mneme_txn_t *txn, uint32_t pgno) {
    if (txn->reuse_count >= txn->reuse_cap) {
        int new_cap       = txn->reuse_cap == 0 ? 64 : txn->reuse_cap * 2;
        uint32_t *new_arr = (uint32_t *)realloc(txn->reuse_pages, (size_t)new_cap * sizeof(uint32_t));
        if (!new_arr)
            return -1;
        txn->reuse_pages = new_arr;
        txn->reuse_cap   = new_cap;
    }
    txn->reuse_pages[txn->reuse_count++] = pgno;
    return 0;
}

static int harvested_append(mneme_txn_t *txn, uint64_t txn_id, uint32_t chunk_idx, const uint32_t *pages,
                            int page_count) {
    if (txn->harvested_count >= txn->harvested_cap) {
        int new_cap   = txn->harvested_cap == 0 ? 16 : txn->harvested_cap * 2;
        void *new_arr = realloc(txn->harvested, (size_t)new_cap * sizeof(*txn->harvested));
        if (!new_arr)
            return -1;
        txn->harvested     = new_arr;
        txn->harvested_cap = new_cap;
    }
    txn->harvested[txn->harvested_count].txn_id     = txn_id;
    txn->harvested[txn->harvested_count].chunk_idx  = chunk_idx;
    txn->harvested[txn->harvested_count].pages      = NULL;
    txn->harvested[txn->harvested_count].page_count = page_count;
    if (page_count > 0) {
        txn->harvested[txn->harvested_count].pages = (uint32_t *)malloc((size_t)page_count * sizeof(uint32_t));
        if (!txn->harvested[txn->harvested_count].pages)
            return -1;
        memcpy(txn->harvested[txn->harvested_count].pages, pages, (size_t)page_count * sizeof(uint32_t));
    }
    txn->harvested_count++;
    return 0;
}

int mneme_retire_harvest_safe(mneme_txn_t *txn, int page_budget) {
    if (!txn || page_budget <= 0)
        return 0;
    if (txn->meta.retired_page_count == 0 || txn->meta.retired_root_pgno == 0)
        return 0;

    mneme_cursor_t *cur = NULL;
    if (mneme_cursor_open(txn, txn->meta.retired_root_pgno, &cur) < 0)
        return -1;
    if (mneme_cursor_first(cur) < 0 || !cur->valid) {
        mneme_cursor_close(cur);
        return 0;
    }

    int harvested_pages = 0;
    while (cur->valid && harvested_pages < page_budget) {
        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)
            break;
        if (key_len != 12)
            break;
        if ((type_tag & MNEME_OVERFLOW_FLAG) != 0)
            break;

        uint64_t retire_txn_id;
        uint32_t chunk_idx;
        mneme_retire_key_unpack(key, &retire_txn_id, &chunk_idx);
        if (retire_txn_id == 0)
            break;
        int probe = mneme_reader_probe_older(txn->db, retire_txn_id - 1);
        if (probe != 0)
            break;

        retire_val_t rv = {0};
        if (retire_val_unpack(val, val_len, &rv) < 0)
            break;

        int take = rv.count;
        if (take > page_budget - harvested_pages)
            take = page_budget - harvested_pages;
        for (int i = 0; i < take; i++) {
            if (reuse_append(txn, rv.pgnos[i]) < 0) {
                retire_val_free(&rv);
                mneme_cursor_close(cur);
                return -1;
            }
        }
        if (harvested_append(txn, retire_txn_id, chunk_idx, rv.pgnos, rv.count) < 0) {
            retire_val_free(&rv);
            mneme_cursor_close(cur);
            return -1;
        }
        harvested_pages += take;
        retire_val_free(&rv);
        mneme_cursor_next(cur);
    }
    mneme_cursor_close(cur);
    return harvested_pages;
}

int mneme_retire_apply_harvested(mneme_txn_t *txn) {
    if (!txn || txn->harvested_count == 0 || txn->meta.retired_root_pgno == 0)
        return 0;

    uint64_t consumed_total = 0;
    if (txn->consumed_reuse_count > 1)
        qsort(txn->consumed_reuse_pages, (size_t)txn->consumed_reuse_count, sizeof(uint32_t), cmp_u32);

    uint32_t *saved_reuse_pages = txn->reuse_pages;
    int saved_reuse_count       = txn->reuse_count;
    int saved_reuse_cap         = txn->reuse_cap;
    txn->reuse_pages            = NULL;
    txn->reuse_count            = 0;
    txn->reuse_cap              = 0;

    for (int i = 0; i < txn->harvested_count; i++) {
        uint8_t key[12];
        mneme_retire_key_pack(key, txn->harvested[i].txn_id, txn->harvested[i].chunk_idx);

        const uint8_t *val;
        uint32_t val_len;
        uint8_t type_tag;
        int64_t ttl;
        int rc = mneme_btree_get(txn, txn->meta.retired_root_pgno, key, sizeof(key), &val, &val_len, &type_tag, &ttl);
        if (rc < 0) {
            txn->reuse_pages = saved_reuse_pages;
            txn->reuse_count = saved_reuse_count;
            txn->reuse_cap   = saved_reuse_cap;
            return -1;
        }
        if (rc != 0)
            continue;
        if ((type_tag & MNEME_OVERFLOW_FLAG) != 0) {
            txn->reuse_pages = saved_reuse_pages;
            txn->reuse_count = saved_reuse_count;
            txn->reuse_cap   = saved_reuse_cap;
            return -1;
        }

        retire_val_t rv = {0};
        if (retire_val_unpack(val, val_len, &rv) < 0) {
            txn->reuse_pages = saved_reuse_pages;
            txn->reuse_count = saved_reuse_count;
            txn->reuse_cap   = saved_reuse_cap;
            return -1;
        }

        int write    = 0;
        int consumed = 0;
        for (int j = 0; j < rv.count; j++) {
            uint32_t pgno = rv.pgnos[j];
            int reused    = u32_contains_sorted(txn->consumed_reuse_pages, txn->consumed_reuse_count, pgno);
            if (reused) {
                consumed++;
                continue;
            }
            rv.pgnos[write++] = pgno;
        }

        if (consumed == 0) {
            retire_val_free(&rv);
            continue;
        }

        uint32_t new_root;
        if (write == 0) {
            new_root = mneme_btree_del(txn, txn->meta.retired_root_pgno, key, sizeof(key));
            if (new_root == 0) {
                retire_val_free(&rv);
                txn->reuse_pages = saved_reuse_pages;
                txn->reuse_count = saved_reuse_count;
                txn->reuse_cap   = saved_reuse_cap;
                return -1;
            }
        } else {
            uint8_t *new_val;
            uint32_t new_val_len;
            if (retire_val_pack(rv.pgnos, write, &new_val, &new_val_len) < 0) {
                retire_val_free(&rv);
                txn->reuse_pages = saved_reuse_pages;
                txn->reuse_count = saved_reuse_count;
                txn->reuse_cap   = saved_reuse_cap;
                return -1;
            }
            new_root = mneme_btree_put(txn, txn->meta.retired_root_pgno, key, sizeof(key), new_val, new_val_len,
                                       MNEME_TYPE_BYTES, 0);
            free(new_val);
            if (new_root == 0) {
                retire_val_free(&rv);
                txn->reuse_pages = saved_reuse_pages;
                txn->reuse_count = saved_reuse_count;
                txn->reuse_cap   = saved_reuse_cap;
                return -1;
            }
        }
        txn->meta.retired_root_pgno = new_root;
        consumed_total += (uint64_t)consumed;
        retire_val_free(&rv);
    }

    txn->reuse_pages = saved_reuse_pages;
    txn->reuse_count = saved_reuse_count;
    txn->reuse_cap   = saved_reuse_cap;

    if (txn->meta.retired_page_count >= consumed_total)
        txn->meta.retired_page_count -= consumed_total;
    else
        txn->meta.retired_page_count = 0;
    return 0;
}

int mneme_retire_persist(mneme_txn_t *txn) {
    if (!txn || txn->pending_retire_count == 0)
        return 0;

    qsort(txn->pending_retire, (size_t)txn->pending_retire_count, sizeof(uint32_t), cmp_u32);
    int uniq = 0;
    for (int i = 0; i < txn->pending_retire_count; i++) {
        if (uniq == 0 || txn->pending_retire[i] != txn->pending_retire[uniq - 1])
            txn->pending_retire[uniq++] = txn->pending_retire[i];
    }
    txn->pending_retire_count = uniq;

    uint32_t root = retire_root_ensure(txn);
    if (root == 0)
        return -1;

    /* The btree_put below COWs retire-tree pages, which appends the
       displaced pages to pending_retire mid-loop. Advance by the actual
       chunk size (not the max batch size) so those late entries land in
       a subsequent chunk instead of being skipped, and count only what
       was actually packed into a batch. The loop converges because puts
       on already-dirty retire-tree pages free no further pages. */
    uint32_t chunk_idx  = 0;
    int persisted_count = 0;
    for (int i = 0; i < txn->pending_retire_count;) {
        int chunk_count = txn->pending_retire_count - i;
        if (chunk_count > MNEME_RETIRE_BATCH_PGNOS)
            chunk_count = MNEME_RETIRE_BATCH_PGNOS;
        uint8_t key[12];
        mneme_retire_key_pack(key, txn->commit_txn_id, chunk_idx++);
        uint8_t *val;
        uint32_t val_len;
        if (retire_val_pack(txn->pending_retire + i, chunk_count, &val, &val_len) < 0)
            return -1;

        uint32_t *saved_reuse_pages = txn->reuse_pages;
        int saved_reuse_count       = txn->reuse_count;
        int saved_reuse_cap         = txn->reuse_cap;
        txn->reuse_pages            = NULL;
        txn->reuse_count            = 0;
        txn->reuse_cap              = 0;
        uint32_t new_root           = mneme_btree_put(txn, root, key, sizeof(key), val, val_len, MNEME_TYPE_BYTES, 0);
        txn->reuse_pages            = saved_reuse_pages;
        txn->reuse_count            = saved_reuse_count;
        txn->reuse_cap              = saved_reuse_cap;

        free(val);
        if (new_root == 0)
            return -1;
        root = new_root;
        i += chunk_count;
        persisted_count += chunk_count;
    }
    txn->meta.retired_root_pgno = root;
    txn->meta.retired_page_count += (uint64_t)persisted_count;
    return 0;
}