import fs from 'node:fs';
import path from 'node:path';
import { performance } from 'node:perf_hooks';
import { CondensedGraph } from '../src/core/CondensedGraph.js';
import { ShardedSnapshotBuilder } from '../src/core/shards/ShardedSnapshotBuilder.js';
import { ShardedSnapshot } from '../src/core/shards/ShardedSnapshot.js';
import { FileShardStorage } from '../src/core/shards/FileShardStorage.js';

function parseArgs(argv) {
  const args = new Map();
  for (let i = 2; i < argv.length; i++) {
    const value = argv[i];
    if (!value.startsWith('--')) continue;
    const [key, inline] = value.slice(2).split('=');
    if (inline !== undefined) {
      args.set(key, inline);
      continue;
    }
    const next = argv[i + 1];
    if (next && !next.startsWith('--')) {
      args.set(key, next);
      i++;
    } else {
      args.set(key, true);
    }
  }
  return args;
}

const args = parseArgs(process.argv);
const edges = Number(args.get('edges') || 50000);
const users = Number(args.get('users') || 5000);
const docs = Number(args.get('docs') || 20000);
const bucketSize = Number(args.get('bucket') || 4096);
const samples = Number(args.get('samples') || 10000);
const output = args.get('output') || 'tmp/shards-bench';

console.log('Sharded snapshot bench');
console.log(`  edges: ${edges}`);
console.log(`  users: ${users}`);
console.log(`  docs: ${docs}`);
console.log(`  bucket: ${bucketSize}`);
console.log(`  samples: ${samples}`);

const graph = new CondensedGraph();
const relations = ['owner', 'editor', 'viewer'];
const edgePairs = [];
const buildStart = performance.now();
for (let i = 0; i < edges; i++) {
  const userIndex = i % users;
  const docIndex = i % docs;
  const relId = i % 3;
  const srcId = graph._ensureNode(`user:${userIndex}`);
  const dstId = graph._ensureNode(`doc:${docIndex}`);
  graph.addEdge(srcId, relations[relId], dstId);
  edgePairs.push({ srcId, dstId, relId });
}
const buildTime = performance.now() - buildStart;

graph.finalizePerfectHash();
graph.finalizeWaveletAdjacency({ dropAdjacencyList: true });
const snapshot = graph.toBinary();
const snapshotPath = path.join(output, 'snapshot.bin');
fs.mkdirSync(output, { recursive: true });
fs.writeFileSync(snapshotPath, new Uint8Array(snapshot));

const shardStart = performance.now();
const builder = new ShardedSnapshotBuilder({ bucketSize, includeDirections: ['out'] });
const manifest = builder.build(graph, output);
const shardTime = performance.now() - shardStart;

const storage = new FileShardStorage(output);
const sharded = new ShardedSnapshot(manifest, storage, { cacheLimit: 8 });
sharded.initializeSync();

const srcIds = Array.from({ length: users }, (_, i) => i);
const dstIds = Array.from({ length: docs }, (_, i) => i + users);

const findStart = performance.now();
let hits = 0;
for (let i = 0; i < samples; i++) {
  const pair = edgePairs[i % edgePairs.length];
  const srcId = pair.srcId;
  const dstId = pair.dstId;
  const relId = pair.relId;
  const edge = sharded.findEdgeSync(srcId, relId, dstId);
  if (edge) hits++;
}
const findTime = performance.now() - findStart;

const outStart = performance.now();
let outEdges = 0;
for (let i = 0; i < samples; i++) {
  const srcId = srcIds[i % srcIds.length];
  const relId = i % 3;
  const edgesOut = sharded.getOutEdgesSync(srcId, relId);
  outEdges += edgesOut.length;
}
const outTime = performance.now() - outStart;

console.log('\nResults');
console.log(`  build time: ${buildTime.toFixed(2)} ms`);
console.log(`  shard time: ${shardTime.toFixed(2)} ms`);
console.log(`  shards: ${manifest.shards.length}`);
console.log(`  snapshot size: ${(snapshot.byteLength / 1024 / 1024).toFixed(2)} MB`);
console.log(`  findEdge avg: ${(findTime / samples * 1000).toFixed(3)} µs`);
console.log(`  getOutEdges avg: ${(outTime / samples * 1000).toFixed(3)} µs`);
console.log(`  hits: ${hits}`);
console.log(`  outEdges: ${outEdges}`);