import { performance } from 'node:perf_hooks';
import { CondensedGraph } from '../src/core/CondensedGraph.js';
import { WaveletTree } from '../src/core/graph/succinct/WaveletTree.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') || 200000);
const users = Number(args.get('users') || 10000);
const docs = Number(args.get('docs') || 50000);
const bucketSize = Number(args.get('bucket') || 4096);

console.log('Shard size bench (relation + source-range)');
console.log(`  edges: ${edges}`);
console.log(`  users: ${users}`);
console.log(`  docs: ${docs}`);
console.log(`  bucket size: ${bucketSize}`);

const graph = new CondensedGraph();
const relations = ['owner', 'editor', 'viewer'];

const buildStart = performance.now();
for (let i = 0; i < edges; i++) {
  graph.addEdge(`user:${i % users}`, relations[i % 3], `doc:${i % docs}`);
}
const buildTime = performance.now() - buildStart;
console.log(`  build time: ${buildTime.toFixed(2)} ms`);

const relCount = graph.nextRelationId;
const bucketCount = Math.ceil(users / bucketSize);
const shardBuckets = new Array(relCount);
for (let r = 0; r < relCount; r++) {
  shardBuckets[r] = new Array(bucketCount);
}

const adjacency = graph.adjacency;
for (let edgeIdx = 0; edgeIdx < graph.edgeIndex; edgeIdx++) {
  const base = edgeIdx * 6;
  const srcId = adjacency[base + 0];
  const relId = adjacency[base + 1];
  const dstId = adjacency[base + 2];
  const bucket = Math.floor(srcId / bucketSize);
  if (!shardBuckets[relId][bucket]) {
    shardBuckets[relId][bucket] = { dstsBySrc: [], edgeCount: 0 };
  }
  const shard = shardBuckets[relId][bucket];
  const local = srcId - bucket * bucketSize;
  let list = shard.dstsBySrc[local];
  if (!list) {
    list = [];
    shard.dstsBySrc[local] = list;
  }
  list.push(dstId);
  shard.edgeCount++;
}

function sizeWaveletAdjacency(params) {
  const {
    sourceCount,
    edgeCount,
    boundaryWordCount,
    boundaryBlockRankLength,
    nodeCount,
    totalWords,
    totalRank
  } = params;
  let offset = 0;
  const align = (n) => {
    const pad = (n - (offset % n)) % n;
    offset += pad;
  };

  offset += 4; // numNodes
  offset += 4; // boundary length
  offset += 4; // boundary word count
  align(4);
  offset += boundaryWordCount * 4;
  offset += 4; // blockSize
  offset += 4; // blockRank length
  align(4);
  offset += boundaryBlockRankLength * 4;
  offset += 4; // edgeCount
  align(4);
  offset += edgeCount * 4; // edge indices
  offset += 1; // hasTree

  if (nodeCount > 0) {
    offset += 4; // nodeCount
    align(4);
    offset += nodeCount * 4 * 11; // 11 uint32/int32 arrays
    offset += 4; // isLeaf length
    offset += nodeCount; // isLeaf bytes
    offset += 4; // words length
    align(4);
    offset += totalWords * 4;
    offset += 4; // rank length
    align(4);
    offset += totalRank * 4;
  }

  return offset;
}

function computeWaveletFlatSize(dstSequence, alphabetSize) {
  if (!dstSequence.length) {
    return { nodeCount: 0, totalWords: 0, totalRank: 0 };
  }
  const tree = new WaveletTree(new Uint32Array(dstSequence), alphabetSize);
  const stack = [tree];
  let nodeCount = 0;
  let totalWords = 0;
  let totalRank = 0;
  while (stack.length) {
    const node = stack.pop();
    if (!node) continue;
    nodeCount++;
    if (node.bitvector) {
      totalWords += node.bitvector.bitvector.length;
      totalRank += node.bitvector.rankSelect.blockRank.length;
    }
    if (node.leftChild) stack.push(node.leftChild);
    if (node.rightChild) stack.push(node.rightChild);
  }
  return { nodeCount, totalWords, totalRank };
}

const shardSizes = [];
const alphabetSize = graph.numNodes;
const blockSize = 512;

for (let relId = 0; relId < relCount; relId++) {
  for (let b = 0; b < bucketCount; b++) {
    const shard = shardBuckets[relId][b];
    if (!shard || shard.edgeCount === 0) continue;
    const sourceStart = b * bucketSize;
    const sourceCount = Math.min(bucketSize, users - sourceStart);
    const dstSequence = [];
    for (let i = 0; i < sourceCount; i++) {
      const list = shard.dstsBySrc[i];
      if (!list) continue;
      for (let j = 0; j < list.length; j++) {
        dstSequence.push(list[j]);
      }
    }

    const boundaryLength = sourceCount + shard.edgeCount + 1;
    const boundaryWordCount = Math.ceil(boundaryLength / 32);
    const boundaryBlockRankLength = Math.ceil(boundaryLength / blockSize) + 1;
    const { nodeCount, totalWords, totalRank } = computeWaveletFlatSize(dstSequence, alphabetSize);
    const bytes = sizeWaveletAdjacency({
      sourceCount,
      edgeCount: shard.edgeCount,
      boundaryWordCount,
      boundaryBlockRankLength,
      nodeCount,
      totalWords,
      totalRank
    });
    shardSizes.push({ relId, bucket: b, edges: shard.edgeCount, bytes });
  }
}

shardSizes.sort((a, b) => a.bytes - b.bytes);
const totalBytes = shardSizes.reduce((sum, shard) => sum + shard.bytes, 0);
const p = (q) => shardSizes[Math.floor((shardSizes.length - 1) * q)] || shardSizes[0];
const toMb = (bytes) => (bytes / 1024 / 1024).toFixed(2);

console.log('\nShard sizes');
console.log(`  shards: ${shardSizes.length}`);
console.log(`  total: ${toMb(totalBytes)} MB`);
console.log(`  avg: ${toMb(totalBytes / shardSizes.length)} MB`);
console.log(`  p50: ${toMb(p(0.5).bytes)} MB`);
console.log(`  p95: ${toMb(p(0.95).bytes)} MB`);
console.log(`  p99: ${toMb(p(0.99).bytes)} MB`);
console.log(`  max: ${toMb(shardSizes[shardSizes.length - 1].bytes)} MB`);
console.log(`  bytes/edge (avg): ${(totalBytes / edges).toFixed(2)}`);