// Ground-truth extractor for the Poseidon2 DuplexChallenger (Fiat-Shamir transcript, component (f))
// and the GrindingChallenger proof-of-work check that SP1/Plonky3 use over BabyBear, for the PQ
// STARK-verify precompile 0x0AE8. It pins the EXACT SP1 inner challenger:
//   Val        = BabyBear
//   Challenge  = BinomialExtensionField<BabyBear, 4>
//   Perm       = Poseidon2<BabyBear, Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, 16, 7>
//   Challenger = DuplexChallenger<BabyBear, Perm, 16, 8>          (WIDTH=16, RATE=8)
// The permutation is built with Xoroshiro128Plus::seed_from_u64(1) via new_from_rng_128, the SAME
// deterministic construction the CONFIRMED Poseidon2/MMCS/FRI references use (a perm_zeros sanity KAT
// is emitted so the harness can assert it before trusting anything).
//
// It emits TWO ground-truth blocks:
//   (1) "duplex_kat": a fully-scripted observe/sample transcript over the pinned DuplexChallenger:
//       observe base elements and an 8-element digest, sample base elements, sample F_{p^4} ext
//       elements (via sample_ext_element, the exact call the FRI verifier uses for betas), sample_bits
//       (the query-index step), and check_witness (grinding). It records every sampled output, the full
//       16-lane sponge_state at a checkpoint, and a table of (witness -> accept?) booleans so the port
//       can reproduce accept AND reject byte-for-byte.
//   (2) "fri_transcript": the SAME three FRI cases the fri-extractor builds (same seeds), rebuilt so we
//       can emit each case's commit_phase_commits, final_poly, pow_bits, pow_witness, and the betas +
//       query_indices that the pinned p3-fri verifier's transcript (verify_shape_and_sample_challenges)
//       derived. This lets the port reproduce the betas/indices FROM THE TRANSCRIPT and close the loop
//       the FRI KAT left open (transcript -> betas/indices -> the confirmed FRI verify_query).
// Field elements are printed as canonical u32.

use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
use p3_challenger::{
    CanObserve, CanSample, CanSampleBits, DuplexChallenger, FieldChallenger, GrindingChallenger,
};
use p3_commit::ExtensionMmcs;
use p3_dft::{Radix2Dit, TwoAdicSubgroupDft};
use p3_field::extension::BinomialExtensionField;
use p3_field::{AbstractExtensionField, AbstractField, Field, PrimeField32, TwoAdicField};
use p3_fri::verifier::verify_shape_and_sample_challenges;
use p3_fri::FriConfig;
use p3_merkle_tree::FieldMerkleTreeMmcs;
use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral};
use p3_symmetric::{PaddingFreeSponge, Permutation, TruncatedPermutation};
use p3_util::reverse_slice_index_bits;
use rand::SeedableRng;
use rand_xoshiro::Xoroshiro128Plus;

type Val = BabyBear;
type Challenge = BinomialExtensionField<BabyBear, 4>;
type Perm = Poseidon2<Val, Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, 16, 7>;
type MyHash = PaddingFreeSponge<Perm, 16, 8, 8>;
type MyCompress = TruncatedPermutation<Perm, 2, 8, 16>;
type ValMmcs =
    FieldMerkleTreeMmcs<<Val as Field>::Packing, <Val as Field>::Packing, MyHash, MyCompress, 8>;
type ChallengeMmcs = ExtensionMmcs<Val, Challenge, ValMmcs>;
type Challenger = DuplexChallenger<Val, Perm, 16, 8>;

fn u32s(row: &[Val]) -> Vec<u32> {
    row.iter().map(|f| f.as_canonical_u32()).collect()
}

fn ef_u32(e: &Challenge) -> Vec<u32> {
    let base: &[Val] = <Challenge as AbstractExtensionField<Val>>::as_base_slice(e);
    base.iter().map(|f| f.as_canonical_u32()).collect()
}

fn ja(v: &[u32]) -> String {
    let s: Vec<String> = v.iter().map(|x| x.to_string()).collect();
    format!("[{}]", s.join(","))
}

fn jaa(v: &[Vec<u32>]) -> String {
    let s: Vec<String> = v.iter().map(|x| ja(x)).collect();
    format!("[{}]", s.join(","))
}

// The fixed digest the pure KAT observes (an 8-element "commitment").
const DIG: [u32; 8] = [101, 102, 103, 104, 105, 106, 107, 108];
// bits for the sample_bits step in the pure KAT (query-index style reduction).
const KAT_SAMPLE_BITS: usize = 10;
// proof-of-work bits for the check_witness sub-KAT.
const KAT_POW_BITS: usize = 4;
// candidate witnesses for the check_witness accept/reject table (index 0 is replaced by a real
// grinding witness at runtime so at least one entry accepts).
const KAT_WITNESS_CANDIDATES: [u32; 7] = [0, 0, 1, 2, 3, 7, 12345];

fn main() {
    let mut rng = Xoroshiro128Plus::seed_from_u64(1);
    let perm = Perm::new_from_rng_128(
        Poseidon2ExternalMatrixGeneral,
        DiffusionMatrixBabyBear,
        &mut rng,
    );

    let mut out = String::new();
    out.push('{');

    // ---- perm sanity KAT ----
    let zeros = perm.permute([Val::zero(); 16]);
    out.push_str(&format!("\"perm_zeros\":{},", ja(&u32s(&zeros))));

    // ---- (1) pure DuplexChallenger transcript KAT ----
    out.push_str(&format!("\"duplex_kat\":{},", duplex_kat(&perm)));

    // ---- (2) FRI transcript closure ----
    let cases = vec![
        (("blowup1_h6_q4"), 1usize, 6usize, 4usize, 1usize, 42u64),
        (("blowup2_h7_q5"), 2, 7, 5, 0, 7),
        (("blowup1_h8_q6"), 1, 8, 6, 3, 12345),
    ];
    out.push_str("\"fri_transcript\":[");
    for (ci, c) in cases.iter().enumerate() {
        if ci > 0 {
            out.push(',');
        }
        out.push_str(&fri_transcript_case(&perm, c.0, c.1, c.2, c.3, c.4, c.5));
    }
    out.push_str("]}");

    println!("{}", out);
}

// A fully-scripted observe/sample transcript over the pinned DuplexChallenger. The op sequence here
// is mirrored EXACTLY by the Python/Node/Java references; this extractor emits only the OUTPUTS.
fn duplex_kat(perm: &Perm) -> String {
    let f = Val::from_canonical_u32;
    let mut ch = Challenger::new(perm.clone());

    // observe two base elements (partial input buffer, < RATE)
    ch.observe(f(11));
    ch.observe(f(22));
    // sample two base elements (forces a duplex with a partial input buffer, then drains output)
    let a: Val = ch.sample();
    let b: Val = ch.sample();
    // sample an ext element (4 base pops)
    let c: Challenge = ch.sample_ext_element();
    // observe three more base elements
    ch.observe(f(33));
    ch.observe(f(44));
    ch.observe(f(55));
    // sample_bits (query-index style; forces a duplex because input buffer is non-empty)
    let d = ch.sample_bits(KAT_SAMPLE_BITS);
    // observe an 8-element digest (fills RATE exactly -> auto-duplex on the 8th observe)
    ch.observe(DIG.map(f));
    // sample an ext element then four base elements then another ext element (crosses a duplex
    // boundary and exercises an empty-output duplex)
    let e: Challenge = ch.sample_ext_element();
    let g0: Val = ch.sample();
    let g1: Val = ch.sample();
    let g2: Val = ch.sample();
    let g3: Val = ch.sample();
    let jext: Challenge = ch.sample_ext_element();

    // checkpoint: full 16-lane sponge state after the main script (before the PoW sub-KAT)
    let state_after: Vec<u32> = ch.sponge_state.iter().map(|x| x.as_canonical_u32()).collect();

    // check_witness accept/reject table from the checkpoint state. Slot 0 gets a REAL grinding
    // witness (guaranteed accept); the rest are fixed candidates (reject unless they happen to hit 0).
    let mut witnesses = KAT_WITNESS_CANDIDATES;
    let good = ch.clone().grind(KAT_POW_BITS).as_canonical_u32();
    witnesses[0] = good;
    let mut wtable: Vec<String> = Vec::new();
    for w in witnesses.iter() {
        let mut cw = ch.clone();
        let accept = cw.check_witness(KAT_POW_BITS, f(*w));
        wtable.push(format!("{{\"witness\":{},\"accept\":{}}}", w, accept));
    }

    format!(
        "{{\"script\":\"observe[11,22];sampleBase A,B;sampleExt C;observe[33,44,55];sampleBits({sb}) D;observeDigest;sampleExt E;sampleBase F,G,H,I;sampleExt J\",\
\"sampleBits\":{sb},\"powBits\":{pw},\"digest\":{dig},\
\"a\":{a},\"b\":{b},\"c\":{c},\"d\":{d},\"e\":{e},\"f\":{gf},\"g\":{gg},\"h\":{gh},\"i\":{gi},\"j\":{jj},\
\"state_after\":{st},\"witness_table\":[{wt}]}}",
        sb = KAT_SAMPLE_BITS,
        pw = KAT_POW_BITS,
        dig = ja(&DIG),
        a = a.as_canonical_u32(),
        b = b.as_canonical_u32(),
        c = ja(&ef_u32(&c)),
        d = d,
        e = ja(&ef_u32(&e)),
        gf = g0.as_canonical_u32(),
        gg = g1.as_canonical_u32(),
        gh = g2.as_canonical_u32(),
        gi = g3.as_canonical_u32(),
        jj = ja(&ef_u32(&jext)),
        st = ja(&state_after),
        wt = wtable.join(",")
    )
}

// Rebuild one FRI case identically to the fri-extractor (same seed/coeff generation), run the pinned
// p3-fri verifier transcript, and emit the transcript inputs + the derived betas/query_indices +
// pow_witness so the port can reproduce them from the transcript alone.
fn fri_transcript_case(
    perm: &Perm,
    name: &str,
    log_blowup: usize,
    log_max_height: usize,
    num_queries: usize,
    pow_bits: usize,
    seed: u64,
) -> String {
    let hash = MyHash::new(perm.clone());
    let compress = MyCompress::new(perm.clone());
    let val_mmcs = ValMmcs::new(hash, compress);
    let challenge_mmcs = ChallengeMmcs::new(val_mmcs);

    let config = FriConfig {
        log_blowup,
        num_queries,
        proof_of_work_bits: pow_bits,
        mmcs: challenge_mmcs,
    };

    let n = 1usize << log_max_height;
    let k = 1usize << (log_max_height - log_blowup);

    let mut state: u64 = seed.wrapping_mul(0x9E3779B97F4A7C15).wrapping_add(1);
    let mut next = || {
        state = state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        ((state >> 33) as u32) % (BabyBear::ORDER_U32)
    };
    let mut coeffs: Vec<Challenge> = Vec::with_capacity(n);
    for i in 0..n {
        if i < k {
            let c0 = Val::from_canonical_u32(next());
            let c1 = Val::from_canonical_u32(next());
            let c2 = Val::from_canonical_u32(next());
            let c3 = Val::from_canonical_u32(next());
            coeffs.push(Challenge::from_base_slice(&[c0, c1, c2, c3]));
        } else {
            coeffs.push(Challenge::zero());
        }
    }
    let dft = Radix2Dit::<Challenge>::default();
    let mut evals = dft.dft(coeffs);
    reverse_slice_index_bits(&mut evals);

    let mut input: [Option<Vec<Challenge>>; 32] = core::array::from_fn(|_| None);
    input[log_max_height] = Some(evals.clone());

    let mut p_challenger = Challenger::new(perm.clone());
    let (proof, _prover_indices) = p3_fri::prover::prove(&config, &input, &mut p_challenger);

    // Re-derive challenges via the pinned verifier transcript (component (f) ground truth).
    let mut v_challenger = Challenger::new(perm.clone());
    let challenges = verify_shape_and_sample_challenges(&config, &proof, &mut v_challenger)
        .expect("verify_shape_and_sample_challenges must succeed");

    let commits: Vec<Vec<u32>> = proof
        .commit_phase_commits
        .iter()
        .map(|c| {
            let arr: [Val; 8] = (*c).into();
            u32s(&arr)
        })
        .collect();
    let betas: Vec<Vec<u32>> = challenges.betas.iter().map(ef_u32).collect();
    let final_poly = ef_u32(&proof.final_poly);
    let pow_witness = proof.pow_witness.as_canonical_u32();
    let indices: Vec<u32> = challenges.query_indices.iter().map(|&i| i as u32).collect();

    format!(
        "{{\"name\":\"{}\",\"log_blowup\":{},\"log_max_height\":{},\"num_queries\":{},\"pow_bits\":{},\
\"commit_phase_commits\":{},\"final_poly\":{},\"pow_witness\":{},\"betas\":{},\"query_indices\":{}}}",
        name,
        log_blowup,
        log_max_height,
        num_queries,
        pow_bits,
        jaa(&commits),
        ja(&final_poly),
        pow_witness,
        jaa(&betas),
        ja(&indices)
    )
}