use dusk_plonk::error::Error as PlonkError;
use dusk_poseidon::tree::{self, PoseidonBranch, PoseidonLeaf, PoseidonTree};
use rand::rngs::OsRng;
use rand::{CryptoRng, RngCore};


use dusk_plonk::prelude::*;
use nstack::annotation::Keyed;

// Depth of the merkle tree
const DEPTH: usize = 17;

// Capacity of the circuit
const CAPACITY: usize = 15;

// Alias for the default tree implementation
type Tree = PoseidonTree<DataLeaf, (), DEPTH>;

// Leaf representation
#[derive(Debug, Default, Clone, Copy, PartialOrd, Ord, PartialEq, Eq)]
pub struct DataLeaf {
    data: BlsScalar,
    pos: u64,
}

// Keyed needs to be implemented for a leaf type and the tree key.
impl Keyed<()> for DataLeaf {
    fn key(&self) -> &() {
        &()
    }
}

impl DataLeaf {
    pub fn random<R: RngCore + CryptoRng>(rng: &mut R) -> Self {
        let data = BlsScalar::random(rng);
        let pos = 0;

        Self { data, pos }
    }
}

// Any leaf of the poseidon tree must implement `PoseidonLeaf`
impl PoseidonLeaf for DataLeaf {
    // Cryptographic hash of the data leaf
    fn poseidon_hash(&self) -> BlsScalar {
        self.data
    }

    // Position on the tree
    fn pos(&self) -> &u64 {
        &self.pos
    }

    // Method used to set the position on the tree after the
    // `PoseidonTree::push` call.
    fn set_pos(&mut self, pos: u64) {
        self.pos = pos;
    }
}

#[derive(Default)]
struct MerkleOpeningCircuit {
    branch: PoseidonBranch<DEPTH>,
}

impl MerkleOpeningCircuit {
    /// Generate N random leaves and append them to the tree
    pub fn random<R: RngCore + CryptoRng>(
        rng: &mut R,
        tree: &mut Tree,
    ) -> Self {
        const N: u64 = 1024;

        // Append 1024 elements to the tree
        for _ in 0..N {
            let leaf = DataLeaf::random(rng);
            tree.push(leaf);
        }

        let branch = tree.branch(N - 1).expect(
            "Failed to fetch the branch of the created leaf from the tree",
        );

        Self { branch }
    }
}

impl Circuit for MerkleOpeningCircuit {
    fn circuit<C>(&self, composer: &mut C) -> Result<(), PlonkError> 
    where 
        C: Composer,
    {
      let leaf: BlsScalar = *self.branch;
      let leaf = composer.append_witness(leaf);
  
      let root = self.branch.root();
      let root = composer.append_witness(*root);
  
      let root_p =
              tree::merkle_opening::<C, DEPTH>(composer, &self.branch, leaf);
  
      composer.assert_equal(root_p, root);
  
      Ok(())
    }
}

// Create a prover and a verifier
let label = b"dusk-network";
let pp = PublicParameters::setup(1 << CAPACITY, &mut OsRng).unwrap();

let (prover, verifier) =
    Compiler::compile(&pp, label).expect("failed to compile circuit");

// Instantiate a new tree
let mut tree = Tree::default();
let circuit = MerkleOpeningCircuit::random(&mut OsRng, &mut tree);

// Generate a ZK opening proof
let (proof, public_inputs) = prover.prove(&mut OsRng, &circuit)
            .expect("proving the circuit should succeed");

// Verify the proof
verifier.verify(&proof, &public_inputs)
    .expect("verifying the proof should succeed");