import fs from 'node:fs';
import os from 'node:os';
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';
import { DeltaShardBinary } from '../src/core/shards/DeltaShardBinary.js';
import { WaveletShardBinary } from '../src/core/shards/WaveletShardBinary.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 makeZipfSampler(count, skew, seed = 1337) {
  const weights = new Float64Array(count);
  let sum = 0;
  for (let i = 1; i <= count; i++) {
    const w = 1 / Math.pow(i, skew);
    weights[i - 1] = w;
    sum += w;
  }
  const cdf = new Float64Array(count);
  let acc = 0;
  for (let i = 0; i < count; i++) {
    acc += weights[i] / sum;
    cdf[i] = acc;
  }
  let state = seed >>> 0;
  const rand = () => {
    state = (state * 1664525 + 1013904223) >>> 0;
    return state / 0xffffffff;
  };
  return () => {
    const r = rand();
    let lo = 0;
    let hi = cdf.length - 1;
    while (lo < hi) {
      const mid = (lo + hi) >> 1;
      if (r <= cdf[mid]) {
        hi = mid;
      } else {
        lo = mid + 1;
      }
    }
    return lo;
  };
}

function buildGraph({ edges, users, docs, seed }) {
  const graph = new CondensedGraph();
  const rng = makeRng(seed);
  for (let i = 0; i < edges; i++) {
    const userId = Math.floor(rng() * users);
    const docId = Math.floor(rng() * docs);
    graph.addEdge(`user:${userId}`, 'owner', `doc:${docId}`);
  }
  graph.finalizePerfectHash();
  graph.finalizeWaveletAdjacency({ dropAdjacencyList: true });
  return graph;
}

function buildDeltaLayer(snapshot, relationId, deltaEdges, seed, dir, skew) {
  const rng = makeRng(seed);
  const nodeCount = snapshot.nodeCount;
  const sampler = skew ? makeZipfSampler(nodeCount, skew, seed + 11) : null;
  const buckets = new Map();
  let addCount = 0;
  let removeCount = 0;
  for (let i = 0; i < deltaEdges; i++) {
    const srcId = sampler ? sampler() : Math.floor(rng() * nodeCount);
    const dstId = sampler ? sampler() : Math.floor(rng() * nodeCount);
    const shardMeta = snapshot._selectShardMeta(relationId, 'out', srcId);
    if (!shardMeta) continue;
    const localSource = snapshot._localSource(srcId, shardMeta);
    const key = shardMeta.cacheKey;
    let entry = buckets.get(key);
    if (!entry) {
      entry = { shardMeta, additions: [], removals: [] };
      buckets.set(key, entry);
    }
    if (rng() < 0.5) {
      entry.removals.push({ srcLocal: localSource, otherId: dstId });
      removeCount++;
    } else {
      entry.additions.push({ srcLocal: localSource, otherId: dstId, possBits: 65535, relBits: 65535 });
      addCount++;
    }
  }

  fs.mkdirSync(dir, { recursive: true });
  const shards = [];
  for (const entry of buckets.values()) {
    const shardKey = `delta-${entry.shardMeta.key}`;
    const buffer = DeltaShardBinary.serialize({
      relationId: entry.shardMeta.relationId,
      direction: entry.shardMeta.direction,
      rangeStart: entry.shardMeta.rangeStart,
      rangeEnd: entry.shardMeta.rangeEnd,
      nodeCount: snapshot.nodeCount,
      additions: entry.additions,
      removals: entry.removals
    });
    fs.writeFileSync(path.join(dir, shardKey), new Uint8Array(buffer));
    shards.push({
      key: shardKey,
      relationId: entry.shardMeta.relationId,
      direction: entry.shardMeta.direction,
      rangeStart: entry.shardMeta.rangeStart,
      rangeEnd: entry.shardMeta.rangeEnd,
      cacheKey: entry.shardMeta.cacheKey
    });
  }
  return { shards, storage: new FileShardStorage(dir), addCount, removeCount };
}

function buildQueries(snapshot, relationId, count, seed) {
  const rng = makeRng(seed);
  const queries = [];
  for (let i = 0; i < count; i++) {
    const srcId = Math.floor(rng() * snapshot.nodeCount);
    const edges = snapshot.getOutEdgesSync(srcId, relationId);
    if (edges.length) {
      const edge = edges[Math.floor(rng() * edges.length)];
      queries.push([srcId, edge.dst]);
    } else {
      const dstId = Math.floor(rng() * snapshot.nodeCount);
      queries.push([srcId, dstId]);
    }
  }
  return queries;
}

function timeFindEdges(snapshot, relationId, queries) {
  const start = performance.now();
  let found = 0;
  for (const [srcId, dstId] of queries) {
    if (snapshot.findEdgeSync(srcId, relationId, dstId)) found++;
  }
  const duration = performance.now() - start;
  return { duration, found };
}

function compactLayer(snapshot, layer, outputDir) {
  fs.mkdirSync(outputDir, { recursive: true });
  let compacted = 0;
  for (const shardMeta of layer.shards) {
    const base = snapshot._cacheIndex.get(shardMeta.cacheKey);
    if (!base) continue;
    const shard = snapshot._loadShardSync(base.relationId, base.direction, base.rangeStart);
    if (!shard) continue;

    const deltaBuffer = layer.storage.getSync(shardMeta.key);
    if (!deltaBuffer) continue;
    const deltaShard = DeltaShardBinary.deserialize(deltaBuffer);
    const rangeSize = shard.rangeEnd - shard.rangeStart;
    const sources = new Array(rangeSize);
    for (let localSource = 0; localSource < rangeSize; localSource++) {
      const range = snapshot._rangeForSource(shard, localSource);
      const list = [];
      if (range) {
        for (let pos = range.start; pos < range.end; pos++) {
          list.push({ otherId: shard.dstIds[pos], possBits: shard.possBits[pos], relBits: shard.relBits[pos] });
        }
      }
      sources[localSource] = list;
    }

    for (const removal of deltaShard.removals) {
      const list = sources[removal.srcLocal];
      if (!list) continue;
      const idx = list.findIndex((item) => item.otherId === removal.otherId);
      if (idx !== -1) list.splice(idx, 1);
    }

    for (const addition of deltaShard.additions) {
      const list = sources[addition.srcLocal] || (sources[addition.srcLocal] = []);
      list.push({ otherId: addition.otherId, possBits: addition.possBits, relBits: addition.relBits });
    }

    const buffer = WaveletShardBinary.serialize({
      relationId: shard.relationId,
      direction: shard.direction,
      rangeStart: shard.rangeStart,
      rangeEnd: shard.rangeEnd,
      nodeCount: snapshot.nodeCount,
      sources
    });
    fs.writeFileSync(path.join(outputDir, base.key), new Uint8Array(buffer));
    compacted++;
  }
  return compacted;
}

function sumDeltaBytes(dir) {
  let total = 0;
  const entries = fs.readdirSync(dir, { withFileTypes: true });
  for (const entry of entries) {
    if (!entry.isFile()) continue;
    const stat = fs.statSync(path.join(dir, entry.name));
    total += stat.size;
  }
  return total;
}

function sumShardBytes(dir, manifest) {
  let total = 0;
  if (manifest.nodeTableKey) {
    const nodeTablePath = path.join(dir, manifest.nodeTableKey);
    if (fs.existsSync(nodeTablePath)) total += fs.statSync(nodeTablePath).size;
  }
  if (manifest.componentKey) {
    const componentPath = path.join(dir, manifest.componentKey);
    if (fs.existsSync(componentPath)) total += fs.statSync(componentPath).size;
  }
  for (const shard of manifest.shards || []) {
    const shardPath = path.join(dir, shard.key);
    if (fs.existsSync(shardPath)) total += fs.statSync(shardPath).size;
  }
  return total;
}

function shardSizeStats(dir, manifest) {
  const sizes = [];
  for (const shard of manifest.shards || []) {
    const shardPath = path.join(dir, shard.key);
    if (fs.existsSync(shardPath)) sizes.push(fs.statSync(shardPath).size);
  }
  sizes.sort((a, b) => a - b);
  if (!sizes.length) return { p50: 0, p95: 0, p99: 0, max: 0, min: 0 };
  const pct = (p) => sizes[Math.min(sizes.length - 1, Math.floor(p * sizes.length))];
  return {
    p50: pct(0.5),
    p95: pct(0.95),
    p99: pct(0.99),
    max: sizes[sizes.length - 1],
    min: sizes[0]
  };
}


const args = parseArgs(process.argv);
const edges = Number(args.get('edges') || 200000);
const users = Number(args.get('users') || 50000);
const docs = Number(args.get('docs') || 100000);
const bucketSize = Number(args.get('bucket') || 65536);
const queryCount = Number(args.get('queries') || 20000);
const seed = Number(args.get('seed') || 1337);
const deltaPowers = String(args.get('delta-powers') || '0,1,2,3').split(',').map((v) => Number(v.trim()));
const compactAt = Number(args.get('compact-at') || 0);
const deltaSkew = Number(args.get('delta-skew') || 0);
const shardMode = String(args.get('shard-mode') || 'chain');
const componentRelations = args.get('component-relations')
  ? String(args.get('component-relations')).split(',').map((v) => v.trim()).filter(Boolean)
  : ['owner'];

console.log('Sharded delta overlay bench');
console.log(`  edges: ${edges}`);
console.log(`  users: ${users}`);
console.log(`  docs: ${docs}`);
console.log(`  queries: ${queryCount}`);
console.log(`  delta powers: ${deltaPowers.join(',')}`);
if (compactAt > 0) console.log(`  compact at: ${compactAt}`);
if (deltaSkew > 0) console.log(`  delta skew: ${deltaSkew}`);
console.log(`  shard mode: ${shardMode}`);
console.log(`  component relations: ${componentRelations.join(',')}`);
console.log('  note: timings are per-query averages; expect noise and cache effects. Compare trends, not single points.');

const graph = buildGraph({ edges, users, docs, seed });
const baseDir = fs.mkdtempSync(path.join(os.tmpdir(), 'sharded-delta-bench-'));
let shardDir = path.join(baseDir, 'base');

const builder = new ShardedSnapshotBuilder({
  bucketSize,
  includeDirections: ['out', 'in'],
  shardMode,
  componentRelations
});
const manifest = builder.build(graph, shardDir);
let storage = new FileShardStorage(shardDir);
let snapshot = new ShardedSnapshot(manifest, storage, { cacheLimit: 64, recentLimit: 256 });
snapshot.initializeSync();

const relationId = snapshot.relationIdToName.indexOf('owner');
const queries = buildQueries(snapshot, relationId, queryCount, seed + 9);

console.log('Delta size vs avg findEdge time');
console.log('edges | base_bytes | shard_files | shard_p50 | shard_p95 | shard_p99 | deltaEdges | avg_us_per_query | delta_bytes | compact_ms | compact_us_per_edge');

let baseEdges = edges;
const baseline = timeFindEdges(snapshot, relationId, queries);
const baseBytes = sumShardBytes(shardDir, manifest);
const shardFiles = countShardFiles(manifest);
const shardStats = shardSizeStats(shardDir, manifest);
console.log(`${baseEdges} | ${baseBytes} | ${shardFiles} | ${shardStats.p50} | ${shardStats.p95} | ${shardStats.p99} | 0 | ${(baseline.duration / queryCount) * 1000} | 0 | - | -`);

for (const power of deltaPowers) {
  const deltaEdges = Math.max(1, Math.floor(Math.pow(10, power)));
  const layerDir = path.join(baseDir, `delta-${power}`);
  const layer = buildDeltaLayer(snapshot, relationId, deltaEdges, seed + power + 1, layerDir, deltaSkew);
  snapshot.setDeltaLayers([layer]);
  snapshot.clearDeltaCache();
  const result = timeFindEdges(snapshot, relationId, queries);
  const avgUs = (result.duration / queryCount) * 1000;
  const deltaBytes = sumDeltaBytes(layerDir);
  let baseBytesNow = baseBytes;
  let shardFilesNow = shardFiles;
  let shardStatsNow = shardStats;
  let compactMs = '-';
  let compactUsPerEdge = '-';
  if (compactAt > 0 && deltaEdges >= compactAt) {
    const compactDir = path.join(baseDir, `compact-${power}`);
    const start = performance.now();
    compactLayer(snapshot, layer, compactDir);
    const totalMs = performance.now() - start;
    compactMs = totalMs.toFixed(2);
    const denom = Math.max(1, layer.addCount + layer.removeCount);
    compactUsPerEdge = ((totalMs * 1000) / denom).toFixed(3);

    baseEdges = Math.max(0, baseEdges + layer.addCount - layer.removeCount);
    const newBaseDir = path.join(baseDir, `base-${power}`);
    fs.rmSync(newBaseDir, { recursive: true, force: true });
    fs.mkdirSync(newBaseDir, { recursive: true });

    if (manifest.nodeTableKey) {
      fs.copyFileSync(path.join(shardDir, manifest.nodeTableKey), path.join(newBaseDir, manifest.nodeTableKey));
    }
    if (manifest.componentKey) {
      fs.copyFileSync(path.join(shardDir, manifest.componentKey), path.join(newBaseDir, manifest.componentKey));
    }

    for (const shard of manifest.shards || []) {
      const basePath = path.join(shardDir, shard.key);
      const compactPath = path.join(compactDir, shard.key);
      const sourcePath = fs.existsSync(compactPath) ? compactPath : basePath;
      fs.copyFileSync(sourcePath, path.join(newBaseDir, shard.key));
    }

    shardDir = newBaseDir;
    storage = new FileShardStorage(shardDir);
    snapshot = new ShardedSnapshot(manifest, storage, { cacheLimit: 64, recentLimit: 256 });
    snapshot.initializeSync();
    baseBytesNow = sumShardBytes(shardDir, manifest);
    shardFilesNow = countShardFiles(manifest);
    shardStatsNow = shardSizeStats(shardDir, manifest);
  }
  console.log(`${baseEdges} | ${baseBytesNow} | ${shardFilesNow} | ${shardStatsNow.p50} | ${shardStatsNow.p95} | ${shardStatsNow.p99} | ${deltaEdges} | ${avgUs} | ${deltaBytes} | ${compactMs} | ${compactUsPerEdge}`);
}

fs.rmSync(baseDir, { recursive: true, force: true });
function countShardFiles(manifest) {
  let total = 0;
  if (manifest.nodeTableKey) total += 1;
  if (manifest.componentKey) total += 1;
  total += (manifest.shards || []).length;
  return total;
}