// Incremental MST updates must produce byte-identical trees to a full
// rebuild. buildMst is the oracle: any divergence in root CID means the
// incremental path corrupted the tree.

import { buildMst, mstAdd, mstDelete, mstUpdate } from '@pdsjs/core/mst';
import { getKeyDepth } from '@pdsjs/core/repo';
import { describe, expect, test } from 'vitest';

/** Deterministic pseudo-random source so failures reproduce. */
function rng(/** @type {number} */ seed) {
  let s = seed >>> 0;
  return () => {
    s = (s * 1664525 + 1013904223) >>> 0;
    return s / 0x100000000;
  };
}

function createStore() {
  /** @type {Map<string, Uint8Array>} */
  const blocks = new Map();
  return {
    blocks,
    ctx: {
      getBlock: async (/** @type {string} */ cid) => blocks.get(cid) ?? null,
      putBlock: async (
        /** @type {string} */ cid,
        /** @type {Uint8Array} */ data,
      ) => {
        blocks.set(cid, data);
      },
    },
  };
}

/** Root CID a full rebuild would produce for this key/value set. */
async function rebuildRoot(
  /** @type {Map<string, string>} */ map,
  /** @type {any} */ store,
) {
  const records = [...map.entries()]
    .map(([key, cid]) => ({ key, cid }))
    .sort((a, b) => (a.key < b.key ? -1 : a.key > b.key ? 1 : 0));
  return buildMst(records, store.ctx.putBlock);
}

const VALUE = 'bafyreib2rxk3rh6kzwq6wxjxvhoiwcsuhkzqvhzcwqz6zvxfhbfgz4pgvy';
const VALUE2 = 'bafyreiczsscbxbwyxfgnb5mhqrpxbzn4bkbmxvhkfkzqbcqvzfkwqvbfha';

/** Find a key whose MST depth is exactly `depth`. */
async function keyWithDepth(
  /** @type {number} */ depth,
  prefix = 'app.bsky.feed.post/',
) {
  for (let i = 0; i < 500000; i++) {
    const key = `${prefix}k${i.toString(36)}`;
    if ((await getKeyDepth(key)) === depth) return key;
  }
  throw new Error(`no key found at depth ${depth}`);
}

describe('incremental MST vs full rebuild', () => {
  test('sequential adds match rebuild at every step', async () => {
    const store = createStore();
    /** @type {Map<string, string>} */
    const map = new Map();
    let root = null;

    for (let i = 0; i < 150; i++) {
      const key = `app.bsky.feed.post/3k${i.toString(36).padStart(6, 'a')}`;
      root = await mstAdd(root, key, VALUE, store.ctx);
      map.set(key, VALUE);
      expect(root, `after adding ${key}`).toBe(await rebuildRoot(map, store));
    }
  });

  test('interleaved adds and deletes match rebuild at every step', async () => {
    const store = createStore();
    /** @type {Map<string, string>} */
    const map = new Map();
    let root = null;
    const rand = rng(20260730);

    for (let step = 0; step < 400; step++) {
      const keys = [...map.keys()];
      const shouldDelete = keys.length > 0 && rand() < 0.35;

      if (shouldDelete) {
        const key = keys[Math.floor(rand() * keys.length)];
        root = await mstDelete(root, key, store.ctx);
        map.delete(key);
      } else {
        const key = `app.bsky.feed.post/3k${Math.floor(rand() * 1e9).toString(36)}`;
        if (map.has(key)) continue;
        root = await mstAdd(root, key, VALUE, store.ctx);
        map.set(key, VALUE);
      }

      expect(root, `step ${step} (size ${map.size})`).toBe(
        await rebuildRoot(map, store),
      );
    }
  });

  test('adding a deeper key grows the tree correctly', async () => {
    const store = createStore();
    /** @type {Map<string, string>} */
    const map = new Map();
    let root = null;

    // Shallow keys first, then one much deeper, forcing new layers above
    for (let i = 0; i < 12; i++) {
      const key = await keyWithDepth(0, `app.bsky.feed.post/s${i}`);
      root = await mstAdd(root, key, VALUE, store.ctx);
      map.set(key, VALUE);
    }
    expect(root).toBe(await rebuildRoot(map, store));

    const deep = await keyWithDepth(3);
    expect(await getKeyDepth(deep)).toBe(3);
    root = await mstAdd(root, deep, VALUE, store.ctx);
    map.set(deep, VALUE);
    expect(root, 'after deep insert').toBe(await rebuildRoot(map, store));

    // ...and removing it shrinks the tree back
    root = await mstDelete(root, deep, store.ctx);
    map.delete(deep);
    expect(root, 'after deep delete').toBe(await rebuildRoot(map, store));
  });

  test('update replaces a value without changing structure', async () => {
    const store = createStore();
    /** @type {Map<string, string>} */
    const map = new Map();
    let root = null;

    for (let i = 0; i < 40; i++) {
      const key = `app.bsky.feed.post/3k${i.toString(36).padStart(6, 'a')}`;
      root = await mstAdd(root, key, VALUE, store.ctx);
      map.set(key, VALUE);
    }

    const target = [...map.keys()][17];
    root = await mstUpdate(root, target, VALUE2, store.ctx);
    map.set(target, VALUE2);
    expect(root).toBe(await rebuildRoot(map, store));
  });

  test('adding a key the tree already holds replaces it', async () => {
    const store = createStore();
    /** @type {Map<string, string>} */
    const map = new Map();
    let root = null;

    for (let i = 0; i < 40; i++) {
      const key = `app.bsky.feed.post/3k${i.toString(36).padStart(6, 'a')}`;
      root = await mstAdd(root, key, VALUE, store.ctx);
      map.set(key, VALUE);
    }

    // A caller that says "create" for a key already present must not get a
    // second entry beside the first. Two entries sharing a key break the
    // strict ordering a verifier checks, and the repo stops being provable.
    const target = [...map.keys()][17];
    root = await mstAdd(root, target, VALUE2, store.ctx);
    map.set(target, VALUE2);
    expect(root).toBe(await rebuildRoot(map, store));
  });

  test('re-adding a key at every depth keeps the tree rebuildable', async () => {
    for (const depth of [0, 1, 2]) {
      const store = createStore();
      /** @type {Map<string, string>} */
      const map = new Map();
      let root = null;

      for (let i = 0; i < 25; i++) {
        const key = `app.bsky.feed.post/d${depth}k${i.toString(36)}`;
        root = await mstAdd(root, key, VALUE, store.ctx);
        map.set(key, VALUE);
      }

      const target = await keyWithDepth(depth);
      root = await mstAdd(root, target, VALUE, store.ctx);
      map.set(target, VALUE);
      root = await mstAdd(root, target, VALUE2, store.ctx);
      map.set(target, VALUE2);

      expect(root).toBe(await rebuildRoot(map, store));
    }
  });

  test('deleting every key returns to an empty tree', async () => {
    const store = createStore();
    /** @type {Map<string, string>} */
    const map = new Map();
    let root = null;

    for (let i = 0; i < 30; i++) {
      const key = `app.bsky.feed.post/3k${i.toString(36).padStart(6, 'a')}`;
      root = await mstAdd(root, key, VALUE, store.ctx);
      map.set(key, VALUE);
    }
    for (const key of [...map.keys()]) {
      root = await mstDelete(root, key, store.ctx);
      map.delete(key);
      expect(root).toBe(await rebuildRoot(map, store));
    }
    expect(root).toBe(null);
  });

  test('touches only a handful of nodes per write', async () => {
    const store = createStore();
    let root = null;
    for (let i = 0; i < 400; i++) {
      const key = `app.bsky.feed.post/3k${i.toString(36).padStart(6, 'a')}`;
      root = await mstAdd(root, key, VALUE, store.ctx);
    }

    let written = 0;
    const counting = {
      getBlock: store.ctx.getBlock,
      putBlock: async (
        /** @type {string} */ cid,
        /** @type {Uint8Array} */ data,
      ) => {
        written++;
        await store.ctx.putBlock(cid, data);
      },
    };
    await mstAdd(root, 'app.bsky.feed.post/3kzzzzzzzz', VALUE, counting);

    // One root-to-leaf path, not the whole tree (a rebuild writes 100+)
    expect(written).toBeLessThan(12);
  });
});