// Independent Node reference for the GENERIC AIR constraint / quotient-consistency check (component
// (e)) that a p3-uni-stark STARK verifier performs, for the PQ STARK-verify precompile 0x0AE8. Mirrors
// air_quotient_reference.py and AirQuotientSelfTest.java. See docs/AERE-STARK-VERIFIER-PORT-SPEC.md
// section 6, spec-stark-air.md.
//
// HONEST SCOPE: this implements the GENERIC quotient-consistency MECHANISM (p3-uni-stark verifier.rs
// lines 90-141) for a SUPPLIED AIR constraint evaluator: fold the AIR's constraints with alpha (Horner),
// reconstruct quotient(zeta) from the chunk openings + split-domain zps, compute Z_H(zeta) and the
// first/last/transition selectors, and check folded_constraints(zeta) == Z_H(zeta) * quotient(zeta),
// over F_{p^4}. The SP1 recursion AIR (its specific constraint set + vkey binding) is NOT ported, so the
// top-level 0x0AE8 STAYS FAIL-CLOSED for real SP1 proofs. Ground truth: real p3-uni-stark 0.4.3-succinct
// proofs of two KNOWN example AIRs (Fibonacci from tests/fib_air.rs; a degree-3 multiply AIR matching
// tests/mul_air.rs), accepted by the library and re-checked here (accept + reject tampered).

import { readFileSync } from "fs";
import { fileURLToPath } from "url";
import { dirname, join } from "path";
import {
  P,
  GENERATOR,
  mul,
  inv,
  powF,
  twoAdicGenerator,
  extFromBase,
  extAdd,
  extSub,
  extMul,
  extInv,
} from "./babybear_field_reference.mjs";

const HERE = dirname(fileURLToPath(import.meta.url));
const GROUND_TRUTH = join(HERE, "air_quotient_ground_truth.json");
const ONE = extFromBase(1n);

const B = (x) => BigInt(x);
const extEq = (a, b) => a.every((_, i) => ((a[i] % P) + P) % P === ((b[i] % P) + P) % P);
const extDiv = (a, b) => extMul(a, extInv(b));

function extExpPow2(x, logN) {
  let r = x.slice();
  for (let k = 0; k < logN; k++) r = extMul(r, r);
  return r;
}

// Z_D(point) for a coset of size 2^logN and (base-field) shift: (point*shift^-1)^(2^logN) - 1.
function zpAtPoint(logN, shiftBase, pointExt) {
  const shiftInv = extFromBase(inv(shiftBase));
  return extSub(extExpPow2(extMul(pointExt, shiftInv), logN), ONE);
}

// domain.rs selectors_at_point for the trace domain (shift=1, logN=degreeBits).
function selectorsAtPoint(degreeBits, zeta) {
  const g = twoAdicGenerator(B(degreeBits));
  const gInv = extFromBase(inv(g));
  const unshifted = zeta.slice();
  const zH = extSub(extExpPow2(unshifted, degreeBits), ONE);
  return {
    isFirst: extDiv(zH, extSub(unshifted, ONE)),
    isLast: extDiv(zH, extSub(unshifted, gInv)),
    isTransition: extSub(unshifted, gInv),
    invZeroifier: extInv(zH),
  };
}

const monomial = (e) => {
  const m = [0n, 0n, 0n, 0n];
  m[e] = 1n;
  return m;
};

function reconstructQuotient(degreeBits, chunks, zeta) {
  const quotientDegree = chunks.length;
  const logQuotientDegree = Math.round(Math.log2(quotientDegree));
  const logNq = degreeBits + logQuotientDegree;
  const genQ = twoAdicGenerator(B(logNq));
  const shifts = [];
  for (let i = 0; i < quotientDegree; i++) shifts.push(mul(GENERATOR, powF(genQ, B(i))));

  const zps = [];
  for (let i = 0; i < quotientDegree; i++) {
    let acc = ONE.slice();
    const firstPointI = extFromBase(shifts[i]);
    for (let j = 0; j < quotientDegree; j++) {
      if (j === i) continue;
      const num = zpAtPoint(degreeBits, shifts[j], zeta);
      const den = zpAtPoint(degreeBits, shifts[j], firstPointI);
      acc = extMul(acc, extMul(num, extInv(den)));
    }
    zps.push(acc);
  }

  let quotient = [0n, 0n, 0n, 0n];
  for (let chI = 0; chI < chunks.length; chI++) {
    for (let eI = 0; eI < chunks[chI].length; eI++) {
      quotient = extAdd(quotient, extMul(extMul(zps[chI], monomial(eI)), chunks[chI][eI]));
    }
  }
  return quotient;
}

// ---- example AIR constraint evaluators (transcribed from their eval()) ----
function horner(constraints, alpha) {
  let acc = [0n, 0n, 0n, 0n];
  for (const c of constraints) acc = extAdd(extMul(acc, alpha), c);
  return acc;
}

function foldFibonacci(tl, tn, pis, sel, alpha) {
  const a = extFromBase(pis[0]);
  const b = extFromBase(pis[1]);
  const x = extFromBase(pis[2]);
  const [left, right] = [tl[0], tl[1]];
  const [nleft, nright] = [tn[0], tn[1]];
  const c1 = extMul(sel.isFirst, extSub(left, a));
  const c2 = extMul(sel.isFirst, extSub(right, b));
  const c3 = extMul(sel.isTransition, extSub(right, nleft));
  const c4 = extMul(sel.isTransition, extSub(extAdd(left, right), nright));
  const c5 = extMul(sel.isLast, extSub(right, x));
  return horner([c1, c2, c3, c4, c5], alpha);
}

function foldMulDeg3(tl, tn, pis, sel, alpha) {
  const [a, b, c] = [tl[0], tl[1], tl[2]];
  const nextA = tn[0];
  const c1 = extSub(extMul(extMul(a, a), b), c);
  const c2 = extMul(sel.isFirst, extSub(extAdd(extMul(a, a), ONE), b));
  const c3 = extMul(sel.isTransition, extSub(extAdd(a, ONE), nextA));
  return horner([c1, c2, c3], alpha);
}

const AIR_FOLDERS = { fibonacci: foldFibonacci, mul_deg3: foldMulDeg3 };

// ---- the generic quotient-consistency check ----
function toBigCase(cas) {
  return {
    name: cas.name,
    air: cas.air,
    degreeBits: cas.degree_bits,
    quotientDegree: cas.quotient_degree,
    alpha: cas.alpha.map(B),
    zeta: cas.zeta.map(B),
    pis: cas.public_values.map(B),
    tl: cas.trace_local.map((v) => v.map(B)),
    tn: cas.trace_next.map((v) => v.map(B)),
    chunks: cas.quotient_chunks.map((ch) => ch.map((v) => v.map(B))),
    libraryAccept: cas.library_accept,
  };
}

function derive(c) {
  const sel = selectorsAtPoint(c.degreeBits, c.zeta);
  const quotient = reconstructQuotient(c.degreeBits, c.chunks, c.zeta);
  const folded = AIR_FOLDERS[c.air](c.tl, c.tn, c.pis, sel, c.alpha);
  return { quotient, folded, invZeroifier: sel.invZeroifier };
}

function checkIdentity(c) {
  const d = derive(c);
  return extEq(extMul(d.folded, d.invZeroifier), d.quotient);
}

function checkCase(c, tamper) {
  const cc = {
    ...c,
    alpha: c.alpha.slice(),
    tl: c.tl.map((v) => v.slice()),
    chunks: c.chunks.map((ch) => ch.map((v) => v.slice())),
  };
  if (tamper === "trace") cc.tl[0][0] = (cc.tl[0][0] + 1n) % P;
  else if (tamper === "quotient") cc.chunks[0][0][0] = (cc.chunks[0][0][0] + 1n) % P;
  else if (tamper === "alpha") cc.alpha[0] = (cc.alpha[0] + 1n) % P;
  return checkIdentity(cc);
}

const asNum = (arr) => arr.map((x) => Number(((x % P) + P) % P));

export function sharedVectors() {
  const gt = JSON.parse(readFileSync(GROUND_TRUTH, "utf8"));
  const cases = gt.cases.map((raw) => {
    const c = toBigCase(raw);
    const d = derive(c);
    return {
      name: c.name,
      air: c.air,
      degreeBits: c.degreeBits,
      quotientDegree: c.quotientDegree,
      quotient: asNum(d.quotient),
      folded: asNum(d.folded),
      invZeroifier: asNum(d.invZeroifier),
      accept: checkCase(c, null),
      rejectTrace: !checkCase(c, "trace"),
      rejectQuotient: !checkCase(c, "quotient"),
      rejectAlpha: !checkCase(c, "alpha"),
    };
  });
  return { valGenerator: Number(GENERATOR), cases };
}

if (import.meta.url === `file://${process.argv[1]}` || process.argv[1] === fileURLToPath(import.meta.url)) {
  process.stdout.write(JSON.stringify(sharedVectors()));
}