// registry-v3.mjs - the Aere validator registry v3 (AIP-22 stage 3), verified in plain JavaScript.
//
// A THIRD implementation, written from the canonical forms in the AIP (section 3.2) and in
// PqRegistryV3.java, sharing no code with either client: Falcon-512 is checked with falcon512.mjs
// (verification from the public spec), SLH-DSA-SHA2-128s and secp256k1 with @noble. Its job is to
// read the same conformance vectors as the two clients (aerenew/vectors/registry-v3/) and reach the
// same outcome on every one: "ok" with the same registry hash and row hashes, or the same named
// refusal (AERE-PQC-REG3-xx). Three independent implementations agreeing byte for byte is what
// makes the forms a specification rather than one client's habit.
//
//   node registry-v3.mjs <dir with the vectors>      exit 0: every vector as named; 1: a mismatch
//   import { load, rowHash, merkleRoot, registryHash, idOf } from './registry-v3.mjs'
import fs from 'node:fs';
import path from 'node:path';
import { keccak_256 } from '@noble/hashes/sha3';
import { secp256k1 } from '@noble/curves/secp256k1';
import { verifyFalcon512 } from './falcon512.mjs';

let slh = null;
try { slh = (await import('@noble/post-quantum/slh-dsa.js')).slh_dsa_sha2_128s; } catch { slh = null; }

const keccak = (b) => Buffer.from(keccak_256(b));
const u32 = (v) => { const b = Buffer.alloc(4); b.writeUInt32BE(v); return b; };
const u64 = (v) => { const b = Buffer.alloc(8); b.writeBigUInt64BE(BigInt(v)); return b; };
const ascii = (s) => Buffer.from(s, 'ascii');

export const SCHEMES = {
  'falcon-512': { wire: 1, pkLen: 897, verify: (pk, m, s) => { try { return verifyFalcon512(new Uint8Array(pk), new Uint8Array(m), new Uint8Array(s)); } catch { return false; } } },
  'slh-dsa-sha2-128s': { wire: 2, pkLen: 32, verify: (pk, m, s) => { if (!slh) throw new Error('NEMASURAT: @noble/post-quantum is not installed; SLH-DSA cannot be checked here'); try { return slh.verify(new Uint8Array(s), new Uint8Array(m), new Uint8Array(pk)); } catch { return false; } } },
};

export class RegistryV3Error extends Error {
  constructor(code, message) { super(code + ': ' + message); this.code = code; }
}
const fail = (n, msg) => { throw new RegistryV3Error('AERE-PQC-REG3-' + String(n).padStart(2, '0'), msg); };

export const idOf = (falcon897) => keccak(Buffer.concat([Buffer.from([0x50, 0x20]), falcon897])).subarray(12, 32);

export function popMessage(chainId, bindHeight, count, index, id, wire, pk) {
  return keccak(Buffer.concat([ascii('AERE-PQ-POP-3'), Buffer.from([3]), u64(chainId), u64(bindHeight), u32(count), u32(index), id,
    Buffer.from([wire]), u32(pk.length), pk]));
}

// row: { index, id, keys: [{ name, wire, pk, pop }] sorted by wire, transport: Buffer|null }
export function rowPreimage(chainId, bindHeight, count, row) {
  const parts = [ascii('AERE-PQ-ROW-3'), Buffer.from([3]), u64(chainId), u64(bindHeight), u32(count), u32(row.index), row.id,
    Buffer.from([row.keys.length])];
  for (const k of row.keys) parts.push(Buffer.from([k.wire]), u32(k.pk.length), k.pk, u32(k.pop.length), k.pop);
  parts.push(row.transport ? Buffer.concat([Buffer.from([1]), row.transport]) : Buffer.from([0]));
  return Buffer.concat(parts);
}
export const rowHash = (chainId, bindHeight, count, row) => keccak(rowPreimage(chainId, bindHeight, count, row));
const node = (l, r) => keccak(Buffer.concat([Buffer.from([0x01]), l, r]));
export function merkleRoot(leaves) {
  let level = leaves;
  while (level.length > 1) { const next = []; for (let i = 0; i < level.length; i += 2) next.push(node(level[i], level[i + 1] ?? level[i])); level = next; }
  return level[0];
}
export function merkleProof(leaves, index) {
  const proof = []; let level = leaves; let i = index;
  while (level.length > 1) {
    const sib = i % 2 === 0 ? i + 1 : i - 1; proof.push(level[sib] ?? level[i]);
    const next = []; for (let j = 0; j < level.length; j += 2) next.push(node(level[j], level[j + 1] ?? level[j])); level = next; i = Math.floor(i / 2);
  }
  return proof;
}
export function verifyMerkleProof(leaf, index, count, proof, root) {
  if (index < 0 || index >= count || count < 1) return false;
  let depth = 0; for (let n = count; n > 1; n = Math.ceil(n / 2)) depth++;
  if (proof.length !== depth) return false;
  let acc = leaf, i = index;
  for (const sib of proof) { acc = i % 2 === 0 ? node(acc, sib) : node(sib, acc); i = Math.floor(i / 2); }
  return acc.equals(root);
}
export const registryHash = (chainId, bindHeight, count, root) =>
  keccak(Buffer.concat([ascii('AERE-PQ-REGISTRY-3'), u64(chainId), u64(bindHeight), u32(count), root]));
export const rotateMessage = (rh) => keccak(Buffer.concat([ascii('AERE-PQ-ROTATE-3'), rh]));
export function claimDigest(chainId, rh) {
  const ctx = keccak(Buffer.concat([ascii('AERE-PQ-CLAIM-1'), u64(chainId)])).subarray(12, 32);
  return keccak(Buffer.concat([Buffer.from([0x19, 0x00]), ctx, rh]));
}
function recover(digest, sig65) {
  if (!sig65 || sig65.length !== 65) return null;
  let rec = sig65[64]; if (rec >= 27) rec -= 27;
  if (rec !== 0 && rec !== 1) return null;
  try {
    const s = secp256k1.Signature.fromCompact(sig65.subarray(0, 64)).addRecoveryBit(rec);
    return keccak(Buffer.from(s.recoverPublicKey(digest).toRawBytes(false).subarray(1))).subarray(12, 32);
  } catch { return null; }
}

function hex(v, len, code, what) {
  if (typeof v !== 'string' || !/^0x([0-9a-f]{2})*$/.test(v)) fail(code, what + ' is not lowercase 0x-hex of even length');
  const b = Buffer.from(v.slice(2), 'hex');
  if (len >= 0 && b.length !== len) fail(code, what + ' is ' + b.length + ' bytes, expected ' + len);
  return b;
}
const isObj = (o) => o !== null && typeof o === 'object' && !Array.isArray(o);

function parseRow(index, e) {
  for (const f of Object.keys(e)) if (!['id', 'keys', 'pop', 'rotation', 'ecdsa', 'transport'].includes(f)) fail(3, 'row ' + index + ' has an unknown field ' + f);
  const id = hex(e.id, 20, 4, 'row ' + index + ' id');
  if (!isObj(e.keys) || Object.keys(e.keys).length === 0) fail(3, 'row ' + index + ' has no keys');
  if (!isObj(e.pop)) fail(8, 'row ' + index + ' has no pop object');
  const keys = [];
  for (const [name, v] of Object.entries(e.keys)) {
    const s = SCHEMES[name]; if (!s) fail(5, 'row ' + index + ' names the unknown scheme ' + name);
    const pk = hex(v, s.pkLen, 7, 'row ' + index + ' ' + name);
    if (name === 'falcon-512' && pk[0] !== 0x09) fail(7, 'row ' + index + ' falcon-512: the standard encoding starts with 0x09');
    if (e.pop[name] === undefined) fail(8, 'row ' + index + ' ' + name + ' has no proof of possession');
    keys.push({ name, wire: s.wire, pk, pop: hex(e.pop[name], -1, 8, 'row ' + index + ' pop') });
  }
  if (Object.keys(e.pop).length !== Object.keys(e.keys).length) fail(8, 'row ' + index + ' carries a proof for a scheme it has no key for');
  keys.sort((a, b) => a.wire - b.wire);
  let transport = null;
  if (e.transport !== undefined) {
    if (!isObj(e.transport) || Object.keys(e.transport).length !== 1 || e.transport.secp256k1Address === undefined) fail(15, 'row ' + index + ' transport must be {secp256k1Address}');
    transport = hex(e.transport.secp256k1Address, 20, 15, 'row ' + index + ' transport');
  }
  const rotation = {};
  if (e.rotation !== undefined) {
    if (!isObj(e.rotation) || Object.keys(e.rotation).length === 0) fail(13, 'row ' + index + ' rotation must be a non-empty map');
    for (const [name, v] of Object.entries(e.rotation)) { if (!SCHEMES[name]) fail(5, 'row ' + index + ' rotation names an unknown scheme'); rotation[name] = hex(v, -1, 12, 'row ' + index + ' rotation'); }
  }
  let ecdsa = null;
  if (e.ecdsa !== undefined) {
    if (!isObj(e.ecdsa) || Object.keys(e.ecdsa).length !== 2 || e.ecdsa.addr === undefined || e.ecdsa.claim === undefined) fail(14, 'row ' + index + ' ecdsa must be {addr, claim}');
    ecdsa = { addr: hex(e.ecdsa.addr, 20, 14, 'ecdsa.addr'), claim: hex(e.ecdsa.claim, 65, 14, 'ecdsa.claim') };
  }
  return { index, id, keys, transport, rotation, ecdsa };
}

const keyOf = (row, name) => row.keys.find(k => k.name === name);
const sameKeys = (a, b) => a.keys.length === b.keys.length && a.keys.every(k => { const o = b.keys.find(x => x.wire === k.wire); return o && o.pk.equals(k.pk); });

export function load(json, chainId, requiredSchemes, previous, ecdsaRequired) {
  if (!isObj(json)) fail(3, 'the manifest is not a JSON object');
  if (json.formatVersion !== 3) fail(1, 'formatVersion must be the number 3');
  if (json.chainId !== chainId) fail(2, 'chainId must be the number ' + chainId);
  if (!Number.isInteger(json.bindHeight) || json.bindHeight < 0) fail(3, 'bindHeight must be a number >= 0');
  if (!Number.isInteger(json.count) || json.count < 1) fail(3, 'count must be a number >= 1');
  const { bindHeight, count } = json;
  const allowed = new Set(['formatVersion', 'chainId', 'bindHeight', 'count', 'registryHash']);
  for (let i = 0; i < count; i++) allowed.add(String(i));
  for (const k of Object.keys(json)) if (!allowed.has(k)) fail(3, 'unexpected top-level field ' + k);
  const rows = [];
  for (let i = 0; i < count; i++) { if (!isObj(json[String(i)])) fail(3, 'row ' + i + ' is missing'); rows.push(parseRow(i, json[String(i)])); }
  for (const row of rows) {
    for (const s of requiredSchemes) if (!keyOf(row, s)) fail(6, 'row ' + row.index + ' has no ' + s + ' key');
    if (!keyOf(row, 'falcon-512')) fail(6, 'row ' + row.index + ' has no falcon-512 key');
    for (const k of row.keys) {
      const m = popMessage(chainId, bindHeight, count, row.index, row.id, k.wire, k.pk);
      if (k.pop.length === 0 || !SCHEMES[k.name].verify(k.pk, m, k.pop)) fail(8, 'row ' + row.index + ' ' + k.name + ': the proof of possession does not verify');
    }
  }
  const ids = new Set(), keys = new Set();
  for (const row of rows) {
    const idh = row.id.toString('hex'); if (ids.has(idh)) fail(10, 'id appears on two rows'); ids.add(idh);
    for (const k of row.keys) { const h = k.pk.toString('hex'); if (keys.has(h)) fail(11, 'row ' + row.index + ' ' + k.name + ': the key appears on another row'); keys.add(h); }
  }
  const hashes = rows.map(r => rowHash(chainId, bindHeight, count, r));
  for (const row of rows) {
    const prev = previous ? previous.rows.find(p => p.id.equals(row.id)) : undefined;
    if (!prev) {
      if (!idOf(keyOf(row, 'falcon-512').pk).equals(row.id)) fail(9, 'row ' + row.index + ': a new id must be the derivation of its falcon key');
      if (Object.keys(row.rotation).length) fail(13, 'row ' + row.index + ' is a new id and carries a rotation');
      continue;
    }
    if (sameKeys(prev, row)) { if (Object.keys(row.rotation).length) fail(13, 'row ' + row.index + ' did not change its keys and carries a rotation'); continue; }
    const rot = rotateMessage(hashes[row.index]);
    for (const old of prev.keys) {
      const sig = row.rotation[old.name];
      if (!sig || !SCHEMES[old.name].verify(old.pk, rot, sig)) fail(12, 'row ' + row.index + ' changed its keys; the previous ' + old.name + ' key did not sign the new row');
    }
    if (Object.keys(row.rotation).length !== prev.keys.length) fail(13, 'row ' + row.index + ' carries a rotation signature for a scheme the previous row did not bind');
  }
  const ecd = new Set();
  for (const row of rows) {
    if (!row.ecdsa) { if (ecdsaRequired) fail(14, 'row ' + row.index + ' has no ecdsa link in the switch epoch'); continue; }
    const rec = recover(claimDigest(chainId, hashes[row.index]), row.ecdsa.claim);
    if (!rec || !rec.equals(row.ecdsa.addr)) fail(14, 'row ' + row.index + ': the ecdsa claim does not recover to its address');
    const a = row.ecdsa.addr.toString('hex'); if (ecd.has(a)) fail(14, 'ecdsa address on two rows'); ecd.add(a);
  }
  const root = merkleRoot(hashes);
  const hash = registryHash(chainId, bindHeight, count, root);
  if (json.registryHash !== undefined && !hex(json.registryHash, 32, 16, 'registryHash').equals(hash)) fail(16, 'the file declares another registry hash');
  return { chainId, bindHeight, count, rows, rowHashes: hashes, root, hash };
}

// ---- CLI: run every vector in a directory ----------------------------------------------------
if (process.argv[1] && path.resolve(process.argv[1]) === path.resolve(new URL(import.meta.url).pathname.replace(/^\/([A-Za-z]:)/, '$1'))) {
  const dir = process.argv[2];
  if (!dir) { console.error('usage: node registry-v3.mjs <vector dir>'); process.exit(2); }
  const files = fs.readdirSync(dir).filter(f => f.endsWith('.json')).sort();
  const cache = new Map();
  const read = (n) => JSON.parse(fs.readFileSync(path.join(dir, n + '.json'), 'utf8'));
  const run = (v) => load(v.manifest, v.chainId, v.requiredSchemes, v.previous ? (cache.get(v.previous) ?? (cache.set(v.previous, run(read(v.previous))), cache.get(v.previous))) : null, v.ecdsaRequired);
  let good = 0, bad = 0;
  for (const f of files) {
    const v = JSON.parse(fs.readFileSync(path.join(dir, f), 'utf8'));
    let got, r = null;
    try { r = run(v); got = 'ok'; } catch (e) { if (e instanceof RegistryV3Error) got = e.code; else { console.log(f + '  NEMASURAT: ' + e.message); process.exitCode = 2; continue; } }
    let same = got === v.expect;
    if (same && got === 'ok') {
      same = '0x' + r.hash.toString('hex') === v.registryHash && r.rowHashes.every((h, i) => '0x' + h.toString('hex') === v.rowHashes[i])
        && r.rows.every((row, i) => verifyMerkleProof(r.rowHashes[i], i, r.count, v.merkleProofs[String(i)].map(x => Buffer.from(x.slice(2), 'hex')), r.root));
    }
    console.log((same ? 'AS NAMED  ' : 'MISMATCH  ') + f.padEnd(34) + ' expected ' + v.expect + ', got ' + got);
    same ? good++ : bad++;
  }
  console.log(`\n${good} of ${files.length} vectors reach the outcome they name` + (bad ? `; ${bad} DO NOT` : ''));
  if (bad) process.exitCode = 1;
}