// Ground-truth extractor for the GENERIC AIR constraint / quotient-consistency check (component (e))
// that a p3-uni-stark STARK verifier performs, for the PQ STARK-verify precompile 0x0AE8.
//
// It builds the EXACT SP1-inner-style config (BabyBear + degree-4 EF + Poseidon2 + FRI + the CONFIRMED
// MMCS/challenger), then for two small, KNOWN example AIRs (the Fibonacci AIR from p3-uni-stark's OWN
// tests/fib_air.rs, and a degree-3 multiplication AIR matching tests/mul_air.rs) it:
//   1. generates a valid trace, runs the pinned p3-uni-stark PROVER to emit a real Proof,
//   2. runs the pinned p3-uni-stark VERIFIER (verify) and asserts it ACCEPTS (the ground truth: the
//      opening argument AND the identity folded_constraints(zeta) == Z_H(zeta) * quotient(zeta) hold),
//   3. re-derives alpha and zeta by REPLAYING the verifier's exact transcript prefix (observe trace
//      commitment; sample_ext alpha; observe quotient_chunks commitment; sample zeta) on a fresh
//      challenger with the same permutation, and
//   4. emits everything the GENERIC quotient-check reference needs: degree_bits, quotient_degree, the
//      out-of-domain trace openings (trace_local at zeta, trace_next at g*zeta), the quotient-chunk
//      openings, alpha, zeta, and the public values. Field elements are canonical u32; EF elements are
//      their 4 canonical base coords [c0,c1,c2,c3] (matching BinomialExtensionField as_base_slice).
//
// The opened_values / commitments fields of Proof are pub(crate); we recover them from the proof's OWN
// serde serialization (BabyBear serializes as canonical u32, the EF as {value:[..]}, the commitment as
// Hash<Val,Val,8>), which is the real emitted proof, not a reconstruction. The reference then reproduces
// the p3-uni-stark verifier's identity check from these inputs (accept), and rejects tampered inputs.

use std::borrow::Borrow;

use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
use p3_challenger::{CanObserve, CanSample, DuplexChallenger, FieldChallenger};
use p3_commit::ExtensionMmcs;
use p3_dft::Radix2DitParallel;
use p3_field::extension::BinomialExtensionField;
use p3_field::{AbstractExtensionField, AbstractField, Field, PrimeField32, PrimeField64};
use p3_fri::{FriConfig, TwoAdicFriPcs};
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 p3_util::log2_ceil_usize;
use rand::SeedableRng;
use rand_xoshiro::Xoroshiro128Plus;
use serde::Deserialize;

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>;

// ---------- serde mirror of Proof<SC> (its fields are pub(crate); recover via serialization) ----------

#[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,
    #[allow(dead_code)]
    opening_proof: serde_json::Value,
    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: 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) ======
// REPETITIONS = 1, degree = 3, boundary + transition constraints. Columns [a, b, c]. Constraints, in
// the SAME order the eval emits them (Horner-folded with alpha by the constraint folder):
//   c1 (all rows):    a^2 * b - c
//   c2 (first row):   is_first_row  * ((a*a + 1) - b)
//   c3 (transition):  is_transition * ((a + 1) - next_a)
// Max constraint degree = 3  =>  log_quotient_degree = ceil(log2(3-1)) = 1  =>  2 quotient chunks.

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];
        // c1: a^2 * b - c  (degree 3), enforced on every row
        builder.assert_zero(a.into().exp_u64(2) * b.into() - c.into());
        // c2: boundary, b == a*a + 1 on the first row
        builder
            .when_first_row()
            .assert_eq(a.into() * a.into() + AB::Expr::one(), b);
        // c3: transition, next_a == a + 1
        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)
}

// ============================ helpers ============================

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(","))
}
fn ef_list(v: &[Challenge]) -> String {
    let s: Vec<String> = v.iter().map(|e| ja(&ef_u32(e))).collect();
    format!("[{}]", s.join(","))
}
fn ef_list2(v: &[Vec<Challenge>]) -> String {
    let s: Vec<String> = v.iter().map(|row| ef_list(row)).collect();
    format!("[{}]", s.join(","))
}

fn config_and_perm(log_n: usize) -> (MyConfig, Perm) {
    // 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 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: 2,
        num_queries: 28,
        proof_of_work_bits: 8,
        mmcs: challenge_mmcs,
    };
    let pcs = Pcs::new(log_n, dft, val_mmcs, fri_config);
    (MyConfig::new(pcs), perm)
}

// Replay the verifier's transcript PREFIX to reproduce (alpha, zeta) exactly.
fn derive_alpha_zeta(perm: &Perm, c: &MirrorCommitments) -> (Challenge, Challenge) {
    let mut ch = Challenger::new(perm.clone());
    ch.observe(c.trace.clone());
    let alpha: Challenge = ch.sample_ext_element();
    ch.observe(c.quotient_chunks.clone());
    let zeta: Challenge = ch.sample();
    (alpha, zeta)
}

fn emit_case(
    name: &str,
    air_id: &str,
    perm: &Perm,
    proof_json: serde_json::Value,
    pis: &[Val],
    library_accept: bool,
) -> String {
    let mp: MirrorProof = serde_json::from_value(proof_json).expect("mirror deserialize");
    let (alpha, zeta) = derive_alpha_zeta(perm, &mp.commitments);
    let quotient_degree = mp.opened_values.quotient_chunks.len();
    let pis_u32: Vec<u32> = pis.iter().map(|p| p.as_canonical_u32()).collect();
    format!(
        "{{\"name\":\"{}\",\"air\":\"{}\",\"degree_bits\":{},\"quotient_degree\":{},\"library_accept\":{},\
         \"public_values\":{},\"alpha\":{},\"zeta\":{},\"trace_local\":{},\"trace_next\":{},\"quotient_chunks\":{}}}",
        name,
        air_id,
        mp.degree_bits,
        quotient_degree,
        library_accept,
        ja(&pis_u32),
        ja(&ef_u32(&alpha)),
        ja(&ef_u32(&zeta)),
        ef_list(&mp.opened_values.trace_local),
        ef_list(&mp.opened_values.trace_next),
        ef_list2(&mp.opened_values.quotient_chunks),
    )
}

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

    // ---- perm sanity so the harness can tie ground truth to the CONFIRMED permutation ----
    let (_cfg0, 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 (Plonky3's own test vector: pis = [0,1,21]) -> 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");
        cases.push(emit_case("fibonacci_n8", "fibonacci", &perm, pj, &pis, accept));
    }

    // ---- case B: degree-3 multiply AIR, n = 16 -> 2 quotient chunks (exercises the zps product) ----
    {
        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");
        cases.push(emit_case("mul_deg3_n16", "mul_deg3", &perm, pj, &pis, accept));
    }

    let out = format!(
        "{{\"perm_zeros\":{},\"val_generator\":{},\"cases\":[{}]}}",
        ja(&perm_zeros),
        val_generator,
        cases.join(",")
    );
    println!("{}", out);
}