import { Arbiter } from '../src/core/Arbiter.js';
import fs from 'node:fs';
import { performance } from 'node:perf_hooks';

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 groups = Number(args.get('groups') || 500);
const samples = Number(args.get('samples') || 5000);
const snapshotPath = args.get('snapshot-path') || null;
const debug = args.has('debug');
const fastConstruction = args.has('no-fast') ? false : true;
const ensureHits = args.has('ensure-hits');
const fixedPair = ensureHits;
const buildIndices = args.has('no-indices') ? false : true;
const hitPairs = [];

console.log('Arbiter snapshot boot bench');
console.log(`  edges: ${edges}`);
console.log(`  users: ${users}`);
console.log(`  docs: ${docs}`);
console.log(`  groups: ${groups}`);
console.log(`  samples: ${samples}`);
if (snapshotPath) {
  console.log(`  snapshot path: ${snapshotPath}`);
}
console.log(`  fast construction: ${fastConstruction}`);
console.log(`  ensure hits: ${ensureHits}`);
console.log(`  fixed pair: ${fixedPair}`);
console.log(`  build indices: ${buildIndices}`);

let snapshotBuffer = null;

if (snapshotPath && fs.existsSync(snapshotPath)) {
  snapshotBuffer = fs.readFileSync(snapshotPath).buffer;
}

let buildTime = 0;
let snapshotTime = 0;

if (!snapshotBuffer) {
  const arbiter = new Arbiter({
    fastConstructionMode: fastConstruction,
    disableCaching: fastConstruction
  });

  if (debug) console.log('  building nodes...');
  const nodeStart = performance.now();
  for (let i = 0; i < users; i++) {
    arbiter.addNode(`user:${i}`, 'user');
  }
  for (let i = 0; i < docs; i++) {
    arbiter.addNode(`doc:${i}`, 'doc');
  }
  for (let i = 0; i < groups; i++) {
    arbiter.addNode(`group:${i}`, 'group');
  }
  const nodeTime = performance.now() - nodeStart;
  if (debug) console.log(`  nodes ready (${nodeTime.toFixed(2)} ms)`);

  if (debug) console.log('  building relations...');
  const buildStart = performance.now();
  for (let i = 0; i < edges; i++) {
    const user = `user:${i % users}`;
    const doc = `doc:${i % docs}`;
    const group = `group:${i % groups}`;
    const rel = ['owner', 'editor', 'viewer'][i % 3];

    if (i % 5 === 0) {
      arbiter.addRelation(user, 'member', group);
    }
    if (i % 7 === 0) {
      arbiter.addRelation(group, 'viewer', doc);
    }
    arbiter.addRelation(user, rel, doc);
  }
  buildTime = performance.now() - buildStart;
  if (debug) console.log(`  relations ready (${buildTime.toFixed(2)} ms)`);

  if (debug) console.log('  compiling rules...');
  arbiter.setRelationConfig('member', { type: 'direct' });
  arbiter.setRelationConfig('viewer', { type: 'direct' });
  arbiter.setRelationConfig('risk', { type: 'direct' });
  arbiter.setRelationConfig('threshold', { type: 'direct' });

  arbiter.setRelationConfig('can_view', {
    union: {
      rules: [
        { type: 'direct', relation: 'owner' },
        { type: 'direct', relation: 'editor' },
        { type: 'tuple_to_userset', tuplesetRelation: 'viewer', tuplesetDirection: 'in', computedRelation: 'member' }
      ]
    }
  });

  arbiter.setRelationConfig('complex_view', {
    intersection: {
      rules: [
        { type: 'tuple_to_userset', tuplesetRelation: 'viewer', tuplesetDirection: 'in', computedRelation: 'member' },
        {
          type: 'chain',
          steps: [
            { relation: 'member', direction: 'out' },
            { relation: 'viewer', direction: 'out' }
          ]
        },
        {
          type: 'relational_comparator',
          comparator: '>=',
          left: {
            rule: { type: 'direct', relation: 'risk' },
            valueRelation: 'risk',
            aggregator: 'max',
            evaluateFrom: 'auto'
          },
          right: {
            rule: { type: 'direct', relation: 'threshold' },
            valueRelation: 'threshold',
            aggregator: 'max',
            evaluateFrom: 'auto'
          }
        }
      ]
    }
  });

  if (!ensureHits) {
    for (let i = 0; i < Math.min(20000, edges); i++) {
      arbiter.addRelation(`user:${i % users}`, 'risk', `doc:${i % docs}`, { value: (i % 100) / 100 });
      arbiter.addRelation(`user:${i % users}`, 'threshold', `doc:${i % docs}`, { value: 0.5 });
    }
  }

  if (ensureHits) {
    const hitCount = Math.min(200, users, docs, groups);
    for (let i = 0; i < hitCount; i++) {
      const user = `user:${i}`;
      const doc = `doc:${i}`;
      const group = `group:${i}`;
      arbiter.addRelation(user, 'member', group);
      arbiter.addRelation(group, 'viewer', doc);
      arbiter.addRelation(user, 'risk', doc, { value: 0.9 });
      arbiter.addRelation(user, 'threshold', doc, { value: 0.1 });
      hitPairs.push({ user, doc });
    }
  }

  if (debug) console.log('  snapshotting...');
  const snapshotStart = performance.now();
  snapshotBuffer = arbiter.toSnapshotBinary({ dropAdjacencyList: true });
  snapshotTime = performance.now() - snapshotStart;
  if (debug) console.log(`  snapshot ready (${snapshotTime.toFixed(2)} ms)`);

  console.log(`  node build time: ${nodeTime.toFixed(2)} ms`);
}

const snapshotSize = snapshotBuffer.byteLength;

if (snapshotPath && !fs.existsSync(snapshotPath)) {
  fs.writeFileSync(snapshotPath, new Uint8Array(snapshotBuffer));
}

const memBefore = process.memoryUsage();
const loadStart = performance.now();
const booted = Arbiter.fromSnapshotBinary(snapshotBuffer, { buildSnapshotIndices: buildIndices });
const loadTime = performance.now() - loadStart;
const memAfter = process.memoryUsage();

const usersArr = Array.from({ length: users }, (_, i) => `user:${i}`);
const docsArr = Array.from({ length: docs }, (_, i) => `doc:${i}`);
let effectiveHitPairs = hitPairs;
if (ensureHits && effectiveHitPairs.length === 0) {
  effectiveHitPairs = [{ user: 'user:0', doc: 'doc:0' }];
}

const hitUsers = ensureHits
  ? (fixedPair ? [effectiveHitPairs[0].user] : effectiveHitPairs.map(pair => pair.user))
  : usersArr;
const hitDocs = ensureHits
  ? (fixedPair ? [effectiveHitPairs[0].doc] : effectiveHitPairs.map(pair => pair.doc))
  : docsArr;
const hitUserIds = hitUsers.map(key => booted.resolveNodeId(key));
const hitDocIds = hitDocs.map(key => booted.resolveNodeId(key));
const complexConfig = booted.relationConfigs.get('complex_view');

const simpleStart = performance.now();
let simpleHits = 0;
for (let i = 0; i < samples; i++) {
  const user = hitUsers[i % hitUsers.length];
  const doc = hitDocs[i % hitDocs.length];
  const result = booted.check(user, 'can_view', doc, { fastPath: true });
  if (result?.possibility > 0) simpleHits++;
}
const simpleTime = performance.now() - simpleStart;

const complexStart = performance.now();
let complexHits = 0;
for (let i = 0; i < samples; i++) {
  const idx = i % hitUsers.length;
  const user = hitUsers[idx];
  const doc = hitDocs[idx];
  const userId = hitUserIds[idx];
  const docId = hitDocIds[idx];
  const result = booted.authChecker.ruleEvaluator.evaluateRule(
    userId,
    user,
    docId,
    doc,
    complexConfig,
    new Set(),
    'complex_view',
    { fastPath: false, collectValues: true, cacheRuleResult: false }
  );
  if (result?.possibility > 0) complexHits++;
}
const complexTime = performance.now() - complexStart;

console.log('\nResults');
console.log(`  build time: ${buildTime.toFixed(2)} ms`);
console.log(`  snapshot time: ${snapshotTime.toFixed(2)} ms`);
console.log(`  snapshot size: ${(snapshotSize / 1024 / 1024).toFixed(2)} MB`);
console.log(`  load time: ${loadTime.toFixed(2)} ms`);
console.log(`  memory rss before: ${(memBefore.rss / 1024 / 1024).toFixed(2)} MB`);
console.log(`  memory rss after: ${(memAfter.rss / 1024 / 1024).toFixed(2)} MB`);
console.log(`  memory heap before: ${(memBefore.heapUsed / 1024 / 1024).toFixed(2)} MB`);
console.log(`  memory heap after: ${(memAfter.heapUsed / 1024 / 1024).toFixed(2)} MB`);
console.log(`  simple avg: ${(simpleTime / samples * 1000).toFixed(3)} µs`);
console.log(`  complex avg: ${(complexTime / samples * 1000).toFixed(3)} µs`);
console.log(`  simple hits: ${simpleHits}`);
console.log(`  complex hits: ${complexHits}`);
if (ensureHits && debug) {
  const sampleUser = hitUsers[0];
  const sampleDoc = hitDocs[0];
  const sampleSimple = booted.check(sampleUser, 'can_view', sampleDoc, { fastPath: true });
  const sampleComplex = booted.check(sampleUser, 'complex_view', sampleDoc, { fastPath: false, collectValues: true, includeMeta: true });
  console.log(`  sample can_view possibility: ${sampleSimple?.possibility ?? 'null'}`);
  console.log(`  sample complex_view possibility: ${sampleComplex?.possibility ?? 'null'}`);
  if (sampleComplex?.meta) {
    console.log(`  sample complex meta: ${JSON.stringify(sampleComplex.meta)}`);
  }

  const complexConfig = booted.relationConfigs.get('complex_view');

  const userId = booted.resolveNodeId(sampleUser);
  const objectId = booted.resolveNodeId(sampleDoc);
  if (userId !== undefined && objectId !== undefined) {
    const evaluator = booted.authChecker.ruleEvaluator;
    const visited = new Set();
    const tupleRule = { type: 'tuple_to_userset', tuplesetRelation: 'viewer', tuplesetDirection: 'in', computedRelation: 'member' };
    const chainRule = {
      type: 'chain',
      steps: [
        { relation: 'member', direction: 'out' },
        { relation: 'viewer', direction: 'out' }
      ]
    };
    const compRule = {
      type: 'relational_comparator',
      comparator: '>=',
      left: {
        rule: { type: 'direct', relation: 'risk' },
        valueRelation: 'risk',
        aggregator: 'max',
        evaluateFrom: 'auto'
      },
      right: {
        rule: { type: 'direct', relation: 'threshold' },
        valueRelation: 'threshold',
        aggregator: 'max',
        evaluateFrom: 'auto'
      }
    };
    const tupleResult = evaluator.evaluateRule(userId, sampleUser, objectId, sampleDoc, tupleRule, visited, 'tuple_debug', { fastPath: false, collectValues: true });
    const chainResult = evaluator.evaluateRule(userId, sampleUser, objectId, sampleDoc, chainRule, visited, 'chain_debug', { fastPath: false, collectValues: true });
    const compResult = evaluator.evaluateRule(userId, sampleUser, objectId, sampleDoc, compRule, visited, 'comp_debug', { fastPath: false, collectValues: true, includeMeta: true });
    console.log(`  tuple_to_userset possibility: ${tupleResult?.possibility ?? 'null'}`);
    console.log(`  chain possibility: ${chainResult?.possibility ?? 'null'}`);
    console.log(`  comparator possibility: ${compResult?.possibility ?? 'null'}`);
    if (compResult?.meta) {
      console.log(`  comparator meta: ${JSON.stringify(compResult.meta)}`);
    }
    const riskRel = booted.relationManager.getDirectRelation(userId, 'risk', objectId);
    const thresholdRel = booted.relationManager.getDirectRelation(userId, 'threshold', objectId);
    console.log(`  risk relation value: ${riskRel?.value ?? 'null'}`);
    console.log(`  threshold relation value: ${thresholdRel?.value ?? 'null'}`);

  if (complexConfig) {
    const compiled = complexConfig._compiled;
      if (compiled && compiled.type === 'logical') {
        const childTypes = compiled.children.map(child => child?.type || 'null');
        console.log(`  complex compiled children: ${JSON.stringify(childTypes)}`);

        const compiledEvaluator = evaluator.compiledEvaluator;
        const childPoss = compiled.children.map(child => {
          const res = compiledEvaluator.evaluate(child, userId, sampleUser, objectId, sampleDoc, new Set(), 'complex_child', { fastPath: false, collectValues: true, includeMeta: true });
          return res?.possibility ?? null;
        });
        console.log(`  complex compiled child possibilities: ${JSON.stringify(childPoss)}`);
      }
      const directComplex = evaluator.evaluateRule(userId, sampleUser, objectId, sampleDoc, complexConfig, new Set(), 'complex_debug', { fastPath: false, collectValues: true, includeMeta: true });
      console.log(`  direct complex possibility: ${directComplex?.possibility ?? 'null'}`);
      if (directComplex?.meta) {
        console.log(`  direct complex meta: ${JSON.stringify(directComplex.meta)}`);
      }
    }
  }
}