// Independent (second-language) reference for the FRI query-index derivation used by the PQ
// STARK-verify precompile 0x0AE8. Mirrors fri_query_index_reference.py and the Java FriQueryIndex
// helper. Its ONLY purpose is cross-language agreement: three independent implementations (Python,
// Node, Java) of the same public Plonky3 p3-fri index spec should produce byte-identical output on
// a shared vector set, which validates the index LOGIC. It verifies NOTHING about a real proof and
// the top-level 0x0AE8 verifier stays fail-closed. Conformance to a real SP1 trace is [MEASURE]
// (see docs/AERE-STARK-VERIFIER-PORT-SPEC.md section 8).
//
// Run standalone to emit its vector set as JSON on stdout:  node fri_query_index_reference.mjs

export const BABYBEAR_P = 2013265921; // 2^31 - 2^27 + 1

// sample_bits: low `bits` bits of a canonical BabyBear representative. [VERIFY] LSB vs MSB against
// p3-challenger; bits < 31 so a single 31-bit sample suffices for the pinned config.
export function sampleBits(canonicalU32, bits) {
  if (bits < 0) throw new Error("bits must be non-negative");
  if (bits >= 31) throw new Error("bits >= 31 needs multiple field samples");
  if (canonicalU32 < 0 || canonicalU32 >= BABYBEAR_P)
    throw new Error("canonicalU32 must be canonical in [0, p)");
  return canonicalU32 & ((1 << bits) - 1);
}

// Per-layer folding-index walk for one FRI query. Arity 2: sibling = index ^ 1,
// index_pair = index >> 1. num_fold_rounds = log_max_height - log_final_poly_len.
export function walk(index, logMaxHeight, logFinalPolyLen = 0) {
  if (logMaxHeight < 0 || logFinalPolyLen < 0) throw new Error("log heights must be non-negative");
  if (logFinalPolyLen > logMaxHeight) throw new Error("log_final_poly_len exceeds log_max_height");
  if (index < 0 || index >= (1 << logMaxHeight)) throw new Error("index out of range");
  const numFoldRounds = logMaxHeight - logFinalPolyLen;
  const steps = [];
  let cur = index;
  for (let layer = 0; layer < numFoldRounds; layer++) {
    steps.push({
      layer,
      index: cur,
      sibling: cur ^ 1,
      index_pair: cur >> 1,
      parity: cur & 1,
    });
    cur = cur >> 1;
  }
  return steps;
}

export function finalIndex(index, logMaxHeight, logFinalPolyLen = 0) {
  return index >> (logMaxHeight - logFinalPolyLen);
}

// A fixed shared vector set the three languages all evaluate, for cross-language agreement.
export function sharedVectors() {
  const walkCases = [
    { index: 11, logMaxHeight: 4, logFinalPolyLen: 0 },
    { index: 22, logMaxHeight: 5, logFinalPolyLen: 1 },
    { index: 0, logMaxHeight: 6, logFinalPolyLen: 0 },
    { index: 63, logMaxHeight: 6, logFinalPolyLen: 0 },
    { index: 12345, logMaxHeight: 20, logFinalPolyLen: 3 },
    { index: 1, logMaxHeight: 1, logFinalPolyLen: 0 },
  ];
  const sampleCases = [
    { v: 20, bits: 2 },
    { v: 20, bits: 3 },
    { v: 20, bits: 5 },
    { v: 0x6ae81234, bits: 5 },
    { v: 0x6ae81234, bits: 20 },
    { v: BABYBEAR_P - 1, bits: 10 },
  ];
  return {
    walks: walkCases.map((c) => ({
      ...c,
      steps: walk(c.index, c.logMaxHeight, c.logFinalPolyLen),
      final_index: finalIndex(c.index, c.logMaxHeight, c.logFinalPolyLen),
    })),
    samples: sampleCases.map((c) => ({ ...c, out: sampleBits(c.v, c.bits) })),
  };
}

// When run directly (not imported), print the shared vector set as canonical JSON. Uses
// pathToFileURL so the main-module check is correct on Windows and POSIX alike.
import { pathToFileURL } from "node:url";
if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) {
  process.stdout.write(JSON.stringify(sharedVectors()));
}