import fs from 'node:fs';
import path from 'node:path';
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;
}

function makeRng(seed = 1337) {
  let state = seed >>> 0;
  return () => {
    state = (state * 1664525 + 1013904223) >>> 0;
    return state / 0xffffffff;
  };
}

function buildRelationMap(relationIdToName) {
  const map = new Map();
  for (let i = 0; i < relationIdToName.length; i++) {
    map.set(relationIdToName[i], i);
  }
  return map;
}

const args = parseArgs(process.argv);
const manifestPath = args.get('manifest') || 'tmp/shards-complex-1m-16384/manifest.json';
const queries = Number(args.get('queries') || 200);
const seed = Number(args.get('seed') || 1234);
const userPool = Number(args.get('user-pool') || 5000);
const docPool = Number(args.get('doc-pool') || 5000);

const manifest = JSON.parse(fs.readFileSync(manifestPath, 'utf8'));
const storage = new FileShardStorage(path.dirname(manifestPath));
const snapshot = new ShardedSnapshot(manifest, storage, { cacheLimit: 32, recentLimit: 128 });
snapshot.initializeSync();

const relationIds = buildRelationMap(manifest.relationIdToName || []);
const rng = makeRng(seed);

const userIds = [];
for (let i = 0; i < userPool; i++) {
  const id = snapshot.resolveNodeId(`user:${i}`);
  if (id !== -1) userIds.push(id);
}

const docIds = [];
for (let i = 0; i < docPool; i++) {
  const id = snapshot.resolveNodeId(`doc:${i}`);
  if (id !== -1) docIds.push(id);
}

const memberRel = relationIds.get('member');
const viewerRel = relationIds.get('viewer');

const chainPairs = [];
for (let i = 0; i < queries; i++) {
  const userId = userIds[Math.floor(rng() * userIds.length)];
  const docId = docIds[Math.floor(rng() * docIds.length)];
  chainPairs.push({ userId, docId });
}

console.log('Sharded plan/execute bench');
console.log(`  manifest: ${manifestPath}`);
console.log(`  queries: ${queries}`);

const plan = new Set();
for (const pair of chainPairs) {
  snapshot.planOutEdges(memberRel, pair.userId, plan);
  snapshot.planOutEdges(viewerRel, pair.userId, plan);
}

const planStart = process.hrtime.bigint();
snapshot.prefetchPlan(plan, (err) => {
  if (err) {
    console.error(err);
    process.exit(1);
  }
  const planTime = Number(process.hrtime.bigint() - planStart) / 1e6;
  console.log(`  prefetch time: ${planTime.toFixed(2)} ms`);
  console.log(`  plan size: ${plan.size}`);
  console.log('Cache stats');
  console.log(snapshot.getStats());
});