import { AESKey } from '../src/crypto/aes-key.js';
import { StorageMemory, EncryptedStorage, type BaseStorage } from '../src/storage/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[], supportsRangeOps = true) {
  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]`);

  // --- Find with $gte / $lt range (full scan) ---
  if (supportsRangeOps) {
    const rangeCount = Math.min(count, 100);
    let rangeTotal = 0;
    const findRangeMs = await time(async () => {
      for (let i = 0; i < rangeCount; i++) {
        const results = await storage.find({ age: { $gte: 30, $lt: 40 } });
        rangeTotal += results.length;
      }
    });
    console.log(`  find $gte/$lt      ${findRangeMs.toFixed(2)}ms  (${fmtOps((rangeCount / findRangeMs) * 1000)} ops/sec)  [${rangeCount} queries, ~${Math.round(rangeTotal / rangeCount)} hits/query]`);

    // --- Find with combined equality + operator (index narrows, operator verifies) ---
    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 — indexed vs non-indexed
// ---------------------------------------------------------------------------

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

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

  const indexed = StorageMemory.from<Doc>(['id', 'name']);
  await benchmarkStorage('StorageMemory (indexed: id, name)', indexed, docs);

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

// ---------------------------------------------------------------------------
// EncryptedStorage — with plaintextKeys for range queries
// ---------------------------------------------------------------------------

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

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

  // No indexes + plaintextKeys — range queries on age work.
  const encBaseNoIdx = StorageMemory.from<Record<string, any>>();
  const encNoIdx = EncryptedStorage.from<Doc>(encBaseNoIdx, encryptionKey, {
    plaintextKeys: ['age'],
  });
  await benchmarkStorage('Encrypted (no indexes, plaintextKeys: age)', encNoIdx, docs);

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

  // Indexed, fully encrypted — same indexes but no plaintext keys.
  const encBaseB = StorageMemory.from<Record<string, any>>(['id', 'name']);
  const encB = EncryptedStorage.from<Doc>(encBaseB, encryptionKey);
  await benchmarkStorage('Encrypted (indexed, fully encrypted)', encB, docs, 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, false);
}

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