import { describe, it, expect, beforeEach } from 'vitest';
import { BPlusTree } from '../../src/utils/btree.js';

describe('BPlusTree', () => {
  let tree: BPlusTree<number, string>;

  beforeEach(() => {
    tree = new BPlusTree<number, string>();
  });

  // -------------------------------------------------------------------------
  // Construction
  // -------------------------------------------------------------------------

  describe('constructor', () => {
    it('should create an empty tree', () => {
      expect(tree.size).toBe(0);
    });

    it('should reject order < 3', () => {
      expect(() => new BPlusTree(2)).toThrow('order must be at least 3');
    });
  });

  // -------------------------------------------------------------------------
  // Insert & Get
  // -------------------------------------------------------------------------

  describe('insert and get', () => {
    it('should insert and retrieve a single entry', () => {
      tree.insert(10, 'a');
      expect(tree.get(10)).toEqual(new Set(['a']));
      expect(tree.size).toBe(1);
    });

    it('should handle multiple distinct keys', () => {
      tree.insert(10, 'a');
      tree.insert(20, 'b');
      tree.insert(5, 'c');
      expect(tree.get(10)).toEqual(new Set(['a']));
      expect(tree.get(20)).toEqual(new Set(['b']));
      expect(tree.get(5)).toEqual(new Set(['c']));
      expect(tree.size).toBe(3);
    });

    it('should return undefined for missing keys', () => {
      tree.insert(10, 'a');
      expect(tree.get(99)).toBeUndefined();
    });

    it('should accumulate duplicate keys into a Set', () => {
      tree.insert(10, 'a');
      tree.insert(10, 'b');
      tree.insert(10, 'c');
      expect(tree.get(10)).toEqual(new Set(['a', 'b', 'c']));
      expect(tree.size).toBe(3);
    });

    it('should not double-count duplicate values for the same key', () => {
      tree.insert(10, 'a');
      tree.insert(10, 'a');
      expect(tree.get(10)).toEqual(new Set(['a']));
      expect(tree.size).toBe(1);
    });
  });

  // -------------------------------------------------------------------------
  // Delete
  // -------------------------------------------------------------------------

  describe('delete', () => {
    it('should delete a specific value from a key', () => {
      tree.insert(10, 'a');
      tree.insert(10, 'b');
      expect(tree.delete(10, 'a')).toBe(true);
      expect(tree.get(10)).toEqual(new Set(['b']));
      expect(tree.size).toBe(1);
    });

    it('should remove the key entry when its last value is deleted', () => {
      tree.insert(10, 'a');
      expect(tree.delete(10, 'a')).toBe(true);
      expect(tree.get(10)).toBeUndefined();
      expect(tree.size).toBe(0);
    });

    it('should delete all values for a key when value is omitted', () => {
      tree.insert(10, 'a');
      tree.insert(10, 'b');
      expect(tree.delete(10)).toBe(true);
      expect(tree.get(10)).toBeUndefined();
      expect(tree.size).toBe(0);
    });

    it('should return false for non-existent key', () => {
      expect(tree.delete(99)).toBe(false);
    });

    it('should return false for non-existent value', () => {
      tree.insert(10, 'a');
      expect(tree.delete(10, 'z')).toBe(false);
      expect(tree.size).toBe(1);
    });
  });

  // -------------------------------------------------------------------------
  // Range queries
  // -------------------------------------------------------------------------

  describe('range', () => {
    beforeEach(() => {
      for (let i = 0; i < 100; i++) {
        tree.insert(i, `v${i}`);
      }
    });

    it('should return all entries when no bounds given', () => {
      const result = tree.range();
      expect(result.length).toBe(100);
      expect(result[0].key).toBe(0);
      expect(result[99].key).toBe(99);
    });

    it('should return entries in key order', () => {
      const keys = tree.range().map((e) => e.key);
      for (let i = 1; i < keys.length; i++) {
        expect(keys[i]).toBeGreaterThan(keys[i - 1]);
      }
    });

    it('should respect lower bound (inclusive by default)', () => {
      const result = tree.range(50);
      expect(result.length).toBe(50);
      expect(result[0].key).toBe(50);
    });

    it('should respect upper bound (exclusive by default)', () => {
      const result = tree.range(undefined, 10);
      expect(result.length).toBe(10);
      expect(result[result.length - 1].key).toBe(9);
    });

    it('should support inclusive upper bound', () => {
      const result = tree.range(undefined, 10, { upperInclusive: true });
      expect(result.length).toBe(11);
      expect(result[result.length - 1].key).toBe(10);
    });

    it('should support exclusive lower bound', () => {
      const result = tree.range(50, undefined, { lowerInclusive: false });
      expect(result.length).toBe(49);
      expect(result[0].key).toBe(51);
    });

    it('should handle combined bounds', () => {
      const result = tree.range(20, 30);
      expect(result.length).toBe(10);
      expect(result[0].key).toBe(20);
      expect(result[result.length - 1].key).toBe(29);
    });

    it('should return empty array for no-result range', () => {
      const result = tree.range(200, 300);
      expect(result).toEqual([]);
    });

    it('should return empty for inverted bounds', () => {
      const result = tree.range(50, 10);
      expect(result).toEqual([]);
    });
  });

  // -------------------------------------------------------------------------
  // Edge cases
  // -------------------------------------------------------------------------

  describe('edge cases', () => {
    it('should handle get on empty tree', () => {
      expect(tree.get(1)).toBeUndefined();
    });

    it('should handle range on empty tree', () => {
      expect(tree.range()).toEqual([]);
    });

    it('should handle delete on empty tree', () => {
      expect(tree.delete(1)).toBe(false);
    });

    it('should handle insert-then-delete-all back to empty', () => {
      for (let i = 0; i < 50; i++) {
        tree.insert(i, `v${i}`);
      }
      for (let i = 0; i < 50; i++) {
        expect(tree.delete(i, `v${i}`)).toBe(true);
      }
      expect(tree.size).toBe(0);
      expect(tree.range()).toEqual([]);
      // Verify we can still insert after emptying.
      tree.insert(1, 'new');
      expect(tree.get(1)).toEqual(new Set(['new']));
    });
  });

  // -------------------------------------------------------------------------
  // Clear
  // -------------------------------------------------------------------------

  describe('clear', () => {
    it('should reset the tree to empty', () => {
      for (let i = 0; i < 100; i++) tree.insert(i, `v${i}`);
      expect(tree.size).toBe(100);
      tree.clear();
      expect(tree.size).toBe(0);
      expect(tree.get(0)).toBeUndefined();
      expect(tree.range()).toEqual([]);
    });
  });

  // -------------------------------------------------------------------------
  // Entries iterator
  // -------------------------------------------------------------------------

  describe('entries', () => {
    it('should yield all entries in key order', () => {
      tree.insert(30, 'c');
      tree.insert(10, 'a');
      tree.insert(20, 'b');

      const result = [...tree.entries()];
      expect(result.map((e) => e.key)).toEqual([10, 20, 30]);
    });

    it('should yield nothing for empty tree', () => {
      expect([...tree.entries()]).toEqual([]);
    });
  });

  // -------------------------------------------------------------------------
  // Large dataset
  // -------------------------------------------------------------------------

  describe('large dataset', () => {
    const N = 10_000;

    it('should correctly store and retrieve N items', () => {
      for (let i = 0; i < N; i++) {
        tree.insert(i, `v${i}`);
      }
      expect(tree.size).toBe(N);

      // Spot-check some values.
      expect(tree.get(0)).toEqual(new Set(['v0']));
      expect(tree.get(N - 1)).toEqual(new Set([`v${N - 1}`]));
      expect(tree.get(Math.floor(N / 2))).toEqual(new Set([`v${Math.floor(N / 2)}`]));
    });

    it('should produce correct range results on large dataset', () => {
      for (let i = 0; i < N; i++) {
        tree.insert(i, `v${i}`);
      }

      const result = tree.range(5000, 5010);
      expect(result.length).toBe(10);
      expect(result[0].key).toBe(5000);
      expect(result[9].key).toBe(5009);
    });

    it('should survive inserting and deleting many items', () => {
      for (let i = 0; i < N; i++) {
        tree.insert(i, `v${i}`);
      }

      // Delete the first half.
      for (let i = 0; i < N / 2; i++) {
        expect(tree.delete(i, `v${i}`)).toBe(true);
      }

      expect(tree.size).toBe(N / 2);
      expect(tree.get(0)).toBeUndefined();
      expect(tree.get(N / 2)).toEqual(new Set([`v${N / 2}`]));

      // Remaining range should start at N/2.
      const remaining = tree.range();
      expect(remaining.length).toBe(N / 2);
      expect(remaining[0].key).toBe(N / 2);
    });
  });

  // -------------------------------------------------------------------------
  // Custom comparator
  // -------------------------------------------------------------------------

  describe('custom comparator', () => {
    it('should support reverse ordering', () => {
      const reverseTree = new BPlusTree<number, string>(32, (a, b) => b - a);
      reverseTree.insert(1, 'a');
      reverseTree.insert(2, 'b');
      reverseTree.insert(3, 'c');

      const entries = [...reverseTree.entries()];
      expect(entries.map((e) => e.key)).toEqual([3, 2, 1]);
    });
  });

  // -------------------------------------------------------------------------
  // Node splitting (small order to force splits)
  // -------------------------------------------------------------------------

  describe('node splitting with small order', () => {
    let smallTree: BPlusTree<number, string>;

    beforeEach(() => {
      smallTree = new BPlusTree<number, string>(4);
    });

    it('should handle splits correctly', () => {
      // Order 4 means max 3 keys per node — splits after the 4th insert.
      for (let i = 0; i < 20; i++) {
        smallTree.insert(i, `v${i}`);
      }
      expect(smallTree.size).toBe(20);

      // All values should be retrievable.
      for (let i = 0; i < 20; i++) {
        expect(smallTree.get(i)).toEqual(new Set([`v${i}`]));
      }
    });

    it('should maintain sorted order after many splits', () => {
      // Insert in random order to stress split logic.
      const values = Array.from({ length: 50 }, (_, i) => i);
      for (let i = values.length - 1; i > 0; i--) {
        const j = Math.floor(Math.random() * (i + 1));
        [values[i], values[j]] = [values[j], values[i]];
      }

      for (const v of values) {
        smallTree.insert(v, `v${v}`);
      }

      const entries = [...smallTree.entries()];
      const keys = entries.map((e) => e.key);
      expect(keys).toEqual([...keys].sort((a, b) => a - b));
    });

    it('should handle delete with merging at small order', () => {
      for (let i = 0; i < 20; i++) {
        smallTree.insert(i, `v${i}`);
      }

      // Delete enough to trigger merges.
      for (let i = 0; i < 15; i++) {
        expect(smallTree.delete(i, `v${i}`)).toBe(true);
      }

      expect(smallTree.size).toBe(5);

      // Remaining keys should be intact.
      for (let i = 15; i < 20; i++) {
        expect(smallTree.get(i)).toEqual(new Set([`v${i}`]));
      }
    });
  });

  // -------------------------------------------------------------------------
  // String keys
  // -------------------------------------------------------------------------

  describe('string keys', () => {
    it('should work with string keys using default comparator', () => {
      const strTree = new BPlusTree<string, number>();
      strTree.insert('banana', 1);
      strTree.insert('apple', 2);
      strTree.insert('cherry', 3);

      const entries = [...strTree.entries()];
      expect(entries.map((e) => e.key)).toEqual(['apple', 'banana', 'cherry']);
      expect(strTree.get('banana')).toEqual(new Set([1]));
    });

    it('should support string range queries', () => {
      const strTree = new BPlusTree<string, number>();
      const words = ['apple', 'banana', 'cherry', 'date', 'elderberry', 'fig'];
      words.forEach((w, i) => strTree.insert(w, i));

      const result = strTree.range('banana', 'elderberry');
      expect(result.map((e) => e.key)).toEqual(['banana', 'cherry', 'date']);
    });
  });
});