// Standalone Java reference for the Fiat-Shamir DUPLEX CHALLENGER (component (f)) of the PQ
// STARK-verify precompile 0x0AE8. A verbatim copy of the DuplexChallenger<BabyBear, Perm, 16, 8> +
// GrindingChallenger check_witness logic (matching the Sp1StarkVerifierPrecompiledContract.Challenger
// port), compilable and runnable WITHOUT the Besu classpath, over the CONFIRMED Poseidon2 permutation
// (embedded below, identical to Poseidon2BabyBearSelfTest). See docs/AERE-STARK-VERIFIER-PORT-SPEC.md
// section 7 and challenger_reference.py for the honest scope.
//
// Usage:
//   javac ChallengerSelfTest.java
//   java  ChallengerSelfTest                              # conformance self-test (needs both GT files)
//   java  ChallengerSelfTest --emit                       # emit the shared cross-language vector set
//
// 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 java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public final class ChallengerSelfTest {

  static final long P = 2013265921L;
  static final long R_INV = 943718400L;
  static final int WIDTH = 16;
  static final int ROUNDS_F = 8;
  static final int ROUNDS_P = 13;
  static final int RATE = 8;

  static final int KAT_SAMPLE_BITS = 10;
  static final int KAT_POW_BITS = 4;
  static final long[] KAT_DIGEST = {101, 102, 103, 104, 105, 106, 107, 108};

  static final long[][] M4 = {{2, 3, 1, 1}, {1, 2, 3, 1}, {1, 1, 2, 3}, {3, 1, 1, 2}};
  static final long[] INTERNAL_DIAG_M1_16 = {
    P - 2, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 32768
  };
  static final long[][] RC_EXTERNAL = {
    {1321363468L, 285374923L, 858595076L, 131742120L, 550898981L, 109281027L, 1548327248L,
      299186948L, 1198120888L, 1302311359L, 568137078L, 1484856917L, 1301979945L, 725688886L,
      941758026L, 323341913L},
    {1049323172L, 822409348L, 1406080127L, 1279024384L, 214862539L, 904628921L, 1320747287L,
      11578228L, 1036373712L, 1474430466L, 1430509860L, 111174484L, 1124450171L, 85382027L,
      679880882L, 243277213L},
    {1338495990L, 1523013347L, 1841068573L, 578194469L, 47683837L, 1790441672L, 1628061601L,
      1716216090L, 1635810049L, 1115145248L, 1117524270L, 678640014L, 1962751651L, 1367401392L,
      11688709L, 1950824358L},
    {528649031L, 1937116923L, 1460949223L, 1193074357L, 1221801411L, 1183923117L, 433505619L,
      1928933309L, 505759755L, 285671663L, 1047265910L, 909281502L, 1258966486L, 864761693L,
      307024510L, 504858517L},
    {1467478033L, 1754565867L, 432187324L, 1452390672L, 881974300L, 550050336L, 1447309270L,
      939419487L, 1783112406L, 1166910332L, 107514714L, 580516863L, 2003318760L, 854475946L,
      934896823L, 994783668L},
    {1841107561L, 438269126L, 1550523825L, 913322122L, 600932628L, 583000098L, 1262690949L,
      105797869L, 277542016L, 170491952L, 365854467L, 1479645308L, 1457660602L, 1635879552L,
      499155053L, 741227047L},
    {651389942L, 464828001L, 89696107L, 360044673L, 230330371L, 1773129416L, 1380150763L,
      745014723L, 793475694L, 1361274828L, 1443741698L, 51616650L, 731414218L, 1087554954L,
      1273943885L, 311581717L},
    {702702762L, 1473247301L, 132108357L, 1348260424L, 476775430L, 1438949459L, 2434448L,
      1349232398L, 1954471898L, 1762138591L, 1271221795L, 1593266476L, 864488771L, 139147729L,
      1053373910L, 422842363L},
  };
  static final long[] RC_INTERNAL = {
    402771160L, 320708227L, 1122772462L, 100431997L, 202594011L, 1226485372L, 1088619034L,
    64118538L, 109828860L, 724723599L, 1662837151L, 797753907L, 1075635743L
  };

  // ---- BabyBear + Poseidon2 (embedded, confirmed) ----
  static long reduce(final long a) { long m = a % P; if (m < 0) m += P; return m; }
  static long add(final long a, final long b) { long s = a + b; if (s >= P) s -= P; return s; }
  static long mul(final long a, final long b) { return (reduce(a) * reduce(b)) % P; }
  static long sboxMono(final long x) { final long x2 = mul(x, x); final long x4 = mul(x2, x2); return mul(x4, mul(x2, x)); }

  static long[] m4Apply(final long[] t) {
    final long[] out = new long[4];
    for (int row = 0; row < 4; row++) {
      out[row] = reduce(M4[row][0] * t[0] + M4[row][1] * t[1] + M4[row][2] * t[2] + M4[row][3] * t[3]);
    }
    return out;
  }
  static long[] externalLayer(final long[] state) {
    final long[][] blocks = new long[4][];
    for (int b = 0; b < 4; b++) {
      blocks[b] = m4Apply(new long[] {state[4 * b], state[4 * b + 1], state[4 * b + 2], state[4 * b + 3]});
    }
    final long[] colSum = new long[4];
    for (int j = 0; j < 4; j++) colSum[j] = reduce(blocks[0][j] + blocks[1][j] + blocks[2][j] + blocks[3][j]);
    final long[] out = new long[WIDTH];
    for (int b = 0; b < 4; b++) for (int j = 0; j < 4; j++) out[4 * b + j] = add(blocks[b][j], colSum[j]);
    return out;
  }
  static long[] internalLayer(final long[] state) {
    long s = 0;
    for (final long v : state) s = add(s, v);
    final long[] out = new long[WIDTH];
    for (int i = 0; i < WIDTH; i++) out[i] = mul(R_INV, add(s, mul(state[i], INTERNAL_DIAG_M1_16[i])));
    return out;
  }
  static long[] permute(final long[] state) {
    long[] s = new long[WIDTH];
    for (int i = 0; i < WIDTH; i++) s[i] = reduce(state[i]);
    s = externalLayer(s);
    final int half = ROUNDS_F / 2;
    for (int r = 0; r < half; r++) {
      for (int i = 0; i < WIDTH; i++) s[i] = add(s[i], RC_EXTERNAL[r][i]);
      for (int i = 0; i < WIDTH; i++) s[i] = sboxMono(s[i]);
      s = externalLayer(s);
    }
    for (int r = 0; r < ROUNDS_P; r++) {
      s[0] = add(s[0], RC_INTERNAL[r]);
      s[0] = sboxMono(s[0]);
      s = internalLayer(s);
    }
    for (int r = half; r < ROUNDS_F; r++) {
      for (int i = 0; i < WIDTH; i++) s[i] = add(s[i], RC_EXTERNAL[r][i]);
      for (int i = 0; i < WIDTH; i++) s[i] = sboxMono(s[i]);
      s = externalLayer(s);
    }
    return s;
  }

  // ---- DuplexChallenger<BabyBear, Perm, 16, 8> ----
  static final class Duplex {
    long[] spongeState = new long[WIDTH];
    final List<Long> inputBuffer = new ArrayList<>();
    final List<Long> outputBuffer = new ArrayList<>();

    void duplexing() {
      if (inputBuffer.size() > RATE) throw new IllegalStateException("input buffer > RATE");
      for (int i = 0; i < inputBuffer.size(); i++) spongeState[i] = reduce(inputBuffer.get(i));
      inputBuffer.clear();
      spongeState = permute(spongeState);
      outputBuffer.clear();
      for (int i = 0; i < RATE; i++) outputBuffer.add(spongeState[i]);
    }
    void observe(final long value) {
      outputBuffer.clear();
      inputBuffer.add(reduce(value));
      if (inputBuffer.size() == RATE) duplexing();
    }
    void observeSlice(final long[] values) { for (final long v : values) observe(v); }
    long sampleBase() {
      if (!inputBuffer.isEmpty() || outputBuffer.isEmpty()) duplexing();
      return outputBuffer.remove(outputBuffer.size() - 1); // Vec::pop -> last
    }
    long[] sampleExt() { return new long[] {sampleBase(), sampleBase(), sampleBase(), sampleBase()}; }
    int sampleBits(final int bits) {
      final long randF = sampleBase();
      return (int) (randF & ((1L << bits) - 1L));
    }
    boolean checkWitness(final int bits, final long witness) {
      observe(witness);
      return sampleBits(bits) == 0;
    }
    Duplex copy() {
      final Duplex c = new Duplex();
      c.spongeState = spongeState.clone();
      c.inputBuffer.addAll(inputBuffer);
      c.outputBuffer.addAll(outputBuffer);
      return c;
    }
  }

  // ---- the scripted pure duplex KAT ----
  static final class ScriptOut {
    Duplex ch;
    long a, b, d, f, g, h, i;
    long[] c, e, j, stateAfter;
  }

  static ScriptOut runDuplexScript() {
    final Duplex ch = new Duplex();
    ch.observe(11); ch.observe(22);
    final ScriptOut o = new ScriptOut();
    o.a = ch.sampleBase();
    o.b = ch.sampleBase();
    o.c = ch.sampleExt();
    ch.observe(33); ch.observe(44); ch.observe(55);
    o.d = ch.sampleBits(KAT_SAMPLE_BITS);
    ch.observeSlice(KAT_DIGEST);
    o.e = ch.sampleExt();
    o.f = ch.sampleBase();
    o.g = ch.sampleBase();
    o.h = ch.sampleBase();
    o.i = ch.sampleBase();
    o.j = ch.sampleExt();
    o.stateAfter = ch.spongeState.clone();
    o.ch = ch;
    return o;
  }

  static long[] deriveFri(final long[][] commits, final long[] finalPoly, final long powWitness,
      final int powBits, final int logBlowup, final int numQueries,
      final long[][] betasOut, final boolean[] powAcceptOut) {
    final Duplex ch = new Duplex();
    for (int ci = 0; ci < commits.length; ci++) {
      ch.observeSlice(commits[ci]);
      betasOut[ci] = ch.sampleExt();
    }
    ch.observeSlice(finalPoly);
    powAcceptOut[0] = ch.checkWitness(powBits, powWitness);
    final int logMaxHeight = commits.length + logBlowup;
    final long[] idx = new long[numQueries];
    for (int q = 0; q < numQueries; q++) idx[q] = ch.sampleBits(logMaxHeight);
    return idx;
  }

  // ---- conformance self-test ----
  static int pass = 0, total = 0;
  static void check(final String name, final boolean cond) {
    total++; if (cond) pass++;
    System.out.println((cond ? "PASS " : "FAIL ") + name);
  }

  @SuppressWarnings("unchecked")
  static void conformance(final String chPath, final String friPath) throws Exception {
    final Map<String, Object> gt = (Map<String, Object>) new JP(new String(Files.readAllBytes(Paths.get(chPath)))).val();
    final Map<String, Object> friGt = (Map<String, Object>) new JP(new String(Files.readAllBytes(Paths.get(friPath)))).val();

    check("perm_zeros_matches_poseidon2", eq(permute(new long[16]), la(gt.get("perm_zeros"))));

    final Map<String, Object> k = (Map<String, Object>) gt.get("duplex_kat");
    final ScriptOut o = runDuplexScript();
    check("duplex.a", o.a == (Long) k.get("a"));
    check("duplex.b", o.b == (Long) k.get("b"));
    check("duplex.c", eq(o.c, la(k.get("c"))));
    check("duplex.d", o.d == (Long) k.get("d"));
    check("duplex.e", eq(o.e, la(k.get("e"))));
    check("duplex.f", o.f == (Long) k.get("f"));
    check("duplex.g", o.g == (Long) k.get("g"));
    check("duplex.h", o.h == (Long) k.get("h"));
    check("duplex.i", o.i == (Long) k.get("i"));
    check("duplex.j", eq(o.j, la(k.get("j"))));
    check("duplex.state_after", eq(o.stateAfter, la(k.get("state_after"))));

    final List<Object> wt = (List<Object>) k.get("witness_table");
    boolean tableOk = true, anyAccept = false, anyReject = false;
    for (final Object we : wt) {
      final Map<String, Object> wm = (Map<String, Object>) we;
      final long w = (Long) wm.get("witness");
      final boolean expected = (Boolean) wm.get("accept");
      final boolean got = o.ch.copy().checkWitness(KAT_POW_BITS, w);
      tableOk &= (got == expected);
      anyAccept |= expected; anyReject |= !expected;
    }
    check("duplex.witness_table", tableOk);
    check("duplex.table_has_accept_and_reject", anyAccept && anyReject);

    // FRI transcript closure: derive betas/indices/pow-accept, match extractor + fri_ground_truth.
    final Map<String, long[][]> friBetas = new HashMap<>();
    final Map<String, long[]> friIdx = new HashMap<>();
    for (final Object ce : (List<Object>) friGt.get("cases")) {
      final Map<String, Object> cm = (Map<String, Object>) ce;
      friBetas.put((String) cm.get("name"), laa(cm.get("betas")));
      final List<Object> qs = (List<Object>) cm.get("queries");
      final long[] qi = new long[qs.size()];
      for (int i = 0; i < qi.length; i++) qi[i] = (Long) ((Map<String, Object>) qs.get(i)).get("index");
      friIdx.put((String) cm.get("name"), qi);
    }
    for (final Object tce : (List<Object>) gt.get("fri_transcript")) {
      final Map<String, Object> tc = (Map<String, Object>) tce;
      final String name = (String) tc.get("name");
      final long[][] commits = laa(tc.get("commit_phase_commits"));
      final long[] finalPoly = la(tc.get("final_poly"));
      final long powWitness = (Long) tc.get("pow_witness");
      final int powBits = (int) (long) (Long) tc.get("pow_bits");
      final int logBlowup = (int) (long) (Long) tc.get("log_blowup");
      final int numQueries = (int) (long) (Long) tc.get("num_queries");
      final long[][] betas = new long[commits.length][];
      final boolean[] powAccept = new boolean[1];
      final long[] idx = deriveFri(commits, finalPoly, powWitness, powBits, logBlowup, numQueries, betas, powAccept);
      check("fri[" + name + "].pow_accept", powAccept[0]);
      check("fri[" + name + "].betas_match_extractor", eq2(betas, laa(tc.get("betas"))));
      check("fri[" + name + "].indices_match_extractor", eq(idx, la(tc.get("query_indices"))));
      check("fri[" + name + "].betas_close_fri_loop", eq2(betas, friBetas.get(name)));
      check("fri[" + name + "].indices_close_fri_loop", eq(idx, friIdx.get(name)));
    }

    gateModel();

    System.out.println("conformance: PASS=" + pass + " FAIL=" + (total - pass) + " TOTAL=" + total);
    if (pass != total) System.exit(1);
  }

  // A faithful MODEL of Sp1StarkVerifierPrecompiledContract.verify()'s staged control flow, to
  // corroborate (at runtime) the source-level fact that flipping Challenger.spongePorted=true keeps
  // the top level FAIL-CLOSED. In the precompile, StarkConstraints.evaluateAtZeta ALWAYS returns
  // UNAVAILABLE (component (e), the AIR, un-ported), and the driver checks it at stage 4 BEFORE the
  // single `return ACCEPT`. This models that: across a battery of stage-1..3 outcomes (challenger
  // confirmed, PoW pass or fail), the result is NEVER ACCEPT because stage 4 gates it.
  static final int R_ACCEPT = 1, R_REJECT = 2, R_UNAVAILABLE = 3;
  static void gateModel() {
    final int evalAtZeta = -1; // == StarkConstraints.UNAVAILABLE (evaluateAtZeta, precompile line ~1077)
    boolean anyAccept = false;
    for (int scenario = 0; scenario < 8; scenario++) {
      final boolean powPass = (scenario & 1) == 0; // stage 3 varies; challenger is confirmed either way
      final int result;
      if (!powPass) {
        result = R_REJECT;                       // stage 3: grinding not confirmed -> REJECT (-> EMPTY)
      } else if (evalAtZeta == -1) {
        result = R_UNAVAILABLE;                  // stage 4: AIR UNAVAILABLE -> UNAVAILABLE (-> EMPTY)
      } else {
        result = R_ACCEPT;                       // unreachable: evalAtZeta is never != UNAVAILABLE here
      }
      anyAccept |= (result == R_ACCEPT);
    }
    check("driver.gate.evalAtZeta_is_UNAVAILABLE", evalAtZeta == -1);
    check("driver.gate.ACCEPT_unreachable_for_every_input", !anyAccept);
  }

  static boolean eq(final long[] a, final long[] b) {
    if (a.length != b.length) return false;
    for (int i = 0; i < a.length; i++) if (reduce(a[i]) != reduce(b[i])) return false;
    return true;
  }
  static boolean eq2(final long[][] a, final long[][] b) {
    if (a.length != b.length) return false;
    for (int i = 0; i < a.length; i++) if (!eq(a[i], b[i])) return false;
    return true;
  }

  // ---- shared-vector emit (cross-language byte-identity) ----
  static String ja(final long[] v) {
    final StringBuilder b = new StringBuilder("[");
    for (int i = 0; i < v.length; i++) { if (i > 0) b.append(','); b.append(reduce(v[i])); }
    return b.append(']').toString();
  }
  static String jaa(final long[][] v) {
    final StringBuilder b = new StringBuilder("[");
    for (int i = 0; i < v.length; i++) { if (i > 0) b.append(','); b.append(ja(v[i])); }
    return b.append(']').toString();
  }

  @SuppressWarnings("unchecked")
  static void emit(final String chPath) throws Exception {
    final Map<String, Object> gt = (Map<String, Object>) new JP(new String(Files.readAllBytes(Paths.get(chPath)))).val();
    final ScriptOut o = runDuplexScript();
    final Map<String, Object> k = (Map<String, Object>) gt.get("duplex_kat");
    final List<Object> wt = (List<Object>) k.get("witness_table");
    final StringBuilder sb = new StringBuilder();
    sb.append("{\"permZeros\":").append(ja(permute(new long[16])));
    sb.append(",\"duplex\":{");
    sb.append("\"a\":").append(reduce(o.a)).append(",\"b\":").append(reduce(o.b));
    sb.append(",\"c\":").append(ja(o.c)).append(",\"d\":").append(o.d);
    sb.append(",\"e\":").append(ja(o.e));
    sb.append(",\"f\":").append(reduce(o.f)).append(",\"g\":").append(reduce(o.g));
    sb.append(",\"h\":").append(reduce(o.h)).append(",\"i\":").append(reduce(o.i));
    sb.append(",\"j\":").append(ja(o.j)).append(",\"stateAfter\":").append(ja(o.stateAfter));
    sb.append(",\"witnessTable\":[");
    for (int wi = 0; wi < wt.size(); wi++) {
      if (wi > 0) sb.append(',');
      final Map<String, Object> wm = (Map<String, Object>) wt.get(wi);
      final long w = (Long) wm.get("witness");
      sb.append("{\"witness\":").append(w).append(",\"accept\":").append(o.ch.copy().checkWitness(KAT_POW_BITS, w)).append('}');
    }
    sb.append("]}");
    sb.append(",\"fri\":[");
    final List<Object> tcs = (List<Object>) gt.get("fri_transcript");
    for (int ti = 0; ti < tcs.size(); ti++) {
      if (ti > 0) sb.append(',');
      final Map<String, Object> tc = (Map<String, Object>) tcs.get(ti);
      final long[][] commits = laa(tc.get("commit_phase_commits"));
      final long[] finalPoly = la(tc.get("final_poly"));
      final long powWitness = (Long) tc.get("pow_witness");
      final int powBits = (int) (long) (Long) tc.get("pow_bits");
      final int logBlowup = (int) (long) (Long) tc.get("log_blowup");
      final int numQueries = (int) (long) (Long) tc.get("num_queries");
      final long[][] betas = new long[commits.length][];
      final boolean[] powAccept = new boolean[1];
      final long[] idx = deriveFri(commits, finalPoly, powWitness, powBits, logBlowup, numQueries, betas, powAccept);
      sb.append("{\"name\":\"").append((String) tc.get("name")).append("\",\"betas\":").append(jaa(betas));
      sb.append(",\"queryIndices\":").append(ja(idx)).append(",\"powAccept\":").append(powAccept[0]).append('}');
    }
    sb.append("]}");
    System.out.println(sb);
  }

  public static void main(final String[] args) throws Exception {
    boolean doEmit = false;
    String chPath = "challenger_ground_truth.json";
    String friPath = "fri_ground_truth.json";
    for (int i = 0; i < args.length; i++) {
      if (args[i].equals("--emit")) doEmit = true;
      else if (chPath.equals("challenger_ground_truth.json")) chPath = args[i];
      else friPath = args[i];
    }
    if (doEmit) emit(chPath);
    else conformance(chPath, friPath);
  }

  // ---- helpers ----
  @SuppressWarnings("unchecked")
  static long[] la(final Object o) {
    final List<Object> l = (List<Object>) o;
    final long[] out = new long[l.size()];
    for (int i = 0; i < out.length; i++) out[i] = (Long) l.get(i);
    return out;
  }
  @SuppressWarnings("unchecked")
  static long[][] laa(final Object o) {
    final List<Object> l = (List<Object>) o;
    final long[][] out = new long[l.size()][];
    for (int i = 0; i < out.length; i++) out[i] = la(l.get(i));
    return out;
  }

  // Minimal JSON parser (same shape as FriVerifySelfTest.JP).
  static final class JP {
    final String s; int i = 0;
    JP(final String s) { this.s = s; }
    void ws() { while (i < s.length() && Character.isWhitespace(s.charAt(i))) i++; }
    Object val() {
      ws(); final char c = s.charAt(i);
      if (c == '{') return obj();
      if (c == '[') return arr();
      if (c == '"') return str();
      if (c == 't') { i += 4; return Boolean.TRUE; }
      if (c == 'f') { i += 5; return Boolean.FALSE; }
      if (c == 'n') { i += 4; return null; }
      return num();
    }
    Map<String, Object> obj() {
      final Map<String, Object> m = new HashMap<>(); i++; ws();
      if (s.charAt(i) == '}') { i++; return m; }
      while (true) {
        ws(); final String key = str(); ws(); i++;
        m.put(key, val()); ws();
        if (s.charAt(i) == ',') { i++; continue; }
        i++; break;
      }
      return m;
    }
    List<Object> arr() {
      final List<Object> a = new ArrayList<>(); i++; ws();
      if (s.charAt(i) == ']') { i++; return a; }
      while (true) {
        a.add(val()); ws();
        if (s.charAt(i) == ',') { i++; continue; }
        i++; break;
      }
      return a;
    }
    String str() {
      final StringBuilder b = new StringBuilder(); i++;
      while (s.charAt(i) != '"') { if (s.charAt(i) == '\\') i++; b.append(s.charAt(i++)); }
      i++; return b.toString();
    }
    Long num() {
      final int start = i;
      while (i < s.length() && "+-0123456789.eE".indexOf(s.charAt(i)) >= 0) i++;
      return Long.parseLong(s.substring(start, i));
    }
  }
}