// SchFS B-tree torture harness.
//
// The B-tree is the one component of SchFS where a subtle bug is unfixable
// later: every other record type is built on it, so a corruption here is a
// corruption in everything. Testing it inside the kernel means rebuilding an
// image and driving a shell for each attempt. Testing it here means a fast
// loop against a reference model, so this runs first and the kernel gets the
// code only once it passes.
//
// What it checks:
//   1. Against a reference model — a plain sorted array — after every batch of
//      random inserts, replaces and deletes: same keys, same values, in order.
//   2. That the tree is structurally sound (magic, self-pointer, checksum,
//      level, key ordering) after every batch.
//   3. Crash safety: writes made after a commit are thrown away and the tree
//      is re-opened from the last committed root. Copy-on-write means it must
//      come back exactly as it was at that commit, never a mixture.
//   4. That a deliberately corrupted block is *detected* rather than silently
//      returning bad data.
//
// Build and run:  make schfs-test

#define SCHFS_HOSTED 1

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

#include "../kernel/fs/schfs/btree.c"

/* ------------------------------------------------------------------ */
/* Memory-backed block device with commit/rollback                     */
/* ------------------------------------------------------------------ */

#define NBLOCKS 8192

struct memdev {
    unsigned char blk[NBLOCKS][SCHFS_BLKSIZE];
    unsigned char committed[NBLOCKS][SCHFS_BLKSIZE];
    int      allocated[NBLOCKS];     // live in the working tree
    int      committed_alloc[NBLOCKS];
    uint64_t next;
    long     writes;
    int      fail_after;             // simulate a crash: -1 = never
};

static int
md_read(void *ctx, uint64_t b, void *buf)
{
    struct memdev *d = ctx;
    if (b == 0 || b >= NBLOCKS) return -1;
    memcpy(buf, d->blk[b], SCHFS_BLKSIZE);
    return 0;
}

static int
md_write(void *ctx, uint64_t b, const void *buf)
{
    struct memdev *d = ctx;
    if (b == 0 || b >= NBLOCKS) return -1;
    if (d->fail_after >= 0 && d->writes >= d->fail_after)
        return -1;                   // the "crash"
    memcpy(d->blk[b], buf, SCHFS_BLKSIZE);
    d->writes++;
    return 0;
}

static uint64_t
md_alloc(void *ctx)
{
    struct memdev *d = ctx;
    for (uint64_t i = d->next; i < NBLOCKS; i++) {
        // Never hand back a block the committed tree still references --
        // that is the core copy-on-write safety rule.
        if (!d->allocated[i] && !d->committed_alloc[i]) {
            d->allocated[i] = 1;
            d->next = i + 1;
            return i;
        }
    }
    for (uint64_t i = SCHFS_UBER_START + 1; i < d->next; i++) {
        if (!d->allocated[i] && !d->committed_alloc[i]) {
            d->allocated[i] = 1;
            return i;
        }
    }
    return 0;                        // out of space
}

static void
md_free(void *ctx, uint64_t b)
{
    struct memdev *d = ctx;
    if (b > 0 && b < NBLOCKS)
        d->allocated[b] = 0;         // reusable only once committed
}

// Publish everything written so far and open a new transaction group.
//
// Advancing the tree's generation is not bookkeeping — it is the rule the tree
// uses to decide whether a node may be rewritten in place. Without it every
// block would look like part of the open transaction and the tree would happily
// overwrite the very blocks this snapshot is supposed to preserve, so the
// rollback test below is what keeps that rule honest.
static void
md_commit(struct memdev *d, struct schfs_btree *t)
{
    memcpy(d->committed, d->blk, sizeof(d->blk));
    memcpy(d->committed_alloc, d->allocated, sizeof(d->allocated));
    t->generation++;
}

static void
md_rollback(struct memdev *d)
{
    memcpy(d->blk, d->committed, sizeof(d->blk));
    memcpy(d->allocated, d->committed_alloc, sizeof(d->allocated));
}

/* ------------------------------------------------------------------ */
/* Reference model                                                     */
/* ------------------------------------------------------------------ */

#define MAXREF 4000

struct ref_ent {
    struct schfs_key key;
    unsigned char    val[64];
    unsigned         len;
    int              live;
};

static struct ref_ent ref[MAXREF];
static int            nref;

// Deliberately does NOT call schfs_key_cmp. The reference model is the oracle,
// so it must not share the code under test: when the model used the real
// comparator, a deliberate bug in schfs_key_cmp broke the tree and the oracle
// in exactly the same way and every check still passed. Written out longhand
// so a bug in the real comparator surfaces as a disagreement.
static int
ref_key_cmp(const struct schfs_key *a, const struct schfs_key *b)
{
    if (a->objectid != b->objectid) return a->objectid < b->objectid ? -1 : 1;
    if (a->type     != b->type)     return a->type     < b->type     ? -1 : 1;
    if (a->offset   != b->offset)   return a->offset   < b->offset   ? -1 : 1;
    return 0;
}

static int
ref_find(const struct schfs_key *k)
{
    for (int i = 0; i < nref; i++)
        if (ref[i].live && ref_key_cmp(&ref[i].key, k) == 0)
            return i;
    return -1;
}

static int
ref_cmp(const void *a, const void *b)
{
    return ref_key_cmp(&((const struct ref_ent *)a)->key,
                       &((const struct ref_ent *)b)->key);
}

/* Compare a full tree walk against the reference model. */
struct cmp_ctx {
    struct ref_ent *sorted;
    int             n;
    int             i;
    int             bad;
};

static int
cmp_one(const struct schfs_key *k, const void *v, unsigned len, void *arg)
{
    struct cmp_ctx *c = arg;
    if (c->i >= c->n) {
        fprintf(stderr, "  tree has extra key (%llu,%u,%llu)\n",
                (unsigned long long)k->objectid, k->type,
                (unsigned long long)k->offset);
        c->bad = 1;
        return 1;
    }
    struct ref_ent *e = &c->sorted[c->i++];
    if (ref_key_cmp(k, &e->key) != 0) {
        fprintf(stderr, "  key mismatch at %d: tree (%llu,%u,%llu) "
                "vs ref (%llu,%u,%llu)\n", c->i - 1,
                (unsigned long long)k->objectid, k->type,
                (unsigned long long)k->offset,
                (unsigned long long)e->key.objectid, e->key.type,
                (unsigned long long)e->key.offset);
        c->bad = 1;
        return 1;
    }
    if (len != e->len || memcmp(v, e->val, len) != 0) {
        fprintf(stderr, "  value mismatch for key (%llu,%u,%llu)\n",
                (unsigned long long)k->objectid, k->type,
                (unsigned long long)k->offset);
        c->bad = 1;
        return 1;
    }
    return 0;
}

static int
verify(struct schfs_btree *t, const char *what)
{
    static struct ref_ent sorted[MAXREF];
    int n = 0;
    for (int i = 0; i < nref; i++)
        if (ref[i].live)
            sorted[n++] = ref[i];
    qsort(sorted, n, sizeof(sorted[0]), ref_cmp);

    if (schfs_btree_check(t) != 0) {
        fprintf(stderr, "FAIL [%s]: structural check failed\n", what);
        return 1;
    }

    struct cmp_ctx c = { sorted, n, 0, 0 };
    schfs_btree_iterate(t, cmp_one, &c);
    if (c.bad) {
        fprintf(stderr, "FAIL [%s]: contents differ from reference\n", what);
        return 1;
    }
    if (c.i != n) {
        fprintf(stderr, "FAIL [%s]: tree has %d records, reference has %d\n",
                what, c.i, n);
        return 1;
    }
    return 0;
}

/* ------------------------------------------------------------------ */

static uint32_t rng_state = 12345;
static uint32_t
rnd(void)
{
    rng_state ^= rng_state << 13;
    rng_state ^= rng_state >> 17;
    rng_state ^= rng_state << 5;
    return rng_state;
}

static struct memdev dev_storage;

// The three scratch blocks every mutation borrows. Statics, not locals, because
// this is exactly the shape the kernel has to use: a block is 4096 bytes and so
// is the whole kernel stack.
static char CTX_A[SCHFS_BLKSIZE];
static char CTX_B[SCHFS_BLKSIZE];
static char CTX_C[SCHFS_BLKSIZE];
static struct schfs_ctx CTX = { CTX_A, CTX_B, CTX_C };

static void
mkdev(struct schfs_bdev *bd, struct memdev *d)
{
    memset(d, 0, sizeof(*d));
    memset(bd, 0, sizeof(*bd));      /* optional hooks must read as absent */
    d->next = SCHFS_UBER_START + 1;
    d->fail_after = -1;
    bd->ctx    = d;
    bd->bread  = md_read;
    bd->bwrite = md_write;
    bd->balloc = md_alloc;
    bd->bfree  = md_free;
    bd->balloc_run = 0;              /* the B-tree never asks for a run */
}

int
main(void)
{
    struct schfs_bdev  bd;
    struct schfs_btree t;
    int failures = 0;

    printf("SchFS B-tree torture test\n");
    printf("  block size %d, max inline value %d, fanout %d/%d (kptr/item)\n\n",
           SCHFS_BLKSIZE, SCHFS_MAX_INLINE, SCHFS_MAX_KPTR, SCHFS_MAX_ITEMS);

    /* ---- 1. random insert / replace / delete against the model ---- */
    mkdev(&bd, &dev_storage);
    memset(&t, 0, sizeof(t));
    t.generation = 1;                    /* the first transaction group */
    if (schfs_btree_create(&t, &bd, &CTX) != 0) {
        fprintf(stderr, "FAIL: create\n");
        return 1;
    }
    nref = 0;

    // Committing part way through matters: between commits the tree rewrites
    // its own nodes in place, and only after one does the copy-on-write path
    // run at all. A test that never commits exercises half the code.
    const int ROUNDS = 60, OPS = 60, COMMIT_EVERY = 20;
    for (int round = 0; round < ROUNDS; round++) {
        if (round && round % COMMIT_EVERY == 0)
            md_commit(&dev_storage, &t);
        for (int op = 0; op < OPS; op++) {
            struct schfs_key k;
            memset(&k, 0, sizeof(k));
            k.objectid = rnd() % 200;
            k.type     = (uint8_t)(1 + rnd() % 3);
            k.offset   = rnd() % 50;

            int action = rnd() % 10;
            if (action < 7) {                    /* insert or replace */
                unsigned char val[64];
                unsigned len = 1 + rnd() % 48;
                for (unsigned i = 0; i < len; i++)
                    val[i] = (unsigned char)(rnd() & 0xff);

                int r = schfs_btree_insert(&t, &CTX, &k, val, len);
                if (r == -3)
                    continue;                    /* out of blocks: fine */
                if (r != 0) {
                    fprintf(stderr, "FAIL: insert returned %d\n", r);
                    return 1;
                }
                int idx = ref_find(&k);
                if (idx < 0) {
                    if (nref >= MAXREF) continue;
                    idx = nref++;
                    ref[idx].key  = k;
                    ref[idx].live = 1;
                }
                memcpy(ref[idx].val, val, len);
                ref[idx].len = len;
            } else {                             /* delete */
                int r = schfs_btree_delete(&t, &CTX, &k);
                int idx = ref_find(&k);
                if (idx >= 0) {
                    if (r != 0) {
                        fprintf(stderr, "FAIL: delete of present key -> %d\n", r);
                        return 1;
                    }
                    ref[idx].live = 0;
                } else if (r != -1) {
                    fprintf(stderr, "FAIL: delete of absent key -> %d "
                            "(expected -1)\n", r);
                    return 1;
                }
            }
        }
        if (verify(&t, "random ops")) {
            failures++;
            break;
        }
    }
    if (!failures)
        printf("  [ok] %d rounds x %d random ops match the reference model\n",
               ROUNDS, OPS);

    /* ---- 2. lookup agrees with the model for every live key ---- */
    {
        int checked = 0, bad = 0;
        for (int i = 0; i < nref; i++) {
            if (!ref[i].live) continue;
            unsigned char buf[64];
            unsigned len = sizeof(buf);
            if (schfs_btree_lookup(&t, &CTX, &ref[i].key, buf, &len) != 0) {
                bad++;
                continue;
            }
            if (len != ref[i].len || memcmp(buf, ref[i].val, len) != 0)
                bad++;
            checked++;
        }
        if (bad) {
            fprintf(stderr, "FAIL: %d of %d lookups wrong\n", bad, checked + bad);
            failures++;
        } else {
            printf("  [ok] %d lookups all return the right value\n", checked);
        }
    }

    /* ---- 3. the cursor: forward, backward, range and predecessor ----
     *
     * The cursor is what the kernel actually reads through, because it holds
     * one caller-supplied block buffer instead of putting 4 KB on a 4 KB
     * kernel stack per level. It re-reads a parent when climbing, so it has
     * genuinely different failure modes from the recursive walker and is
     * tested against the same independent model rather than against iterate.
     */
    {
        static struct ref_ent sorted[MAXREF];
        int n = 0;
        for (int i = 0; i < nref; i++)
            if (ref[i].live)
                sorted[n++] = ref[i];
        qsort(sorted, n, sizeof(sorted[0]), ref_cmp);

        static char curblk[SCHFS_BLKSIZE];
        struct schfs_cursor c;
        schfs_cursor_init(&c, &t, curblk);

        /* (a) forward from the start */
        int bad = 0, seen = 0;
        if (schfs_cursor_seek(&c, NULL) != 0 && n > 0)
            bad = 1;
        while (!bad && seen < n) {
            const struct schfs_key *k = schfs_cursor_key(&c);
            unsigned vlen;
            const void *v = schfs_cursor_val(&c, &vlen);
            if (!k || !v) { bad = 1; break; }
            if (ref_key_cmp(k, &sorted[seen].key) != 0 ||
                vlen != sorted[seen].len ||
                memcmp(v, sorted[seen].val, vlen) != 0) {
                fprintf(stderr, "  cursor diverges at record %d\n", seen);
                bad = 1;
                break;
            }
            seen++;
            int r = schfs_cursor_next(&c);
            if (r == -1) break;
            if (r != 0) { bad = 1; break; }
        }
        if (bad || seen != n) {
            fprintf(stderr, "FAIL: forward cursor walked %d of %d records\n",
                    seen, n);
            failures++;
        } else {
            printf("  [ok] cursor walks all %d records forward in key order\n", n);
        }

        /* (b) backward from the end. Getting to the end uses lookup_le with a
         *     key past every real one, which also exercises the "greater than
         *     everything" branch. */
        struct schfs_key maxk;
        memset(&maxk, 0, sizeof(maxk));
        maxk.objectid = ~0ULL; maxk.type = 255; maxk.offset = ~0ULL;

        bad = 0; seen = 0;
        if (n > 0 && schfs_btree_lookup_le(&c, &maxk) != 0)
            bad = 1;
        while (!bad && seen < n) {
            const struct schfs_key *k = schfs_cursor_key(&c);
            if (!k) { bad = 1; break; }
            if (ref_key_cmp(k, &sorted[n - 1 - seen].key) != 0) {
                fprintf(stderr, "  reverse cursor diverges at record %d\n", seen);
                bad = 1;
                break;
            }
            seen++;
            int r = schfs_cursor_prev(&c);
            if (r == -1) break;
            if (r != 0) { bad = 1; break; }
        }
        if (bad || seen != n) {
            fprintf(stderr, "FAIL: reverse cursor walked %d of %d records\n",
                    seen, n);
            failures++;
        } else {
            printf("  [ok] cursor walks all %d records backward\n", n);
        }

        /* (c) range scans from arbitrary start keys must yield exactly the
         *     model's suffix. This is the operation readdir depends on. */
        bad = 0;
        for (int trial = 0; trial < 200 && !bad; trial++) {
            struct schfs_key start;
            memset(&start, 0, sizeof(start));
            start.objectid = rnd() % 220;
            start.type     = (uint8_t)(rnd() % 5);
            start.offset   = rnd() % 60;

            int expect = 0;
            while (expect < n && ref_key_cmp(&sorted[expect].key, &start) < 0)
                expect++;

            int r = schfs_cursor_seek(&c, &start);
            if (r == -2) { bad = 1; break; }
            if (expect == n) {
                if (r == 0 && schfs_cursor_key(&c) != NULL) bad = 1;
                continue;
            }
            if (r != 0) { bad = 1; break; }
            const struct schfs_key *k = schfs_cursor_key(&c);
            if (!k || ref_key_cmp(k, &sorted[expect].key) != 0) {
                fprintf(stderr, "  seek(%llu,%u,%llu) landed wrong\n",
                        (unsigned long long)start.objectid, start.type,
                        (unsigned long long)start.offset);
                bad = 1;
            }
        }
        if (bad) {
            fprintf(stderr, "FAIL: range seek disagrees with the model\n");
            failures++;
        } else {
            printf("  [ok] 200 range seeks land on the right record\n");
        }

        /* (d) lookup_le must find the greatest key <= the probe, which is how
         *     a file read finds the extent covering a byte offset. */
        bad = 0;
        for (int trial = 0; trial < 200 && !bad; trial++) {
            struct schfs_key probe;
            memset(&probe, 0, sizeof(probe));
            probe.objectid = rnd() % 220;
            probe.type     = (uint8_t)(rnd() % 5);
            probe.offset   = rnd() % 60;

            int expect = -1;
            for (int i = 0; i < n; i++)
                if (ref_key_cmp(&sorted[i].key, &probe) <= 0)
                    expect = i;
                else
                    break;

            int r = schfs_btree_lookup_le(&c, &probe);
            if (expect < 0) {
                if (r != -1) {
                    fprintf(stderr, "  lookup_le below all keys returned %d\n", r);
                    bad = 1;
                }
                continue;
            }
            if (r != 0) {
                fprintf(stderr, "  lookup_le returned %d, expected a hit\n", r);
                bad = 1;
                break;
            }
            const struct schfs_key *k = schfs_cursor_key(&c);
            if (!k || ref_key_cmp(k, &sorted[expect].key) != 0) {
                fprintf(stderr, "  lookup_le(%llu,%u,%llu) landed wrong\n",
                        (unsigned long long)probe.objectid, probe.type,
                        (unsigned long long)probe.offset);
                bad = 1;
            }
        }
        if (bad) {
            fprintf(stderr, "FAIL: lookup_le disagrees with the model\n");
            failures++;
        } else {
            printf("  [ok] 200 predecessor lookups match the model\n");
        }
    }

    /* ---- 4. crash safety ---- */
    {
        struct memdev *d = &dev_storage;
        md_commit(d, &t);
        uint64_t saved_root = t.root;
        uint32_t saved_crc  = t.root_crc;
        uint64_t saved_gen  = t.generation;

        /* Snapshot what the committed tree contains. */
        static struct ref_ent snap[MAXREF];
        int nsnap = 0;
        for (int i = 0; i < nref; i++)
            if (ref[i].live) snap[nsnap++] = ref[i];

        /* Now do a pile of work and then throw the writes away. */
        for (int op = 0; op < 500; op++) {
            struct schfs_key k;
            memset(&k, 0, sizeof(k));
            k.objectid = 1000 + rnd() % 500;
            k.type     = 1;
            k.offset   = rnd() % 100;
            unsigned char val[32];
            memset(val, (int)(rnd() & 0xff), sizeof(val));
            schfs_btree_insert(&t, &CTX, &k, val, sizeof(val));
        }
        md_rollback(d);
        t.root       = saved_root;
        t.root_crc   = saved_crc;
        t.generation = saved_gen;

        /* The committed tree must be intact and unchanged. */
        nref = 0;
        for (int i = 0; i < nsnap; i++) ref[nref++] = snap[i];

        if (verify(&t, "after simulated crash")) {
            fprintf(stderr, "FAIL: tree did not survive rollback\n");
            failures++;
        } else {
            printf("  [ok] tree intact after discarding 500 uncommitted ops\n");
        }
    }

    /* ---- 5. corruption is detected, not silently returned ---- */
    {
        struct memdev *d = &dev_storage;
        unsigned char save[SCHFS_BLKSIZE];
        uint64_t victim = t.root;
        memcpy(save, d->blk[victim], SCHFS_BLKSIZE);

        d->blk[victim][sizeof(struct schfs_node_hdr) + 4] ^= 0xff;
        int r = schfs_btree_check(&t);
        if (r != -2) {
            fprintf(stderr, "FAIL: flipped a bit in the root and the tree "
                    "still verified (got %d, wanted -2)\n", r);
            failures++;
        } else {
            printf("  [ok] bit flip in a node is caught by the checksum\n");
        }
        memcpy(d->blk[victim], save, SCHFS_BLKSIZE);

        /* A block written to the wrong address still checksums correctly --
           only the self-pointer catches it. */
        uint64_t elsewhere = md_alloc(d);
        if (elsewhere) {
            memcpy(d->blk[elsewhere], d->blk[victim], SCHFS_BLKSIZE);
            char tmp[SCHFS_BLKSIZE];
            if (node_read(&bd, elsewhere, 0, tmp) != -2) {
                fprintf(stderr, "FAIL: misdirected write not detected\n");
                failures++;
            } else {
                printf("  [ok] misdirected write caught by the self-pointer\n");
            }
        }
    }

    printf("\n%s\n", failures ? "FAILED" : "all checks passed");
    return failures ? 1 : 0;
}