/// Plonky2 circuit to verify the fibonacci-starky proofs
use anyhow::Result;
use plonky2::iop::{target::Target, witness::PartialWitness};
use plonky2::plonk::circuit_builder::CircuitBuilder;
use starky::proof::{StarkProofWithPublicInputs, StarkProofWithPublicInputsTarget};
use starky::recursive_verifier::{
    add_virtual_stark_proof_with_pis, set_stark_proof_with_pis_target, verify_stark_proof_circuit,
};
use starky::stark::Stark;
use std::marker::PhantomData;

use crate::{C, D, F, STARK_CONFIG};

/// Plonky2 circuit to verify a Starky proof of fibstarky
pub struct FibStarkVerifierTargets<S: Stark<F, D> + Copy> {
    _s: PhantomData<S>,

    proof_target: StarkProofWithPublicInputsTarget<D>,
    zero_target: Target,
}
impl<S: Stark<F, D> + Copy> FibStarkVerifierTargets<S> {
    pub fn add_targets(
        builder: &mut CircuitBuilder<F, D>,
        stark: S,
        degree_bits: usize,
    ) -> Result<Self> {
        let proof_target =
            add_virtual_stark_proof_with_pis(builder, &stark, &STARK_CONFIG, degree_bits, 0, 0);

        verify_stark_proof_circuit::<F, C, S, D>(
            builder,
            stark,
            proof_target.clone(),
            &STARK_CONFIG,
            None,
        );
        let zero_target = builder.zero();

        Ok(Self {
            _s: PhantomData,
            proof_target,
            zero_target,
        })
    }
    pub fn set_targets(
        &self,
        pw: &mut PartialWitness<F>,
        proof: &StarkProofWithPublicInputs<F, C, D>,
        degree_bits: usize,
    ) -> Result<()> {
        set_stark_proof_with_pis_target(
            pw,
            &self.proof_target,
            &proof,
            degree_bits,
            self.zero_target,
        )?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use plonky2::field::types::Field;
    use plonky2::plonk::circuit_data::CircuitConfig;
    use plonky2::util::timing::TimingTree;
    use starky::prover::prove;
    use starky::verifier::verify_stark_proof;

    use super::*;
    use crate::fibstark::FibonacciStark;

    fn set_log() {
        let _ = env_logger::builder()
            .filter_level(log::LevelFilter::Debug)
            .is_test(true)
            .try_init();
    }

    #[test]
    fn test_stark_verif_in_plonky2() -> Result<()> {
        set_log();

        // 1) first generate the fibstark starky proof
        let num_rows = 1 << 10;
        let x0 = F::from_canonical_u32(2);
        let x1 = F::from_canonical_u32(7);

        let stark = FibonacciStark::<F, D>::new(num_rows);
        let public_inputs = [x0, x1, stark.compute_native(x0, x1)];
        let trace = stark.generate_trace(public_inputs[0], public_inputs[1]);

        let stark_proof = prove::<F, C, FibonacciStark<F, D>, D>(
            stark,
            &STARK_CONFIG,
            trace,
            &public_inputs,
            None,
            &mut TimingTree::default(),
        )
        .expect("We should have a valid proof!");

        verify_stark_proof(stark, stark_proof.clone(), &STARK_CONFIG, None)
            .expect("We should be able to verify this proof!");

        // 2) now build a plonky2 circuit that verifies the starky proof
        let degree_bits = stark_proof.proof.recover_degree_bits(&STARK_CONFIG);

        // construct the circuit
        let circuit_config = CircuitConfig::standard_recursion_config();
        let mut builder = CircuitBuilder::<F, D>::new(circuit_config);
        let plonky_targets =
            FibStarkVerifierTargets::add_targets(&mut builder, stark, degree_bits)?;

        builder.print_gate_counts(0);

        // set concrete values to targets
        let mut pw = PartialWitness::new();
        plonky_targets.set_targets(&mut pw, &stark_proof, degree_bits)?;

        // generate the plonky2 proof
        let data = builder.build::<C>();
        let proof = data.prove(pw)?;

        // 3) verify the plonky2 proof
        data.verify(proof)?;

        Ok(())
    }
}