// Ground-truth extractor for the FRI low-degree test (component (d)) that SP1/Plonky3 use over
// BabyBear, for the PQ STARK-verify precompile 0x0AE8. It builds the EXACT SP1 inner FRI config
//   Val        = BabyBear
//   Challenge  = BinomialExtensionField<BabyBear, 4>            (the FRI folding field EF)
//   Perm       = Poseidon2<BabyBear, Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, 16, 7>
//   MyHash     = PaddingFreeSponge<Perm, 16, 8, 8>
//   MyCompress = TruncatedPermutation<Perm, 2, 8, 16>
//   ValMmcs    = FieldMerkleTreeMmcs<Packing, Packing, MyHash, MyCompress, 8>   (component (c), CONFIRMED)
//   FriMmcs    = ExtensionMmcs<Val, Challenge, ValMmcs>          (the commit-phase MMCS over EF)
//   Challenger = DuplexChallenger<BabyBear, Perm, 16, 8>         (SP1 inner challenger, component (f))
// The permutation is built with Xoroshiro128Plus::seed_from_u64(1) via new_from_rng_128, the SAME
// deterministic construction the CONFIRMED Poseidon2/MMCS references use (a perm_zeros sanity KAT is
// emitted so the harness can assert it before trusting anything).
//
// For each FRI case it: constructs a genuinely low-degree codeword (an RS/LDE of a random low-degree
// polynomial over the two-adic subgroup, in bit-reversed order), runs the pinned p3-fri PROVER to emit
// a real FriProof (commit-phase roots, per-query openings + MMCS proofs, betas, final poly, pow
// witness), then runs the pinned p3-fri VERIFIER (verify_shape_and_sample_challenges +
// verify_challenges) to confirm the library itself accepts. It emits every value the reference FRI
// verifier needs to reproduce verify_query: the betas and query indices (which the real transcript /
// component (f) derives, supplied here as inputs), the reduced openings, and per-query per-layer
// sibling values + MMCS opening proofs. Field elements are printed as canonical u32.

use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
use p3_challenger::DuplexChallenger;
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_challenges, 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()
}

// An extension field element as its 4 canonical base coords [c0, c1, c2, c3].
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(","))
}

struct FriCase {
    name: &'static str,
    log_blowup: usize,
    log_max_height: usize,
    num_queries: usize,
    pow_bits: usize,
    seed: u64,
}

fn main() {
    // Deterministic perm identical to the CONFIRMED Poseidon2 reference (seed_from_u64(1)).
    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: permute([0;16]) must equal the confirmed Poseidon2 "zeros" vector ----
    let zeros = perm.permute([Val::zero(); 16]);
    out.push_str(&format!("\"perm_zeros\":{},", ja(&u32s(&zeros))));

    // ---- two-adic generators g(bits) for bits 0..=27 (confirms component (a) twoAdicGenerator) ----
    let gens: Vec<u32> = (0..=27)
        .map(|b| Val::two_adic_generator(b).as_canonical_u32())
        .collect();
    out.push_str(&format!("\"two_adic_generators\":{},", ja(&gens)));

    // ---- FRI cases ----
    let cases = vec![
        FriCase { name: "blowup1_h6_q4", log_blowup: 1, log_max_height: 6, num_queries: 4, pow_bits: 1, seed: 42 },
        FriCase { name: "blowup2_h7_q5", log_blowup: 2, log_max_height: 7, num_queries: 5, pow_bits: 0, seed: 7 },
        FriCase { name: "blowup1_h8_q6", log_blowup: 1, log_max_height: 8, num_queries: 6, pow_bits: 3, seed: 12345 },
    ];

    out.push_str("\"cases\":[");
    for (ci, case) in cases.iter().enumerate() {
        if ci > 0 {
            out.push(',');
        }
        out.push_str(&run_case(&perm, case));
    }
    out.push_str("]}");

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

fn run_case(perm: &Perm, case: &FriCase) -> 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: case.log_blowup,
        num_queries: case.num_queries,
        proof_of_work_bits: case.pow_bits,
        mmcs: challenge_mmcs,
    };

    let n = 1usize << case.log_max_height;
    let k = 1usize << (case.log_max_height - case.log_blowup); // message length (degree bound)

    // Build a genuinely low-degree codeword: coeffs of a degree-<k poly (rest zero), DFT to evals over
    // the size-n subgroup, then bit-reverse (the order the FRI prover's `input` expects).
    let mut state: u64 = case.seed.wrapping_mul(0x9E3779B97F4A7C15).wrapping_add(1);
    let mut next = || {
        // simple deterministic LCG-ish generator for reproducible coeffs
        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); // natural order over the subgroup
    reverse_slice_index_bits(&mut evals); // bit-reversed order (prover input convention)

    // input[log_max_height] = Some(codeword); all other heights None (single-codeword LDT).
    let mut input: [Option<Vec<Challenge>>; 32] = core::array::from_fn(|_| None);
    input[case.log_max_height] = Some(evals.clone());

    // ---- PROVE (pinned p3-fri prover) ----
    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's transcript (component (f)) ----
    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");
    assert_eq!(
        challenges.query_indices, prover_indices,
        "verifier-derived query indices must equal the prover's"
    );

    // ---- reduced openings: single codeword => ro[log_max_height] = codeword value at index ----
    let mut ros: Vec<[Challenge; 32]> = Vec::new();
    for &index in &challenges.query_indices {
        let mut ro = [Challenge::zero(); 32];
        ro[case.log_max_height] = evals[index];
        ros.push(ro);
    }

    // ---- VERIFY (pinned p3-fri verifier: the fold + opening relations) : must accept ----
    verify_challenges(&config, &proof, &challenges, &ros)
        .expect("library verify_challenges must accept the honest proof");

    // ---- emit everything the reference FRI verifier needs ----
    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);

    // per-query records
    let mut queries_json: Vec<String> = Vec::new();
    for (qi, &index) in challenges.query_indices.iter().enumerate() {
        let ro_top = ef_u32(&ros[qi][case.log_max_height]);
        let qp = &proof.query_proofs[qi];
        let mut layers_json: Vec<String> = Vec::new();
        for step in &qp.commit_phase_openings {
            let sib = ef_u32(&step.sibling_value);
            let op: Vec<Vec<u32>> = step.opening_proof.iter().map(|d| u32s(d)).collect();
            layers_json.push(format!(
                "{{\"sibling_value\":{},\"opening_proof\":{}}}",
                ja(&sib),
                jaa(&op)
            ));
        }
        queries_json.push(format!(
            "{{\"index\":{},\"ro_top\":{},\"layers\":[{}]}}",
            index,
            ja(&ro_top),
            layers_json.join(",")
        ));
    }

    format!(
        "{{\"name\":\"{}\",\"log_blowup\":{},\"log_max_height\":{},\"num_queries\":{},\"pow_bits\":{},\"commit_phase_commits\":{},\"betas\":{},\"final_poly\":{},\"queries\":[{}]}}",
        case.name,
        case.log_blowup,
        case.log_max_height,
        case.num_queries,
        case.pow_bits,
        jaa(&commits),
        jaa(&betas),
        ja(&final_poly),
        queries_json.join(",")
    )
}