#!/usr/bin/env python3
# Independent reference for the Fiat-Shamir DUPLEX CHALLENGER (component (f)) of the PQ STARK-verify
# precompile 0x0AE8: the Poseidon2 duplex sponge over BabyBear that SP1/Plonky3 use to derive the FRI
# folding challenges (betas), the query indices, and the grinding proof-of-work acceptance from the
# proof data. See docs/AERE-STARK-VERIFIER-PORT-SPEC.md section 7.
#
# ============================ HONEST SCOPE (read first) ============================
# This module implements DuplexChallenger<BabyBear, Perm, WIDTH=16, RATE=8> and its GrindingChallenger
# check_witness EXACTLY as the pinned p3-challenger 0.4.3-succinct does, over the CONFIRMED Poseidon2
# permutation (component (b)). It is CONFORMANCE-CONFIRMED against ground truth emitted by executing the
# pinned crates:
#   - a fully-scripted observe/sample transcript ("duplex_kat"): observe base elements + a digest,
#     sample base + F_{p^4} ext elements, sample_bits, and a check_witness accept/reject table; this
#     reference reproduces every sampled value, the full 16-lane sponge state, and every accept flag.
#   - the FRI transcript closure ("fri_transcript"): for the SAME three FRI cases the confirmed FRI
#     ground truth (component (d)) uses, this reference derives the betas and query indices FROM THE
#     TRANSCRIPT (commit_phase_commits + final_poly + pow_witness) and reproduces them byte-for-byte,
#     matching BOTH the challenger extractor AND fri_ground_truth.json. That closes the loop the FRI KAT
#     left open (transcript -> betas/indices -> the CONFIRMED FRI verify_query), so the betas/indices are
#     now DERIVED in-circuit, not taken as inputs.
#
# What this does NOT do: it does not verify any STARK proof and does not make the top-level 0x0AE8
# accept anything. Component (f) being confirmed is necessary but not sufficient: the SP1 recursion-AIR
# constraint evaluation (component (e)) is un-ported (StarkConstraints = UNAVAILABLE), so the precompile
# stays FAIL-CLOSED (returns EMPTY for every input). See the README and the port spec.
#
# Source (pinned, checksum-matched to the repo Cargo.lock):
#   p3-challenger 0.4.3-succinct  duplex_challenger.rs (DuplexChallenger observe/duplexing/sample/
#                                 sample_bits) + grinding_challenger.rs (check_witness).
#   The transcript observe/sample ORDER is p3-fri verifier.rs verify_shape_and_sample_challenges.
#   Sponge permutation: the CONFIRMED Poseidon2-BabyBear width-16 permutation (component (b)).

import json
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import poseidon2_babybear_reference as P2  # component (b): the CONFIRMED Poseidon2 permutation

P = P2.P
WIDTH = 16   # DuplexChallenger WIDTH (sponge state lanes). CONFIRMED from source.
RATE = 8     # DuplexChallenger RATE (absorb width + squeeze width). CONFIRMED from source.

CHALLENGER_GT = os.path.join(HERE, "challenger_ground_truth.json")
FRI_GT = os.path.join(HERE, "fri_ground_truth.json")

# Pure-KAT script constants (mirrored EXACTLY in the Rust extractor and the Node/Java references).
KAT_DIGEST = [101, 102, 103, 104, 105, 106, 107, 108]
KAT_SAMPLE_BITS = 10
KAT_POW_BITS = 4


class DuplexChallenger:
    """Port of p3-challenger DuplexChallenger<F, Perm, 16, 8>. Faithful to duplex_challenger.rs:
    overwrite-mode absorb of up to RATE inputs, permute, squeeze the first RATE lanes; sample pops
    from the END of the output buffer (Vec::pop)."""

    def __init__(self):
        self.sponge_state = [0] * WIDTH   # F::default() == 0
        self.input_buffer = []            # Vec<F>
        self.output_buffer = []           # Vec<F>

    def duplexing(self):
        assert len(self.input_buffer) <= RATE
        # Overwrite the first len(input_buffer) lanes with the buffered inputs (drain).
        for i, val in enumerate(self.input_buffer):
            self.sponge_state[i] = val % P
        self.input_buffer = []
        # Apply the permutation.
        self.sponge_state = P2.permute(self.sponge_state)
        # output_buffer = sponge_state[0..RATE]
        self.output_buffer = list(self.sponge_state[0:RATE])

    def observe(self, value):
        # Any buffered output is now invalid.
        self.output_buffer = []
        self.input_buffer.append(value % P)
        if len(self.input_buffer) == RATE:
            self.duplexing()

    def observe_slice(self, values):
        for v in values:
            self.observe(v)

    def observe_ext(self, ext):
        """observe_ext_element: absorb the 4 base coords [c0, c1, c2, c3] in order."""
        self.observe_slice(ext)

    def sample_base(self):
        """CanSample<F>::sample (from_base_fn, D=1): duplex if inputs are buffered or outputs are
        exhausted, then pop the last output element. Returns a canonical field element in [0, p)."""
        if self.input_buffer or not self.output_buffer:
            self.duplexing()
        return self.output_buffer.pop()  # Vec::pop -> last element

    def sample_ext(self):
        """sample_ext_element::<F_{p^4}>(): sample_vec(4) then from_base_slice -> [c0, c1, c2, c3]."""
        return [self.sample_base() for _ in range(4)]

    def sample_bits(self, bits):
        """CanSampleBits::sample_bits: low `bits` bits of a base-field sample's canonical rep."""
        rand_f = self.sample_base()
        return rand_f & ((1 << bits) - 1)

    def check_witness(self, bits, witness):
        """GrindingChallenger::check_witness: observe(witness); sample_bits(bits) == 0."""
        self.observe(witness)
        return self.sample_bits(bits) == 0


# ============================ pure duplex KAT (mirrors the Rust script) ============================

def run_duplex_script():
    """Run the exact scripted transcript the extractor emits outputs for. Returns a dict of outputs."""
    ch = DuplexChallenger()
    ch.observe(11)
    ch.observe(22)
    a = ch.sample_base()
    b = ch.sample_base()
    c = ch.sample_ext()
    ch.observe(33)
    ch.observe(44)
    ch.observe(55)
    d = ch.sample_bits(KAT_SAMPLE_BITS)
    ch.observe_slice(KAT_DIGEST)
    e = ch.sample_ext()
    f = ch.sample_base()
    g = ch.sample_base()
    h = ch.sample_base()
    i = ch.sample_base()
    j = ch.sample_ext()
    state_after = list(ch.sponge_state)
    return {"ch": ch, "a": a, "b": b, "c": c, "d": d, "e": e,
            "f": f, "g": g, "h": h, "i": i, "j": j, "state_after": state_after}


def duplex_witness_table(ch, witnesses):
    """Reproduce the check_witness accept/reject table from the checkpoint state (each on a clone)."""
    out = []
    for w in witnesses:
        clone = _clone(ch)
        out.append({"witness": w, "accept": clone.check_witness(KAT_POW_BITS, w)})
    return out


def _clone(ch):
    c = DuplexChallenger()
    c.sponge_state = list(ch.sponge_state)
    c.input_buffer = list(ch.input_buffer)
    c.output_buffer = list(ch.output_buffer)
    return c


# ============================ FRI transcript closure ============================

def derive_fri_challenges(commit_phase_commits, final_poly, pow_witness, pow_bits,
                          log_blowup, num_queries):
    """Reproduce p3-fri verify_shape_and_sample_challenges over the DuplexChallenger:
      for each commit: observe(commit digest); beta = sample_ext()
      observe_ext(final_poly)
      pow_accept = check_witness(pow_bits, pow_witness)
      log_max_height = len(commits) + log_blowup
      query_indices = [sample_bits(log_max_height) for _ in range(num_queries)]
    Returns (betas, query_indices, pow_accept). The betas/indices are DERIVED here, not taken as
    inputs; this is exactly what component (f) adds on top of the confirmed FRI verify_query."""
    ch = DuplexChallenger()
    betas = []
    for comm in commit_phase_commits:
        ch.observe_slice(comm)     # observe the 8-element commitment digest
        betas.append(ch.sample_ext())
    ch.observe_ext(final_poly)     # observe the F_{p^4} final polynomial
    pow_accept = ch.check_witness(pow_bits, pow_witness)
    log_max_height = len(commit_phase_commits) + log_blowup
    query_indices = [ch.sample_bits(log_max_height) for _ in range(num_queries)]
    return betas, query_indices, pow_accept


# ============================ conformance KAT (real ground truth) ============================

def conformance_kats():
    """Return (passed, total, details). Ties to BOTH the challenger ground truth (pure duplex KAT +
    FRI transcript) AND fri_ground_truth.json (closing the FRI betas/indices loop)."""
    with open(CHALLENGER_GT) as fp:
        gt = json.load(fp)
    with open(FRI_GT) as fp:
        fri_gt = json.load(fp)

    passed = 0
    total = 0
    details = []

    def check(name, cond):
        nonlocal passed, total
        total += 1
        if cond:
            passed += 1
        details.append({"check": name, "pass": bool(cond)})

    # (0) sanity: the extractor's perm matches the confirmed Poseidon2 permutation.
    check("perm_zeros_matches_poseidon2", P2.permute([0] * 16) == gt["perm_zeros"])

    # (1) pure duplex KAT.
    k = gt["duplex_kat"]
    r = run_duplex_script()
    check("duplex.a", r["a"] == k["a"])
    check("duplex.b", r["b"] == k["b"])
    check("duplex.c", r["c"] == k["c"])
    check("duplex.d", r["d"] == k["d"])
    check("duplex.e", r["e"] == k["e"])
    check("duplex.f", r["f"] == k["f"])
    check("duplex.g", r["g"] == k["g"])
    check("duplex.h", r["h"] == k["h"])
    check("duplex.i", r["i"] == k["i"])
    check("duplex.j", r["j"] == k["j"])
    check("duplex.state_after", r["state_after"] == k["state_after"])
    # check_witness accept/reject table
    witnesses = [w["witness"] for w in k["witness_table"]]
    my_table = duplex_witness_table(r["ch"], witnesses)
    table_ok = all(my_table[n]["accept"] == k["witness_table"][n]["accept"]
                   for n in range(len(witnesses)))
    check("duplex.witness_table", table_ok)
    # sanity: the table exhibits BOTH accept and reject (a meaningful grinding test).
    accepts = [w["accept"] for w in k["witness_table"]]
    check("duplex.table_has_accept_and_reject", any(accepts) and not all(accepts))

    # (2) FRI transcript closure: derive betas/indices/pow-accept and match extractor + fri_ground_truth.
    fri_by_name = {c["name"]: c for c in fri_gt["cases"]}
    for tc in gt["fri_transcript"]:
        name = tc["name"]
        betas, indices, pow_accept = derive_fri_challenges(
            tc["commit_phase_commits"], tc["final_poly"], tc["pow_witness"],
            tc["pow_bits"], tc["log_blowup"], tc["num_queries"])
        check(f"fri[{name}].pow_accept", pow_accept is True)
        check(f"fri[{name}].betas_match_extractor", betas == tc["betas"])
        check(f"fri[{name}].indices_match_extractor", indices == tc["query_indices"])
        # close the loop: the DERIVED betas/indices equal the ones the CONFIRMED FRI KAT consumed.
        fc = fri_by_name[name]
        check(f"fri[{name}].betas_close_fri_loop", betas == fc["betas"])
        check(f"fri[{name}].indices_close_fri_loop", indices == [q["index"] for q in fc["queries"]])

    return passed, total, details


# ============================ shared cross-language vector set ============================

def shared_vectors():
    r = run_duplex_script()
    with open(CHALLENGER_GT) as fp:
        gt = json.load(fp)
    witnesses = [w["witness"] for w in gt["duplex_kat"]["witness_table"]]
    fri_out = []
    for tc in gt["fri_transcript"]:
        betas, indices, pow_accept = derive_fri_challenges(
            tc["commit_phase_commits"], tc["final_poly"], tc["pow_witness"],
            tc["pow_bits"], tc["log_blowup"], tc["num_queries"])
        fri_out.append({"name": tc["name"], "betas": betas, "queryIndices": indices,
                        "powAccept": pow_accept})
    return {
        "permZeros": P2.permute([0] * 16),
        "duplex": {"a": r["a"], "b": r["b"], "c": r["c"], "d": r["d"], "e": r["e"],
                   "f": r["f"], "g": r["g"], "h": r["h"], "i": r["i"], "j": r["j"],
                   "stateAfter": r["state_after"],
                   "witnessTable": duplex_witness_table(r["ch"], witnesses)},
        "fri": fri_out,
    }


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "kat":
        p, t, det = conformance_kats()
        for d in det:
            print(("PASS" if d["pass"] else "FAIL"), d["check"])
        print(f"conformance: PASS={p} FAIL={t - p} TOTAL={t}")
        sys.exit(0 if p == t else 1)
    json.dump(shared_vectors(), sys.stdout)