use std::iter;

use itertools::Itertools;
use plonky2::{
    hash::hash_types::{HashOut, HashOutTarget},
    iop::witness::{PartialWitness, WitnessWrite},
    plonk::circuit_builder::CircuitBuilder,
};
use serde::{Deserialize, Serialize};

use crate::{
    backends::plonky2::{
        basetypes::D,
        circuits::common::{
            CircuitBuilderPod, PredicateTarget, StatementArgTarget, StatementTarget, ValueTarget,
        },
        error::Result,
        primitives::{
            merkletree::{
                verify_merkle_proof_existence_circuit, MerkleClaimAndProof,
                MerkleProofExistenceTarget,
            },
            signature::{verify_signature_circuit, SignatureVerifyTarget},
        },
        signedpod::SignedPod,
    },
    measure_gates_begin, measure_gates_end,
    middleware::{
        hash_str, Key, NativePredicate, Params, PodType, RawValue, Value, F, KEY_SIGNER, KEY_TYPE,
        SELF,
    },
};

pub fn verify_signed_pod_circuit(
    builder: &mut CircuitBuilder<F, D>,
    signed_pod: &SignedPodVerifyTarget,
) -> Result<()> {
    let params = &signed_pod.params;
    let measure = measure_gates_begin!(builder, "SignedPodVerify");
    // 1. Verify id
    assert_eq!(params.max_signed_pod_values, signed_pod.mt_proofs.len());
    for mt_proof in &signed_pod.mt_proofs {
        verify_merkle_proof_existence_circuit(builder, mt_proof);
        builder.connect_hashes(signed_pod.id, mt_proof.root);
        // mt_proofs.push(mt_proof);
    }

    // 2. Verify type
    let type_mt_proof = &signed_pod.mt_proofs[0];
    let key_type = builder.constant_value(hash_str(KEY_TYPE).into());
    builder.connect_values(type_mt_proof.key, key_type);
    let value_type = builder.constant_value(Value::from(PodType::Signed).raw());
    builder.connect_values(type_mt_proof.value, value_type);

    // 3.a. Verify signature
    verify_signature_circuit(builder, &signed_pod.signature);

    // 3.b. Verify signer (ie. hash(signature.pk) == merkletree.signer_leaf)
    let signer_mt_proof = &signed_pod.mt_proofs[1];
    let key_signer = builder.constant_value(Key::from(KEY_SIGNER).raw());
    let pk_hash = signed_pod.signature.pk.to_value(builder);
    builder.connect_values(signer_mt_proof.key, key_signer);
    builder.connect_values(signer_mt_proof.value, pk_hash);

    // 3.c. connect signed message to pod.id
    builder.connect_values(
        ValueTarget::from_slice(&signed_pod.id.elements),
        signed_pod.signature.msg,
    );

    measure_gates_end!(builder, measure);
    Ok(())
}

#[derive(Clone, Serialize, Deserialize)]
pub struct SignedPodVerifyTarget {
    params: Params,
    id: HashOutTarget,
    // the KEY_TYPE entry must be the first one
    // the KEY_SIGNER entry must be the second one
    mt_proofs: Vec<MerkleProofExistenceTarget>,
    pub(crate) signature: SignatureVerifyTarget,
}

impl SignedPodVerifyTarget {
    pub fn new_virtual(params: &Params, builder: &mut CircuitBuilder<F, D>) -> Self {
        SignedPodVerifyTarget {
            params: params.clone(),
            id: builder.add_virtual_hash(),
            mt_proofs: (0..params.max_signed_pod_values)
                .map(|_| {
                    MerkleProofExistenceTarget::new_virtual(params.max_depth_mt_containers, builder)
                })
                .collect(),
            signature: SignatureVerifyTarget::new_virtual(builder),
        }
    }
    pub fn pub_statements(
        &self,
        builder: &mut CircuitBuilder<F, D>,
        self_id: bool,
    ) -> Vec<StatementTarget> {
        let mut statements = Vec::new();
        let predicate = PredicateTarget::new_native(builder, &self.params, NativePredicate::Equal);
        let pod_id = if self_id {
            builder.constant_value(SELF.0.into())
        } else {
            ValueTarget {
                elements: self.id.elements,
            }
        };
        for mt_proof in &self.mt_proofs {
            let args = [
                StatementArgTarget::anchored_key(builder, &pod_id, &mt_proof.key),
                StatementArgTarget::literal(builder, &mt_proof.value),
            ]
            .into_iter()
            .chain(iter::repeat_with(|| StatementArgTarget::none(builder)))
            .take(self.params.max_statement_args)
            .collect();
            let statement = StatementTarget {
                predicate: predicate.clone(),
                args,
            };
            statements.push(statement);
        }
        statements
    }

    pub fn set_targets(&self, pw: &mut PartialWitness<F>, pod: &SignedPod) -> Result<()> {
        // set the self.mt_proofs witness with the following order:
        // - KEY_TYPE leaf proof
        // - KEY_SIGNER leaf proof
        // - rest of leaves
        // - empty leaves (if needed)

        // add proof verification of KEY_TYPE & KEY_SIGNER leaves
        let key_type_key = Key::from(KEY_TYPE);
        let key_signer_key = Key::from(KEY_SIGNER);
        [&key_type_key, &key_signer_key]
            .iter()
            .enumerate()
            .try_for_each(|(i, k)| {
                let (v, proof) = pod.dict.prove(k)?;
                self.mt_proofs[i].set_targets(
                    pw,
                    true,
                    &MerkleClaimAndProof::new(pod.dict.commitment(), k.raw(), Some(v.raw()), proof),
                )
            })?;

        // add the verification of the rest of leaves
        let mut curr = 2; // since we already added key_type and key_signer
        for (k, v) in pod.dict.kvs().iter().sorted_by_key(|kv| kv.0.hash()) {
            if *k == key_type_key || *k == key_signer_key {
                // skip the key_type & key_signer leaves, since they have
                // already been checked
                continue;
            }

            let (obtained_v, proof) = pod.dict.prove(k)?;
            assert_eq!(obtained_v, v); // sanity check

            self.mt_proofs[curr].set_targets(
                pw,
                true,
                &MerkleClaimAndProof::new(pod.dict.commitment(), k.raw(), Some(v.raw()), proof),
            )?;
            curr += 1;
        }
        // sanity check
        assert!(curr <= self.params.max_signed_pod_values);

        // add the proofs of empty leaves (if needed), till the max_signed_pod_values
        let mut mp = MerkleClaimAndProof::empty();
        mp.root = pod.dict.commitment();
        for i in curr..self.params.max_signed_pod_values {
            self.mt_proofs[i].set_targets(pw, false, &mp)?;
        }

        // get the signer pk
        let pk = pod.signer;
        // the msg signed is the pod.id
        let msg = RawValue::from(pod.id.0);

        // set signature targets values
        self.signature
            .set_targets(pw, true, pk, msg, pod.signature.clone())?;

        // set the id target value
        pw.set_hash_target(self.id, HashOut::from_vec(pod.id.0 .0.to_vec()))?;

        Ok(())
    }
}

#[cfg(test)]
pub mod tests {
    use std::any::Any;

    use plonky2::plonk::{circuit_builder::CircuitBuilder, circuit_data::CircuitConfig};

    use super::*;
    use crate::{
        backends::plonky2::{
            basetypes::C,
            primitives::ec::schnorr::SecretKey,
            signedpod::{SignedPod, Signer},
        },
        middleware::F,
    };

    #[test]
    fn test_signed_pod_verify() -> Result<()> {
        let params = Params {
            max_signed_pod_values: 6,
            ..Default::default()
        };
        // set max_signed_pod_values to 6, and we insert 3 leaves, so that the
        // circuit has enough space for the 3 leaves plus the KEY_TYPE and
        // KEY_SIGNER and one empty leaf.

        // prepare a signedpod
        let mut pod = crate::frontend::SignedPodBuilder::new(&params);
        pod.insert("idNumber", "4242424242");
        pod.insert("dateOfBirth", 1169909384);
        pod.insert("socialSecurityNumber", "G2121210");
        let sk = SecretKey::new_rand();
        let signer = Signer(sk);
        let pod = pod.sign(&signer).unwrap();
        let signed_pod = (pod.pod as Box<dyn Any>).downcast::<SignedPod>().unwrap();

        // use the pod in the circuit
        let config = CircuitConfig::standard_recursion_config();
        let mut builder = CircuitBuilder::<F, D>::new(config);
        let mut pw = PartialWitness::<F>::new();

        // build the circuit logic
        let signed_pod_verify = SignedPodVerifyTarget::new_virtual(&params, &mut builder);
        verify_signed_pod_circuit(&mut builder, &signed_pod_verify)?;

        // set the signed_pod as target values for the circuit
        signed_pod_verify.set_targets(&mut pw, &signed_pod)?;

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

        Ok(())
    }
}