// Independent Node reference for the Fiat-Shamir DUPLEX CHALLENGER (component (f)) of the PQ
// STARK-verify precompile 0x0AE8. A genuinely different-language second implementation of
// DuplexChallenger<BabyBear, Perm, 16, 8> + GrindingChallenger check_witness over the CONFIRMED
// Poseidon2 permutation, used to cross-check the Python and Java references byte-for-byte. See
// docs/AERE-STARK-VERIFIER-PORT-SPEC.md section 7 and challenger_reference.py for the honest scope.
//
// Pinned source: p3-challenger 0.4.3-succinct duplex_challenger.rs / grinding_challenger.rs; the
// transcript observe/sample order is p3-fri verifier.rs verify_shape_and_sample_challenges.

import { readFileSync } from "node:fs";
import { fileURLToPath } from "node:url";
import { dirname, join } from "node:path";
import * as P2 from "./poseidon2_babybear_reference.mjs";

const HERE = dirname(fileURLToPath(import.meta.url));
const P = P2.P; // BigInt BabyBear prime
const WIDTH = 16;
const RATE = 8;

const KAT_DIGEST = [101, 102, 103, 104, 105, 106, 107, 108];
const KAT_SAMPLE_BITS = 10;
const KAT_POW_BITS = 4;

const mod = (a) => ((a % P) + P) % P;

class DuplexChallenger {
  constructor() {
    this.spongeState = new Array(WIDTH).fill(0n);
    this.inputBuffer = [];
    this.outputBuffer = [];
  }
  duplexing() {
    if (this.inputBuffer.length > RATE) throw new Error("input buffer > RATE");
    for (let i = 0; i < this.inputBuffer.length; i++) {
      this.spongeState[i] = mod(this.inputBuffer[i]);
    }
    this.inputBuffer = [];
    this.spongeState = P2.permute(this.spongeState);
    this.outputBuffer = this.spongeState.slice(0, RATE);
  }
  observe(value) {
    this.outputBuffer = [];
    this.inputBuffer.push(mod(BigInt(value)));
    if (this.inputBuffer.length === RATE) this.duplexing();
  }
  observeSlice(values) {
    for (const v of values) this.observe(v);
  }
  observeExt(ext) {
    this.observeSlice(ext);
  }
  sampleBase() {
    if (this.inputBuffer.length > 0 || this.outputBuffer.length === 0) this.duplexing();
    return this.outputBuffer.pop(); // Vec::pop -> last element
  }
  sampleExt() {
    return [this.sampleBase(), this.sampleBase(), this.sampleBase(), this.sampleBase()];
  }
  sampleBits(bits) {
    const randF = this.sampleBase();
    return Number(randF & ((1n << BigInt(bits)) - 1n));
  }
  checkWitness(bits, witness) {
    this.observe(witness);
    return this.sampleBits(bits) === 0;
  }
  clone() {
    const c = new DuplexChallenger();
    c.spongeState = this.spongeState.slice();
    c.inputBuffer = this.inputBuffer.slice();
    c.outputBuffer = this.outputBuffer.slice();
    return c;
  }
}

const n = (x) => Number(x); // canonical values fit in 2^31, safe as JS Number
const na = (arr) => arr.map(n);

function runDuplexScript() {
  const ch = new DuplexChallenger();
  ch.observe(11);
  ch.observe(22);
  const a = ch.sampleBase();
  const b = ch.sampleBase();
  const c = ch.sampleExt();
  ch.observe(33);
  ch.observe(44);
  ch.observe(55);
  const d = ch.sampleBits(KAT_SAMPLE_BITS);
  ch.observeSlice(KAT_DIGEST);
  const e = ch.sampleExt();
  const f = ch.sampleBase();
  const g = ch.sampleBase();
  const h = ch.sampleBase();
  const i = ch.sampleBase();
  const j = ch.sampleExt();
  const stateAfter = ch.spongeState.slice();
  return { ch, a, b, c, d, e, f, g, h, i, j, stateAfter };
}

function duplexWitnessTable(ch, witnesses) {
  return witnesses.map((w) => ({ witness: w, accept: ch.clone().checkWitness(KAT_POW_BITS, w) }));
}

function deriveFriChallenges(commits, finalPoly, powWitness, powBits, logBlowup, numQueries) {
  const ch = new DuplexChallenger();
  const betas = [];
  for (const comm of commits) {
    ch.observeSlice(comm);
    betas.push(ch.sampleExt());
  }
  ch.observeExt(finalPoly);
  const powAccept = ch.checkWitness(powBits, powWitness);
  const logMaxHeight = commits.length + logBlowup;
  const queryIndices = [];
  for (let q = 0; q < numQueries; q++) queryIndices.push(ch.sampleBits(logMaxHeight));
  return { betas, queryIndices, powAccept };
}

export function sharedVectors() {
  const gt = JSON.parse(readFileSync(join(HERE, "challenger_ground_truth.json"), "utf8"));
  const r = runDuplexScript();
  const witnesses = gt.duplex_kat.witness_table.map((w) => w.witness);
  const fri = gt.fri_transcript.map((tc) => {
    const { betas, queryIndices, powAccept } = deriveFriChallenges(
      tc.commit_phase_commits, tc.final_poly, tc.pow_witness, tc.pow_bits,
      tc.log_blowup, tc.num_queries);
    return { name: tc.name, betas: betas.map(na), queryIndices, powAccept };
  });
  return {
    permZeros: na(P2.permute(new Array(16).fill(0n))),
    duplex: {
      a: n(r.a), b: n(r.b), c: na(r.c), d: r.d, e: na(r.e),
      f: n(r.f), g: n(r.g), h: n(r.h), i: n(r.i), j: na(r.j),
      stateAfter: na(r.stateAfter),
      witnessTable: duplexWitnessTable(r.ch, witnesses),
    },
    fri,
  };
}

process.stdout.write(JSON.stringify(sharedVectors()));