import { AESKey } from '../src/crypto/aes-key.js';
import { StorageMemory, EncryptedStorage, type BaseStorage } from '../src/storage/index.js';
import { BTreeCache, KvCache } from '../src/cache/index.js';

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

type Doc = {
  id: string;
  name: string;
  age: number;
  email: string;
  createdAt: Date;
};

/**
 * Generate a batch of unique documents.
 */
function generateDocs(count: number): Doc[] {
  const docs: Doc[] = [];
  for (let i = 0; i < count; i++) {
    docs.push({
      id: `id-${i}`,
      name: `user-${i}`,
      age: 20 + (i % 50),
      email: `user-${i}@test.com`,
      createdAt: new Date(Date.now() + i),
    });
  }
  return docs;
}

/**
 * Time an async operation and return elapsed milliseconds.
 */
async function time(fn: () => Promise<void>): Promise<number> {
  const start = performance.now();
  await fn();
  return performance.now() - start;
}

/**
 * Format ops/sec with thousands separators.
 */
function fmtOps(ops: number): string {
  return Math.round(ops).toLocaleString('en-US');
}

/**
 * Run a full suite of benchmarks against a given storage instance.
 */
async function benchmarkStorage(
  label: string,
  storage: BaseStorage<Doc>,
  docs: Doc[],
  options: { supportsRangeOps?: boolean; hasAgeIndex?: boolean } = {},
) {
  const { supportsRangeOps = true, hasAgeIndex = false } = options;
  const count = docs.length;
  console.log(`\n${'='.repeat(60)}`);
  console.log(` ${label}  (${count.toLocaleString()} documents)`);
  console.log('='.repeat(60));

  // --- Insert ---
  const insertMs = await time(async () => {
    await storage.insertMany(docs);
  });
  console.log(`  insertMany          ${insertMs.toFixed(2)}ms  (${fmtOps((count / insertMs) * 1000)} ops/sec)`);

  // --- Find all (no filter) ---
  const findAllMs = await time(async () => {
    await storage.find();
  });
  console.log(`  find()              ${findAllMs.toFixed(2)}ms  (${fmtOps((count / findAllMs) * 1000)} docs/sec)`);

  // --- Find by indexed field (equality) ---
  const lookupCount = Math.min(count, 1_000);
  const findIndexedMs = await time(async () => {
    for (let i = 0; i < lookupCount; i++) {
      await storage.findOne({ id: `id-${i}` });
    }
  });
  console.log(`  findOne indexed      ${findIndexedMs.toFixed(2)}ms  (${fmtOps((lookupCount / findIndexedMs) * 1000)} ops/sec)  [${lookupCount} lookups]`);

  // --- Find by non-indexed field (full scan) ---
  const scanCount = Math.min(count, 1_000);
  const findScanMs = await time(async () => {
    for (let i = 0; i < scanCount; i++) {
      await storage.findOne({ email: `user-${i}@test.com` });
    }
  });
  console.log(`  findOne scan         ${findScanMs.toFixed(2)}ms  (${fmtOps((scanCount / findScanMs) * 1000)} ops/sec)  [${scanCount} lookups]`);

  // --- Range queries ---
  if (supportsRangeOps) {
    // Wide range: 20% selectivity (10 out of 50 age values).
    const rangeCount = Math.min(count, 100);
    let rangeWideTotal = 0;
    const findRangeWideMs = await time(async () => {
      for (let i = 0; i < rangeCount; i++) {
        const results = await storage.find({ age: { $gte: 30, $lt: 40 } });
        rangeWideTotal += results.length;
      }
    });
    const indexLabel = hasAgeIndex ? 'B+Tree' : 'scan';
    console.log(`  find wide  [${indexLabel}]  ${findRangeWideMs.toFixed(2)}ms  (${fmtOps((rangeCount / findRangeWideMs) * 1000)} ops/sec)  [${rangeCount}x, ~${Math.round(rangeWideTotal / rangeCount)} hits, 20% sel.]`);

    // Narrow range: 2% selectivity (1 out of 50 age values).
    let rangeNarrowTotal = 0;
    const findRangeNarrowMs = await time(async () => {
      for (let i = 0; i < rangeCount; i++) {
        const results = await storage.find({ age: { $gte: 42, $lt: 43 } });
        rangeNarrowTotal += results.length;
      }
    });
    console.log(`  find narrow [${indexLabel}]  ${findRangeNarrowMs.toFixed(2)}ms  (${fmtOps((rangeCount / findRangeNarrowMs) * 1000)} ops/sec)  [${rangeCount}x, ~${Math.round(rangeNarrowTotal / rangeCount)} hits, 2% sel.]`);

    // --- Combined equality + operator ---
    const comboCount = Math.min(count, 1_000);
    const findComboMs = await time(async () => {
      for (let i = 0; i < comboCount; i++) {
        await storage.find({ id: `id-${i}`, age: { $gte: 20 } });
      }
    });
    console.log(`  find idx+operator    ${findComboMs.toFixed(2)}ms  (${fmtOps((comboCount / findComboMs) * 1000)} ops/sec)  [${comboCount} queries]`);
  }

  // --- Update by indexed field ---
  const updateCount = Math.min(count, 1_000);
  const updateMs = await time(async () => {
    for (let i = 0; i < updateCount; i++) {
      await storage.updateOne({ id: `id-${i}` }, { name: `updated-${i}` });
    }
  });
  console.log(`  updateOne indexed    ${updateMs.toFixed(2)}ms  (${fmtOps((updateCount / updateMs) * 1000)} ops/sec)  [${updateCount} updates]`);

  // --- Delete by indexed field ---
  const deleteCount = Math.min(count, 1_000);
  const deleteMs = await time(async () => {
    for (let i = 0; i < deleteCount; i++) {
      await storage.deleteOne({ id: `id-${i}` });
    }
  });
  console.log(`  deleteOne indexed    ${deleteMs.toFixed(2)}ms  (${fmtOps((deleteCount / deleteMs) * 1000)} ops/sec)  [${deleteCount} deletes]`);

  // --- Verify remaining count ---
  const remaining = await storage.find();
  console.log(`  remaining docs:      ${remaining.length.toLocaleString()}`);
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

// ---------------------------------------------------------------------------
// StorageMemory — B+ Tree range queries vs full scan
// ---------------------------------------------------------------------------

const DOC_COUNTS = [1_000, 10_000, 50_000];

for (const count of DOC_COUNTS) {
  const docs = generateDocs(count);

  // Indexes on id, name, AND age with explicit B+ Tree cache.
  const indexedWithAgeBTree = StorageMemory.from<Doc>(
    ['id', 'name', 'age'],
    new BTreeCache(),
  );
  await benchmarkStorage('StorageMemory (B+Tree cache, indexed: id,name,age)', indexedWithAgeBTree, docs, { hasAgeIndex: true });

  // Same indexes, but KV cache (no range support).
  const indexedWithAgeKv = StorageMemory.from<Doc>(
    ['id', 'name', 'age'],
    new KvCache(),
  );
  await benchmarkStorage('StorageMemory (KV cache, indexed: id,name,age)', indexedWithAgeKv, docs);

  // Indexes on id, name only — range queries on age fall back to full scan.
  const indexed = StorageMemory.from<Doc>(['id', 'name']);
  await benchmarkStorage('StorageMemory (indexed: id,name)', indexed, docs);

  // No indexes at all.
  const noIndex = StorageMemory.from<Doc>();
  await benchmarkStorage('StorageMemory (no indexes)', noIndex, docs);
}

// ---------------------------------------------------------------------------
// EncryptedStorage
// ---------------------------------------------------------------------------

const ENCRYPTED_DOC_COUNTS = [100, 1_000];
const encryptionKey = await AESKey.fromSeed('benchmark-key');

for (const count of ENCRYPTED_DOC_COUNTS) {
  const docs = generateDocs(count);

  // Indexed + plaintextKeys (age) — range queries on age use B+ Tree via backing store.
  const encBaseA = StorageMemory.from<Record<string, any>>(['id', 'name', 'age']);
  const encA = EncryptedStorage.from<Doc>(encBaseA, encryptionKey, {
    plaintextKeys: ['age'],
  });
  await benchmarkStorage('Encrypted (indexed+age, plaintextKeys: age)', encA, docs, { hasAgeIndex: true });

  // Indexed, fully encrypted — no range ops.
  const encBaseB = StorageMemory.from<Record<string, any>>(['id', 'name']);
  const encB = EncryptedStorage.from<Doc>(encBaseB, encryptionKey);
  await benchmarkStorage('Encrypted (indexed, fully encrypted)', encB, docs, { supportsRangeOps: false });

  // No indexes, fully encrypted — worst case.
  const encBaseC = StorageMemory.from<Record<string, any>>();
  const encC = EncryptedStorage.from<Doc>(encBaseC, encryptionKey);
  await benchmarkStorage('Encrypted (no indexes, fully encrypted)', encC, docs, { supportsRangeOps: false });
}

console.log('\nDone.\n');