// END-TO-END ground-truth extractor for the assembled BabyBear + FRI STARK verifier (0x0AE8 reference).
//
// Unlike airquotient-extractor (which emits only the OOD openings + alpha/zeta and IGNORES the FRI
// opening_proof), this binary emits the COMPLETE p3-uni-stark Proof so the assembled Python reference can
// run the WHOLE verify pipeline end to end from proof + config alone:
//   transcript -> observe(trace) -> sample alpha -> observe(quotient) -> sample zeta ->
//   PCS.verify (sample FRI alpha, verify_shape_and_sample_challenges -> betas/indices/pow,
//               reduced-opening combination + input-batch MMCS opening, verify_challenges -> verify_query)
//   -> AIR quotient consistency check.
//
// For each of two KNOWN example AIRs (Fibonacci from p3-uni-stark's own tests/fib_air.rs, and a degree-3
// multiply AIR matching tests/mul_air.rs) it: (1) runs the pinned p3-uni-stark PROVER, (2) runs the pinned
// p3-uni-stark VERIFIER and asserts ACCEPT (the ground truth), (3) recovers the full proof via the real
// Deserialize types and emits everything as canonical u32 JSON. NOTE: this is an EXAMPLE p3-uni-stark AIR
// and an EXAMPLE FRI config (log_blowup=2, num_queries=28, pow_bits=8), NOT SP1 6.1.0 (which is Hypercube),
// and NOT an Aere production circuit (Aere's zk-circuits are SP1 guest programs). See the port spec.

use std::borrow::Borrow;

use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
use p3_challenger::DuplexChallenger;
use p3_commit::ExtensionMmcs;
use p3_dft::Radix2DitParallel;
use p3_field::extension::BinomialExtensionField;
use p3_field::{AbstractExtensionField, AbstractField, Field, PrimeField32};
use p3_fri::{BatchOpening, FriConfig, TwoAdicFriPcs, TwoAdicFriPcsProof};
use p3_matrix::dense::RowMajorMatrix;
use p3_matrix::Matrix;
use p3_merkle_tree::FieldMerkleTreeMmcs;
use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral};
use p3_symmetric::{Hash, PaddingFreeSponge, TruncatedPermutation};
use p3_uni_stark::{prove, verify, StarkConfig};
use rand::SeedableRng;
use rand_xoshiro::Xoroshiro128Plus;
use serde::Deserialize;
use serde_json::{json, Value};

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>;
type Dft = Radix2DitParallel;
type Pcs = TwoAdicFriPcs<Val, Dft, ValMmcs, ChallengeMmcs>;
type MyConfig = StarkConfig<Pcs, Challenge, Challenger>;
type Com = Hash<Val, Val, 8>;

// FRI config used here (example, NOT SP1). Emitted in the JSON so the Python reference reads it.
const LOG_BLOWUP: usize = 2;
const NUM_QUERIES: usize = 28;
const POW_BITS: usize = 8;

// ---- full deserialize mirror of Proof<SC> (fields are pub(crate); recovered via serde round-trip) ----
#[derive(Deserialize)]
struct MirrorCommitments {
    trace: Com,
    quotient_chunks: Com,
}
#[derive(Deserialize)]
struct MirrorOpened {
    trace_local: Vec<Challenge>,
    trace_next: Vec<Challenge>,
    quotient_chunks: Vec<Vec<Challenge>>,
}
#[derive(Deserialize)]
struct MirrorProof {
    commitments: MirrorCommitments,
    opened_values: MirrorOpened,
    // The REAL opening-proof type (all fields pub), so we can walk fri_proof + query_openings.
    opening_proof: TwoAdicFriPcsProof<Val, Challenge, ValMmcs, ChallengeMmcs>,
    degree_bits: usize,
}

// ============================ example AIR 1: Fibonacci (verbatim from tests/fib_air.rs) ============
const NUM_FIBONACCI_COLS: usize = 2;
pub struct FibonacciAir {}
impl<F> BaseAir<F> for FibonacciAir {
    fn width(&self) -> usize {
        NUM_FIBONACCI_COLS
    }
}
impl<AB: AirBuilderWithPublicValues> Air<AB> for FibonacciAir {
    fn eval(&self, builder: &mut AB) {
        let main = builder.main();
        let pis = builder.public_values();
        let a = pis[0];
        let b = pis[1];
        let x = pis[2];
        let (local, next) = (main.row_slice(0), main.row_slice(1));
        let local: &FibonacciRow<AB::Var> = (*local).borrow();
        let next: &FibonacciRow<AB::Var> = (*next).borrow();
        let mut when_first_row = builder.when_first_row();
        when_first_row.assert_eq(local.left, a);
        when_first_row.assert_eq(local.right, b);
        let mut when_transition = builder.when_transition();
        when_transition.assert_eq(local.right, next.left);
        when_transition.assert_eq(local.left + local.right, next.right);
        builder.when_last_row().assert_eq(local.right, x);
    }
}
pub fn fib_trace<F: p3_field::PrimeField64>(a: u64, b: u64, n: usize) -> RowMajorMatrix<F> {
    assert!(n.is_power_of_two());
    let mut trace = RowMajorMatrix::new(vec![F::zero(); n * NUM_FIBONACCI_COLS], NUM_FIBONACCI_COLS);
    let (prefix, rows, suffix) = unsafe { trace.values.align_to_mut::<FibonacciRow<F>>() };
    assert!(prefix.is_empty() && suffix.is_empty());
    assert_eq!(rows.len(), n);
    rows[0] = FibonacciRow::new(F::from_canonical_u64(a), F::from_canonical_u64(b));
    for i in 1..n {
        rows[i].left = rows[i - 1].right;
        rows[i].right = rows[i - 1].left + rows[i - 1].right;
    }
    trace
}
pub struct FibonacciRow<F> {
    pub left: F,
    pub right: F,
}
impl<F> FibonacciRow<F> {
    const fn new(left: F, right: F) -> FibonacciRow<F> {
        FibonacciRow { left, right }
    }
}
impl<F> Borrow<FibonacciRow<F>> for [F] {
    fn borrow(&self) -> &FibonacciRow<F> {
        debug_assert_eq!(self.len(), NUM_FIBONACCI_COLS);
        let (prefix, shorts, suffix) = unsafe { self.align_to::<FibonacciRow<F>>() };
        debug_assert!(prefix.is_empty() && suffix.is_empty());
        debug_assert_eq!(shorts.len(), 1);
        &shorts[0]
    }
}

// ============================ example AIR 2: degree-3 multiply AIR (matches tests/mul_air.rs) ======
const MUL_WIDTH: usize = 3;
pub struct MulAir {}
impl<F> BaseAir<F> for MulAir {
    fn width(&self) -> usize {
        MUL_WIDTH
    }
}
impl<AB: AirBuilder> Air<AB> for MulAir {
    fn eval(&self, builder: &mut AB) {
        let main = builder.main();
        let local = main.row_slice(0);
        let next = main.row_slice(1);
        let a = local[0];
        let b = local[1];
        let c = local[2];
        builder.assert_zero(a.into().exp_u64(2) * b.into() - c.into());
        builder
            .when_first_row()
            .assert_eq(a.into() * a.into() + AB::Expr::one(), b);
        let next_a = next[0];
        builder
            .when_transition()
            .assert_eq(a.into() + AB::Expr::one(), next_a);
    }
}
fn mul_trace(n: usize) -> RowMajorMatrix<Val> {
    let mut v = vec![Val::zero(); n * MUL_WIDTH];
    for i in 0..n {
        let a = Val::from_canonical_u64(i as u64);
        let b = if i == 0 {
            a * a + Val::one()
        } else {
            Val::from_canonical_u64(2 * (i as u64) + 5)
        };
        let c = a * a * b;
        v[i * MUL_WIDTH] = a;
        v[i * MUL_WIDTH + 1] = b;
        v[i * MUL_WIDTH + 2] = c;
    }
    RowMajorMatrix::new(v, MUL_WIDTH)
}

// ============================ canonical-u32 emit helpers ============================
fn val_u32(v: &Val) -> u32 {
    v.as_canonical_u32()
}
fn ef_u32(e: &Challenge) -> Vec<u32> {
    <Challenge as AbstractExtensionField<Val>>::as_base_slice(e)
        .iter()
        .map(|f| f.as_canonical_u32())
        .collect()
}
fn com_u32(c: &Com) -> Vec<u32> {
    let a: [Val; 8] = (*c).into();
    a.iter().map(|f| f.as_canonical_u32()).collect()
}
// A ValMmcs / ChallengeMmcs proof is Vec<[Val; 8]> (list of sibling digests).
fn digests_u32(p: &[[Val; 8]]) -> Vec<Vec<u32>> {
    p.iter()
        .map(|d| d.iter().map(|f| f.as_canonical_u32()).collect())
        .collect()
}

fn config_and_perm(log_n: usize) -> (MyConfig, Perm) {
    let mut rng = Xoroshiro128Plus::seed_from_u64(1); // CONFIRMED Poseidon2 seed.
    let perm = Perm::new_from_rng_128(Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, &mut rng);
    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.clone());
    let dft = Dft {};
    let fri_config = FriConfig {
        log_blowup: LOG_BLOWUP,
        num_queries: NUM_QUERIES,
        proof_of_work_bits: POW_BITS,
        mmcs: challenge_mmcs,
    };
    let pcs = Pcs::new(log_n, dft, val_mmcs, fri_config);
    (MyConfig::new(pcs), perm)
}

fn emit_case(name: &str, air: &str, mp: &MirrorProof, pis: &[Val], accept: bool) -> Value {
    // opened values
    let trace_local: Vec<Vec<u32>> = mp.opened_values.trace_local.iter().map(ef_u32).collect();
    let trace_next: Vec<Vec<u32>> = mp.opened_values.trace_next.iter().map(ef_u32).collect();
    let quotient_chunks: Vec<Vec<Vec<u32>>> = mp
        .opened_values
        .quotient_chunks
        .iter()
        .map(|ch| ch.iter().map(ef_u32).collect())
        .collect();

    let fp = &mp.opening_proof.fri_proof;
    let commit_phase_commits: Vec<Vec<u32>> = fp.commit_phase_commits.iter().map(com_u32).collect();
    let final_poly = ef_u32(&fp.final_poly);
    let pow_witness = val_u32(&fp.pow_witness);
    let query_proofs: Vec<Value> = fp
        .query_proofs
        .iter()
        .map(|qp| {
            let steps: Vec<Value> = qp
                .commit_phase_openings
                .iter()
                .map(|st| {
                    json!({
                        "sibling_value": ef_u32(&st.sibling_value),
                        "opening_proof": digests_u32(&st.opening_proof),
                    })
                })
                .collect();
            json!({ "commit_phase_openings": steps })
        })
        .collect();

    let query_openings: Vec<Value> = mp
        .opening_proof
        .query_openings
        .iter()
        .map(|per_query| {
            let batches: Vec<Value> = per_query
                .iter()
                .map(|bo: &BatchOpening<Val, ValMmcs>| {
                    let ov: Vec<Vec<u32>> = bo
                        .opened_values
                        .iter()
                        .map(|row| row.iter().map(val_u32).collect())
                        .collect();
                    json!({
                        "opened_values": ov,
                        "opening_proof": digests_u32(&bo.opening_proof),
                    })
                })
                .collect();
            Value::Array(batches)
        })
        .collect();

    let pis_u32: Vec<u32> = pis.iter().map(val_u32).collect();
    json!({
        "name": name,
        "air": air,
        "degree_bits": mp.degree_bits,
        "log_blowup": LOG_BLOWUP,
        "num_queries": NUM_QUERIES,
        "pow_bits": POW_BITS,
        "library_accept": accept,
        "public_values": pis_u32,
        "commitments": {
            "trace": com_u32(&mp.commitments.trace),
            "quotient_chunks": com_u32(&mp.commitments.quotient_chunks),
        },
        "opened_values": {
            "trace_local": trace_local,
            "trace_next": trace_next,
            "quotient_chunks": quotient_chunks,
        },
        "opening_proof": {
            "fri_proof": {
                "commit_phase_commits": commit_phase_commits,
                "final_poly": final_poly,
                "pow_witness": pow_witness,
                "query_proofs": query_proofs,
            },
            "query_openings": query_openings,
        },
    })
}

fn main() {
    let mut cases: Vec<Value> = Vec::new();

    // perm sanity: tie ground truth to the CONFIRMED Poseidon2 permutation + Val generator.
    let (_c0, perm0) = config_and_perm(3);
    let perm_zeros: Vec<u32> = {
        use p3_symmetric::Permutation;
        perm0.permute([Val::zero(); 16]).iter().map(|f| f.as_canonical_u32()).collect()
    };
    let val_generator = <Val as AbstractField>::generator().as_canonical_u32();

    // case A: Fibonacci AIR, n = 8, pis = [0,1,21] (Plonky3's own test vector) -> 1 quotient chunk.
    {
        let log_n = 3;
        let (config, perm) = config_and_perm(log_n);
        let n = 1usize << log_n;
        let trace = fib_trace::<Val>(0, 1, n);
        let pis = vec![Val::from_canonical_u64(0), Val::from_canonical_u64(1), Val::from_canonical_u64(21)];
        let mut p_ch = Challenger::new(perm.clone());
        let proof = prove(&config, &FibonacciAir {}, &mut p_ch, trace, &pis);
        let mut v_ch = Challenger::new(perm.clone());
        let accept = verify(&config, &FibonacciAir {}, &mut v_ch, &proof, &pis).is_ok();
        assert!(accept, "library must accept the honest Fibonacci proof");
        let pj = serde_json::to_value(&proof).expect("serialize proof");
        let mp: MirrorProof = serde_json::from_value(pj).expect("mirror deserialize (fib)");
        cases.push(emit_case("fibonacci_n8", "fibonacci", &mp, &pis, accept));
    }

    // case B: degree-3 multiply AIR, n = 16 -> 2 quotient chunks (exercises multi-chunk PCS batch).
    {
        let log_n = 4;
        let (config, perm) = config_and_perm(log_n);
        let n = 1usize << log_n;
        let trace = mul_trace(n);
        let pis: Vec<Val> = vec![];
        let mut p_ch = Challenger::new(perm.clone());
        let proof = prove(&config, &MulAir {}, &mut p_ch, trace, &pis);
        let mut v_ch = Challenger::new(perm.clone());
        let accept = verify(&config, &MulAir {}, &mut v_ch, &proof, &pis).is_ok();
        assert!(accept, "library must accept the honest MulAir proof");
        let pj = serde_json::to_value(&proof).expect("serialize proof");
        let mp: MirrorProof = serde_json::from_value(pj).expect("mirror deserialize (mul)");
        cases.push(emit_case("mul_deg3_n16", "mul_deg3", &mp, &pis, accept));
    }

    let out = json!({
        "perm_zeros": perm_zeros,
        "val_generator": val_generator,
        "cases": cases,
    });
    println!("{}", serde_json::to_string(&out).unwrap());
}