//!  Multi-party key resharing
//!
//! Key resharing process resembles the key generation protocol, however no new random shards are sampled in the beginning. Instead, existing Shamir's shares are reshared with the set of new parties using new randomly generated polynomials.
//! New shares are generated by the set of current shareholders, after which existing shares are destroyed. Eventually old shareholders do not have access to the key anymore.
//!
//! The special case where the old party set is equal to the new party set is called "key refresh". In this case every party just updates own shard with new value.
use crate::protocol::PartyIndex;

use failure::Fail;

pub use super::messages::resharing::{InMsg, Message, OutMsg, Phase1Broadcast};

use crate::ecdsa::messages::SecretShare;
use curv::{BigInt, FE};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeSet, HashMap};
use std::fmt::Debug;

use trace::trace;

/// Contains  zero knowledge proof of Paillier key's correctness
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CorrectKeyProof(Vec<BigInt>);

/// Enumerates errors which can be reported by resharing protocol
#[derive(Debug, Fail)]
#[allow(clippy::large_enum_variant)]
pub enum ResharingError {
    #[fail(display = "resharing: timeout in {}", phase)]
    Timeout { phase: String },
    #[fail(display = "resharing: invalid empty message set {}", desc)]
    EmptyMessageSet { desc: String },
    #[fail(
        display = "invalid decommitment {}, commitment {}, party {}",
        decomm, comm, party
    )]
    InvalidComm {
        comm: String,
        decomm: String,
        party: PartyIndex,
    },
    #[fail(display = "public keys of old committee members are not same")]
    InconsistentPublicKeys,
    #[fail(display = "invalid secret sharing {}, party {}", vss, party)]
    InvalidVSS { vss: String, party: PartyIndex },
    #[fail(display = "received point has wrong X coordinate: {}", x_coord)]
    WrongXCoordinate { x_coord: usize },
    #[fail(display = "unexpected message {:?}, party {}", message_type, party)]
    UnknownMessageType {
        message_type: Message,
        party: PartyIndex,
    },
    #[fail(display = "invalid proof {}, party {}", proof, party)]
    InvalidDlogProof { proof: String, party: PartyIndex },
    #[fail(display = "invalid correct key proof {}, party {}", proof, party)]
    InvalidCorrectKeyProof { proof: String, party: PartyIndex },
    #[fail(display = "missing range proof from {} ", party)]
    RangeProofSetupMissing { party: PartyIndex },
    #[fail(display = "unexpected range proof from {}, proof {:?} ", party, proof)]
    RangeProofSetupUnexpected { proof: String, party: PartyIndex },
    #[fail(
        display = "range proof setup: dlog proof failed , party {}, proof {} ",
        party, proof
    )]
    RangeProofSetupDlogProofFailed { proof: String, party: PartyIndex },
    #[fail(display = "protocol setup error: {}", _0)]
    ProtocolSetupError(String),
    #[fail(display = "{}", _0)]
    GeneralError(String),
}

/// Contains a vector of possible resharing errors
#[derive(Debug)]
pub struct ErrorState {
    pub errors: Vec<ResharingError>,
}

/// the state the machine returns in case of error(s)
impl ErrorState {
    pub fn new(errors: Vec<ResharingError>) -> Self {
        ErrorState { errors }
    }
    pub fn append(self, rhs: ErrorState) -> Self {
        let mut errors = self.errors;
        errors.extend(rhs.errors.into_iter());
        Self { errors }
    }
}

/// Checks whether all expected messages have been received so far from other parties  
fn is_broadcast_input_complete(
    current_msg_set: &[InMsg],
    other_parties: &BTreeSet<PartyIndex>,
) -> bool {
    let senders = current_msg_set.iter().map(|m| m.sender).collect::<Vec<_>>();
    other_parties.iter().all(|p| senders.contains(p))
}

/// Extracts payloads form enum variants of input message into the hash map
#[trace(disable(current_msg_set), res = "{:?}")]
fn to_hash_map_gen<K, V>(current_msg_set: Vec<InMsg>) -> Result<HashMap<K, V>, Vec<ResharingError>>
where
    K: std::cmp::Eq + std::hash::Hash + std::convert::From<PartyIndex> + std::fmt::Debug,
    V: std::fmt::Debug,
    Option<V>: std::convert::From<Message>,
{
    let (converted_messages, errors) =
        current_msg_set
            .iter()
            .fold((vec![], vec![]), |(mut values, mut errors), m| {
                let body: Option<V> = m.body.clone().into();
                match body {
                    Some(b) => values.push((m.sender, b)),
                    None => errors.push(ResharingError::UnknownMessageType {
                        message_type: m.body.clone(),
                        party: m.sender,
                    }),
                };
                (values, errors)
            });

    if errors.is_empty() {
        Ok(converted_messages
            .into_iter()
            .map(|(party, body)| (party.into(), body))
            .collect::<HashMap<K, V>>())
    } else {
        Err(errors)
    }
}

fn map_parties_to_shares(
    parties: &[PartyIndex],
    outgoing_shares: &[FE],
) -> HashMap<PartyIndex, SecretShare> {
    assert_eq!(outgoing_shares.len(), parties.len()); // hence unwrap() safely
    let outgoing_shares = (1..=outgoing_shares.len())
        .zip(outgoing_shares.iter())
        .collect::<Vec<_>>();

    let party_indexes_ascending = parties.iter().collect::<BTreeSet<_>>();

    party_indexes_ascending
        .iter()
        .zip(outgoing_shares.iter())
        .map(|(party, share)| (**party, (share.0, *share.1)))
        .collect::<HashMap<_, _>>()
}

/// Contains the protocol part performed by a member of old committee
pub mod old_member {

    use super::ErrorState;
    use crate::ecdsa::keygen::MultiPartyInfo;
    use crate::ecdsa::messages::resharing::{InMsg, Message, OutMsg, Phase1Broadcast, VSS};

    use crate::ecdsa::resharing::{map_parties_to_shares, ResharingError};
    use crate::protocol::{Address, PartyIndex};
    use crate::state_machine::{State, StateMachineTraits, Transition};
    use crate::Parameters;
    use curv::cryptographic_primitives::hashing::hash_sha256::HSha256;
    use curv::cryptographic_primitives::hashing::traits::Hash;
    use curv::cryptographic_primitives::secret_sharing::feldman_vss::VerifiableSS;
    use curv::elliptic::curves::traits::ECScalar;
    use curv::{BigInt, FE, GE};
    use std::cell::RefCell;
    use std::collections::BTreeSet;
    use std::iter::FromIterator;
    use std::time::Duration;
    use trace::trace;
    use zeroize::Zeroize;

    #[derive(Clone, Debug, super::Serialize, super::Deserialize)]
    pub struct FinalState;

    type OutMsgVec = Vec<OutMsg>;
    pub type MachineResult = Result<FinalState, ErrorState>;

    /// Type definitions
    #[derive(Debug)]
    pub struct KeyResharingTraits;

    impl StateMachineTraits for KeyResharingTraits {
        type InMsg = InMsg;
        type OutMsg = OutMsg;
        type FinalState = FinalState;
        type ErrorState = ErrorState;
    }

    /// Initial phase of the protcol
    ///
    /// * generates new Shamir's shares of existing share to share among members of new committee
    /// * generates Feldman's VSS
    /// * broadcasts commitment to public key and commitment to Feldman's VSS
    /// * collects ACK messages
    #[derive(Debug)]
    pub struct Phase1 {
        new_committee: BTreeSet<PartyIndex>,
        vss_scheme: VerifiableSS,
        outgoing_shares: Vec<FE>,
        vss_comm: BigInt,
        y: GE,
        timeout: Option<Duration>,
    }

    #[trace(pretty, prefix = "Phase1::")]
    impl Phase1 {
        pub fn new(
            multi_party_info: &MultiPartyInfo,
            new_params: &Parameters,
            old_committee: &[PartyIndex],
            new_committee: &[PartyIndex],
            timeout: Option<Duration>,
        ) -> Result<Self, ResharingError> {
            //check if old committee is sized correctly
            if old_committee.len() <= multi_party_info.key_params.threshold() {
                return Err(ResharingError::ProtocolSetupError(
                    "old committee too small".to_string(),
                ));
            }
            if old_committee.len() > multi_party_info.key_params.share_count() {
                return Err(ResharingError::ProtocolSetupError(
                    "old committee too big".to_string(),
                ));
            }

            // check if new committee is sized correctly
            if new_committee.len() <= new_params.threshold() {
                return Err(ResharingError::ProtocolSetupError(
                    "new committee too small".to_string(),
                ));
            }
            if new_committee.len() > new_params.share_count() {
                return Err(ResharingError::ProtocolSetupError(
                    "new committee too big".to_string(),
                ));
            }

            // check if new committee has duplicates
            let new_parties_as_set = BTreeSet::from_iter(new_committee.iter().cloned());
            if new_parties_as_set.len() != new_committee.len() {
                return Err(ResharingError::ProtocolSetupError(
                    "duplicate entries in new committee's list".to_string(),
                ));
            }

            let own_x: FE = ECScalar::from(&BigInt::from(multi_party_info.own_point() as u64));
            let multiplier = multi_party_info
                .party_to_point_map
                .calculate_lagrange_multiplier(old_committee, own_x);
            let w_i = multi_party_info.own_share() * multiplier;

            let (vss_scheme, outgoing_shares) =
                VerifiableSS::share(new_params.threshold, new_params.share_count, &w_i);
            let vss_refs = vss_scheme.commitments.iter().collect::<Vec<_>>();
            let vss_comm = HSha256::create_hash_from_ge(&vss_refs).to_big_int();

            Ok(Phase1 {
                new_committee: new_parties_as_set,
                vss_scheme,
                outgoing_shares,
                vss_comm,
                y: multi_party_info.public_key,
                timeout,
            })
        }
    }

    #[trace(pretty, prefix = "Phase1::")]
    impl State<KeyResharingTraits> for Phase1 {
        fn start(&mut self) -> Option<OutMsgVec> {
            log::info!("Phase1 (old member) starts");

            let output = self
                .new_committee
                .iter()
                .map(|p| OutMsg {
                    recipient: Address::Peer(*p),
                    body: Message::R1(Phase1Broadcast {
                        y: self.y,
                        vss_commitment: self.vss_comm.clone(),
                    }),
                })
                .collect::<Vec<_>>();

            Some(output)
        }

        #[trace(disable(current_msg_set))]
        fn is_message_expected(&self, msg: &InMsg, current_msg_set: &[InMsg]) -> bool {
            matches!(msg.body, Message::Ack if self.new_committee.contains(&msg.sender) && !msg.is_duplicate(current_msg_set))
        }

        #[trace(disable(current_msg_set))]
        fn is_input_complete(&self, current_msg_set: &[InMsg]) -> bool {
            super::is_broadcast_input_complete(current_msg_set, &self.new_committee)
        }

        fn consume(&self, _current_msg_set: Vec<InMsg>) -> Transition<KeyResharingTraits> {
            Transition::NewState(Box::new(Phase2 {
                new_committee: self.new_committee.clone(),
                vss_scheme: self.vss_scheme.clone(),
                outgoing_shares: RefCell::new(self.outgoing_shares.clone()),
                timeout: self.timeout,
            }))
        }

        fn timeout(&self) -> Option<Duration> {
            self.timeout
        }

        fn timeout_outcome(&self, _current_msg_set: Vec<InMsg>) -> MachineResult {
            Err(ErrorState::new(vec![super::ResharingError::Timeout {
                phase: "Old.phase1".to_owned(),
            }]))
        }
    }

    /// Second phase of the protocol
    ///
    /// * Shares new Shamir's secrets and their respective Feldman's VSS with members of new committee
    /// * Collect `FinalAck` messages and exits  
    struct Phase2 {
        new_committee: BTreeSet<PartyIndex>,
        vss_scheme: VerifiableSS,
        outgoing_shares: RefCell<Vec<FE>>,
        timeout: Option<Duration>,
    }

    impl Phase2 {
        fn zeroize_secret_shares(&self) {
            // Simultaneously remove and zeroize elements.
            self.outgoing_shares
                .borrow_mut()
                .drain(..)
                .for_each(|mut share| share.zeroize());
        }
    }

    #[trace(pretty, prefix = "Phase2::")]
    impl State<KeyResharingTraits> for Phase2 {
        fn start(&mut self) -> Option<Vec<OutMsg>> {
            log::debug!("Phase2 (old member) starts");
            let output = map_parties_to_shares(
                &self.new_committee.iter().cloned().collect::<Vec<_>>(),
                &self.outgoing_shares.borrow(),
            );
            Some(
                output
                    .into_iter()
                    .map(|(p, (x, share))| OutMsg {
                        recipient: Address::Peer(p),
                        body: Message::R3(VSS {
                            share: (x, share),
                            vss: self.vss_scheme.clone(),
                        }),
                    })
                    .collect::<Vec<_>>(),
            )
        }

        #[trace(disable(current_msg_set))]
        fn is_message_expected(&self, msg: &InMsg, current_msg_set: &[InMsg]) -> bool {
            matches!(msg.body, Message::FinalAck if self.new_committee.contains(&msg.sender) && !msg.is_duplicate(current_msg_set))
        }

        #[trace(disable(current_msg_set))]
        fn is_input_complete(&self, current_msg_set: &[InMsg]) -> bool {
            super::is_broadcast_input_complete(current_msg_set, &self.new_committee)
        }

        fn consume(&self, _current_msg_set: Vec<InMsg>) -> Transition<KeyResharingTraits> {
            log::info!("Phase2 succeeded");
            self.zeroize_secret_shares();
            Transition::FinalState(Ok(FinalState {}))
        }

        fn timeout(&self) -> Option<Duration> {
            self.timeout
        }

        fn timeout_outcome(&self, _current_msg_set: Vec<InMsg>) -> MachineResult {
            Err(ErrorState::new(vec![super::ResharingError::Timeout {
                phase: "Old.phase2".to_owned(),
            }]))
        }
    }

    /// Helper function which returns true if a given message of the resharing protocol should be forwarded to a member of old committee
    pub fn is_message_to_committee(msg: &Message) -> bool {
        matches!(msg, Message::Ack | Message::FinalAck)
    }
}
/// Contains the protocol part performed by a member of new committee
pub mod new_member {
    use crate::algorithms::nizk_rsa;
    use crate::ecdsa::keygen::{CorrectKeyProof, MultiPartyInfo, Party2PointMap, RangeProofSetups};
    use crate::ecdsa::messages::resharing::{Phase1Broadcast, Phase2Broadcast, VSS};
    use crate::ecdsa::resharing::{
        map_parties_to_shares, to_hash_map_gen, ErrorState, InMsg, Message, OutMsg, ResharingError,
    };
    use crate::ecdsa::{all_mapped_equal, PaillierKeys, PRIME_BIT_LENGTH_IN_PAILLIER_SCHEMA};
    use crate::protocol::{Address, PartyIndex};
    use crate::state_machine::{State, StateMachineTraits, Transition};
    use crate::Parameters;
    use curv::cryptographic_primitives::hashing::hash_sha256::HSha256;
    use curv::cryptographic_primitives::hashing::traits::Hash;

    use curv::elliptic::curves::traits::{ECPoint, ECScalar};
    use curv::{BigInt, FE, GE};

    use paillier::{EncryptionKey, KeyGeneration, Paillier};

    use crate::algorithms::zkp::{ZkpPublicSetup, ZkpSetup};
    use crate::ecdsa::messages::SecretShare;
    use failure::_core::time::Duration;
    use std::collections::{BTreeSet, HashMap};
    use std::iter::FromIterator;
    use trace::trace;

    /// Result of resharing protocol
    #[derive(Clone, Debug, super::Serialize, super::Deserialize)]
    pub struct FinalState {
        pub info: MultiPartyInfo,
    }

    #[doc(hidden)]
    type OutMsgVec = Vec<OutMsg>;

    /// Type definitions
    #[derive(Debug)]
    pub struct KeyResharingTraits;

    impl StateMachineTraits for KeyResharingTraits {
        type InMsg = InMsg;
        type OutMsg = OutMsg;
        type FinalState = FinalState;
        type ErrorState = ErrorState;
    }

    pub type MachineResult = Result<FinalState, ErrorState>;

    /// Starting phase of resharing protocol
    ///
    /// * Sends nothing out
    /// * collects commitments to public key and to Feldman's VSS
    /// * verifies that all public keys are same
    #[derive(Clone, Debug)]
    pub struct Phase1 {
        old_params: Parameters,
        new_params: Parameters,
        old_committee: BTreeSet<PartyIndex>,
        others_from_new_committee: BTreeSet<PartyIndex>,
        own_party_index: PartyIndex,
        range_proof_setup: Option<ZkpSetup>,
        timeout: Option<Duration>,
    }

    #[trace(pretty, prefix = "Phase1::")]
    impl Phase1 {
        pub fn new(
            old_params: &Parameters,
            new_params: &Parameters,
            old_committee: &[PartyIndex],
            new_committee: &[PartyIndex],
            own_party_index: PartyIndex,
            range_proof_setup: Option<ZkpSetup>,
            timeout: Option<Duration>,
        ) -> Result<Self, ResharingError> {
            // check if old committee is sized correctly
            if old_committee.len() <= old_params.threshold() {
                return Err(ResharingError::ProtocolSetupError(
                    "old committee too small".to_string(),
                ));
            }
            if old_committee.len() > old_params.share_count() {
                return Err(ResharingError::ProtocolSetupError(
                    "old committee too big".to_string(),
                ));
            }

            // check if new committee is sized correctly
            if new_committee.len() <= new_params.threshold() {
                return Err(ResharingError::ProtocolSetupError(
                    "new committee too small".to_string(),
                ));
            }
            if new_committee.len() > new_params.share_count() {
                return Err(ResharingError::ProtocolSetupError(
                    "new committee too big".to_string(),
                ));
            }

            // check if new committee has duplicates
            let old_parties_as_set = BTreeSet::from_iter(old_committee.iter().cloned());
            if old_parties_as_set.len() != old_committee.len() {
                return Err(ResharingError::ProtocolSetupError(
                    "duplicate entries in new committee's list".to_string(),
                ));
            }

            // check if new committee has duplicates
            let new_parties_as_set = BTreeSet::from_iter(new_committee.iter().cloned());
            if new_parties_as_set.len() != new_committee.len() {
                return Err(ResharingError::ProtocolSetupError(
                    "duplicate entries in new committee's list".to_string(),
                ));
            }

            if new_parties_as_set.get(&own_party_index).is_none() {
                return Err(ResharingError::ProtocolSetupError(
                    "own party index not in new committee list".to_string(),
                ));
            }

            let mut others_from_new_committee = new_parties_as_set;
            others_from_new_committee.remove(&own_party_index);

            Ok(Phase1 {
                old_params: *old_params,
                new_params: *new_params,
                old_committee: BTreeSet::from_iter(old_committee.iter().cloned()),
                others_from_new_committee,
                own_party_index,
                range_proof_setup,
                timeout,
            })
        }
    }

    #[trace(pretty, prefix = "Phase1::")]
    impl State<KeyResharingTraits> for Phase1 {
        fn start(&mut self) -> Option<OutMsgVec> {
            log::info!("Phase1 (new member) starts");
            None
        }

        #[trace(disable(current_msg_set))]
        fn is_message_expected(&self, msg: &InMsg, current_msg_set: &[InMsg]) -> bool {
            matches!(msg.body, Message::R1(_)) && !msg.is_duplicate(current_msg_set)
        }

        #[trace(disable(current_msg_set))]
        fn is_input_complete(&self, current_msg_set: &[InMsg]) -> bool {
            super::is_broadcast_input_complete(current_msg_set, &self.old_committee)
        }

        fn consume(&self, current_msg_set: Vec<InMsg>) -> Transition<KeyResharingTraits> {
            match to_hash_map_gen::<PartyIndex, Phase1Broadcast>(current_msg_set) {
                Ok(input) => {
                    if input.is_empty() {
                        let error_state = ErrorState::new(vec![ResharingError::EmptyMessageSet {
                            desc: "Phase1, new committee".to_string(),
                        }]);
                        log::error!("Phase1 returns {:?}", error_state);
                        return Transition::FinalState(Err(error_state));
                    }

                    let different_public_keys = !all_mapped_equal(input.iter(), |(_, msg)| {
                        msg.y.bytes_compressed_to_big_int()
                    });
                    if different_public_keys {
                        let error_state =
                            ErrorState::new(vec![ResharingError::InconsistentPublicKeys]);
                        log::error!("Phase1 returns {:?}", error_state);
                        return Transition::FinalState(Err(error_state));
                    }

                    // Actual consuming happens here
                    let (ek, dk) = Paillier::keypair_with_modulus_size(
                        2 * PRIME_BIT_LENGTH_IN_PAILLIER_SCHEMA,
                    )
                    .keys();
                    let y = input.iter().next().map(|(_, msg)| msg.y).unwrap();
                    let vss_comms = input
                        .into_iter()
                        .map(|(p, m)| (p, m.vss_commitment))
                        .collect::<HashMap<_, _>>();

                    Transition::NewState(Box::new(Phase2 {
                        previous_phase: (*self).clone(),
                        y,
                        vss_comms,
                        my_paillier_keys: PaillierKeys { dk, ek },
                    }))
                }
                Err(e) => {
                    let error_state = ErrorState::new(e);
                    log::error!("Phase 1 returns {:?}", error_state);
                    Transition::FinalState(Err(error_state))
                }
            }
        }

        fn timeout_outcome(&self, _current_msg_set: Vec<InMsg>) -> MachineResult {
            Err(ErrorState::new(vec![super::ResharingError::Timeout {
                phase: "Old.phase1".to_owned(),
            }]))
        }

        fn timeout(&self) -> Option<Duration> {
            self.timeout
        }
    }

    /// Second phase of the resharing protocol
    ///
    /// * Broadcasts public Paillier key, ZK proof of its correctness and optional `RangeProof` setup to other members of new committee
    /// * Collects and verifies same items from other parties
    #[derive(Clone)]
    pub struct Phase2 {
        previous_phase: Phase1,
        y: GE,
        vss_comms: HashMap<PartyIndex, BigInt>,
        my_paillier_keys: PaillierKeys,
    }

    #[trace(pretty, prefix = "Phase2::")]
    impl Phase2 {
        fn verify_range_proof_setups(
            &self,
            input: &HashMap<PartyIndex, Phase2Broadcast>,
        ) -> Result<Option<RangeProofSetups>, Vec<ResharingError>> {
            let my_range_proof_setup = &self.previous_phase.range_proof_setup;
            let verification_errors = input
                .iter()
                .filter_map(
                    |(&p, m)| match (my_range_proof_setup, &m.range_proof_setup) {
                        (Some(_), None) => {
                            Some(ResharingError::RangeProofSetupMissing { party: p })
                        }
                        (None, Some(s)) => Some(ResharingError::RangeProofSetupUnexpected {
                            party: p,
                            proof: format!("{:?}", s.dlog_proof),
                        }),
                        (Some(_), Some(setup)) => {
                            if setup.verify().is_err() {
                                Some(ResharingError::RangeProofSetupDlogProofFailed {
                                    proof: format!("{:?}", setup.dlog_proof),
                                    party: p,
                                })
                            } else {
                                None
                            }
                        }
                        _ => None,
                    },
                )
                .collect::<Vec<_>>();

            if verification_errors.is_empty() {
                // either local setup exists and all parties shared their setups
                // or local setup does not exist and all parties did not share theirs setups
                Ok(my_range_proof_setup.as_ref().map(|s| RangeProofSetups {
                    my_setup: s.clone(),
                    party_setups: input
                        .iter()
                        .map(|(&p, m)| {
                            (
                                p,
                                m.range_proof_setup
                                    .as_ref()
                                    .expect("range proof setup can't be None here")
                                    .clone(),
                            )
                        })
                        .collect::<HashMap<_, _>>(),
                }))
            } else {
                Err(verification_errors)
            }
        }
    }

    #[trace(pretty, prefix = "Phase2::")]
    impl State<KeyResharingTraits> for Phase2 {
        fn start(&mut self) -> Option<OutMsgVec> {
            log::debug!("Phase2 (new member) starts");
            let range_proof_setup = self
                .previous_phase
                .range_proof_setup
                .as_ref()
                .map(|s| ZkpPublicSetup::from_private_zkp_setup(s));
            let proof = nizk_rsa::gen_proof(self.my_paillier_keys.dk.clone());
            #[allow(clippy::if_not_else)]
            let output = self
                .previous_phase
                .others_from_new_committee
                .iter()
                .map(|p| OutMsg {
                    recipient: Address::Peer(*p),
                    body: Message::R2(Phase2Broadcast {
                        ek: self.my_paillier_keys.ek.clone(),
                        correct_key_proof: CorrectKeyProof(proof.clone()),
                        range_proof_setup: range_proof_setup.clone(),
                    }),
                })
                .collect();
            Some(output)
        }

        #[trace(disable(current_msg_set))]
        fn is_message_expected(&self, msg: &InMsg, current_msg_set: &[InMsg]) -> bool {
            matches!(msg.body,
            Message::R2(_)
            if self.previous_phase.others_from_new_committee.contains(&msg.sender)
            && !msg.is_duplicate(current_msg_set))
        }

        #[trace(disable(current_msg_set))]
        fn is_input_complete(&self, current_msg_set: &[InMsg]) -> bool {
            super::is_broadcast_input_complete(
                current_msg_set,
                &self.previous_phase.others_from_new_committee,
            )
        }

        fn consume(&self, current_msg_set: Vec<InMsg>) -> Transition<KeyResharingTraits> {
            match to_hash_map_gen::<PartyIndex, Phase2Broadcast>(current_msg_set) {
                Ok(input) => {
                    let mut errors = input
                        .iter()
                        .filter_map(|(party, msg)| {
                            if nizk_rsa::verify(&msg.ek, &msg.correct_key_proof.0).is_err() {
                                Some(ResharingError::InvalidCorrectKeyProof {
                                    proof: format!("{:?}", msg.correct_key_proof),
                                    party: *party,
                                })
                            } else {
                                None
                            }
                        })
                        .collect::<Vec<_>>();

                    let range_proof_setups = match self.verify_range_proof_setups(&input) {
                        Err(ve) => {
                            errors.extend(ve);
                            None
                        }
                        Ok(setups) => setups,
                    };

                    let mut other_paillier_keys = input
                        .into_iter()
                        .map(|(p, msg)| (p, msg.ek))
                        .collect::<HashMap<_, _>>();
                    other_paillier_keys.insert(
                        self.previous_phase.own_party_index,
                        self.my_paillier_keys.ek.clone(),
                    );

                    if !errors.is_empty() {
                        let error_state = ErrorState::new(errors);
                        log::error!("Phase 2 returns {:?}", error_state);
                        return Transition::FinalState(Err(error_state));
                    }
                    Transition::NewState(Box::new(Phase3 {
                        previous_phase: (*self).clone(),
                        other_paillier_keys,
                        range_proof_setups,
                    }))
                }
                Err(e) => {
                    let error_state = ErrorState::new(e);
                    log::error!("Phase 2 return {:?}", error_state);
                    Transition::FinalState(Err(error_state))
                }
            }
        }

        fn timeout(&self) -> Option<Duration> {
            self.previous_phase.timeout
        }

        fn timeout_outcome(&self, _current_msg_set: Vec<InMsg>) -> MachineResult {
            Err(ErrorState::new(vec![super::ResharingError::Timeout {
                phase: "New.phase2".to_owned(),
            }]))
        }
    }

    /// Third phase of the protocol
    ///
    /// * sends ACK to old members
    /// * collects new Shamir's secrets and FVSS from them
    /// * verifies FVSS
    #[derive(Clone)]
    struct Phase3 {
        previous_phase: Phase2,
        other_paillier_keys: HashMap<PartyIndex, EncryptionKey>,
        range_proof_setups: Option<RangeProofSetups>,
    }

    #[trace(pretty, prefix = "Phase3::")]
    impl State<KeyResharingTraits> for Phase3 {
        fn start(&mut self) -> Option<Vec<OutMsg>> {
            log::debug!("Phase3 (new member) starts");
            Some(
                self.previous_phase
                    .previous_phase
                    .old_committee
                    .iter()
                    .map(|p| OutMsg {
                        recipient: Address::Peer(*p),
                        body: Message::Ack,
                    })
                    .collect::<Vec<_>>(),
            )
        }

        #[trace(disable(current_msg_set))]
        fn is_message_expected(&self, msg: &InMsg, current_msg_set: &[InMsg]) -> bool {
            matches!(msg.body,
            Message::R3(_)
            if self.previous_phase.previous_phase.old_committee.contains(&msg.sender)
            && !msg.is_duplicate(current_msg_set))
        }

        #[trace(disable(current_msg_set))]
        fn is_input_complete(&self, current_msg_set: &[InMsg]) -> bool {
            super::is_broadcast_input_complete(
                current_msg_set,
                &self.previous_phase.previous_phase.old_committee,
            )
        }

        fn consume(&self, current_msg_set: Vec<InMsg>) -> Transition<KeyResharingTraits> {
            match to_hash_map_gen::<PartyIndex, VSS>(current_msg_set) {
                Err(e) => {
                    let error_state = ErrorState::new(e);
                    log::error!("Phase 3 returns {:?}", error_state);
                    Transition::FinalState(Err(error_state))
                }

                Ok(input) => {
                    // check input for emptiness
                    if input.is_empty() {
                        let error_state = ErrorState::new(vec![ResharingError::EmptyMessageSet {
                            desc: "Phase3, new committee".to_string(),
                        }]);
                        log::error!("Phase 3 returns {:?}", error_state);
                        return Transition::FinalState(Err(error_state));
                    }
                    // check if all x coordinates are same
                    if !all_mapped_equal(input.iter(), |(_, vss)| vss.share.0) {
                        let error_state = ErrorState::new(vec![ResharingError::GeneralError(
                            "Shares have different x coordinate".to_string(),
                        )]);
                        log::error!("Phase 3 returns {:?}", error_state);
                        return Transition::FinalState(Err(error_state));
                    }

                    let (my_x, my_share) = (
                        input.iter().next().map(|(_, vss)| vss.share.0).unwrap(),
                        input
                            .iter()
                            .fold(FE::zero(), |acc, (_, vss)| acc + vss.share.1),
                    );

                    // check commitment errors
                    let commitment_errors = input
                        .into_iter()
                        .filter_map(|(p, vss)| {
                            let ((_, x_i), vss) = (vss.share, vss.vss);
                            let vss_refs = vss.commitments.iter().collect::<Vec<_>>();
                            let decomm = HSha256::create_hash_from_ge(&vss_refs).to_big_int();
                            match self.previous_phase.vss_comms.get(&p) {
                                Some(comm) => {
                                    if *comm == decomm {
                                        match vss.validate_share(&x_i, my_x) {
                                            Ok(()) => None,
                                            Err(_) => Some(ResharingError::InvalidVSS {
                                                vss: format!("{:?}", vss),
                                                party: p,
                                            }),
                                        }
                                    } else {
                                        Some(ResharingError::InvalidComm {
                                            comm: format!("{:?}", comm),
                                            decomm: format!("{:?}", decomm),
                                            party: p,
                                        })
                                    }
                                }
                                None => Some(ResharingError::GeneralError(format!(
                                    "commitment not found for party {}",
                                    p
                                ))),
                            }
                        })
                        .collect::<Vec<_>>();

                    if !commitment_errors.is_empty() {
                        let error_state = ErrorState::new(commitment_errors);
                        log::error!("Phase3 returns {:?}", error_state);
                        return Transition::FinalState(Err(error_state));
                    }
                    Transition::NewState(Box::new(Phase4 {
                        previous_phase: (*self).clone(),
                        share: (my_x, my_share),
                    }))
                }
            }
        }

        fn timeout_outcome(&self, _current_msg_set: Vec<InMsg>) -> MachineResult {
            Err(ErrorState::new(vec![super::ResharingError::Timeout {
                phase: "Old.phase3".to_owned(),
            }]))
        }

        fn timeout(&self) -> Option<Duration> {
            self.previous_phase.previous_phase.timeout
        }
    }

    /// Last phase of the protocol
    ///
    /// * sends `FinalAck` messages to all parties, including members of old and new committees
    /// * collects `FinalAck` from membeers of new committee and exits  
    struct Phase4 {
        previous_phase: Phase3,
        share: SecretShare,
    }

    #[trace(pretty, prefix = "Phase4::")]
    impl State<KeyResharingTraits> for Phase4 {
        fn start(&mut self) -> Option<Vec<OutMsg>> {
            log::debug!("Phase4 (new member) starts");
            let self_setup = &self.previous_phase.previous_phase.previous_phase;
            Some(
                #[allow(clippy::filter_map)]
                self_setup
                    .old_committee
                    .iter()
                    .map(|p| OutMsg {
                        recipient: Address::Peer(*p),
                        body: Message::FinalAck,
                    })
                    .chain(self_setup.others_from_new_committee.iter().map(|p| OutMsg {
                        recipient: Address::Peer(*p),
                        body: Message::FinalAck,
                    }))
                    .collect::<Vec<_>>(),
            )
        }

        #[trace(disable(current_msg_set))]
        fn is_message_expected(&self, msg: &InMsg, current_msg_set: &[InMsg]) -> bool {
            matches!(msg.body, Message::FinalAck
            if self.previous_phase.previous_phase.previous_phase.others_from_new_committee.contains(&msg.sender)
            && !msg.is_duplicate(current_msg_set))
        }

        #[trace(disable(current_msg_set))]
        fn is_input_complete(&self, current_msg_set: &[InMsg]) -> bool {
            super::is_broadcast_input_complete(
                current_msg_set,
                &self
                    .previous_phase
                    .previous_phase
                    .previous_phase
                    .others_from_new_committee,
            )
        }

        fn consume(&self, _current_msg_set: Vec<InMsg>) -> Transition<KeyResharingTraits> {
            let mut new_committee = self
                .previous_phase
                .previous_phase
                .previous_phase
                .others_from_new_committee
                .iter()
                .cloned()
                .collect::<Vec<_>>();

            new_committee.push(
                self.previous_phase
                    .previous_phase
                    .previous_phase
                    .own_party_index,
            );

            let party_mapping_to_points = Party2PointMap {
                // using dummy shares  as we need x-coords only
                points: map_parties_to_shares(
                    new_committee.as_slice(),
                    &vec![FE::zero(); new_committee.len()],
                )
                .into_iter()
                .map(|(party, (point, _))| (party, point))
                .collect::<HashMap<_, _>>(),
            };
            log::info!("Phase4 succeeded");
            Transition::FinalState(Ok(FinalState {
                info: MultiPartyInfo {
                    key_params: self.previous_phase.previous_phase.previous_phase.new_params,
                    own_party_index: self
                        .previous_phase
                        .previous_phase
                        .previous_phase
                        .own_party_index,
                    secret_share: self.share,
                    public_key: self.previous_phase.previous_phase.y,
                    own_he_keys: self.previous_phase.previous_phase.my_paillier_keys.clone(),
                    party_he_keys: self.previous_phase.other_paillier_keys.clone(),
                    party_to_point_map: party_mapping_to_points,
                    range_proof_setups: self.previous_phase.range_proof_setups.clone(),
                },
            }))
        }

        fn timeout(&self) -> Option<Duration> {
            self.previous_phase.previous_phase.previous_phase.timeout
        }

        fn timeout_outcome(&self, _current_msg_set: Vec<InMsg>) -> MachineResult {
            Err(ErrorState::new(vec![super::ResharingError::Timeout {
                phase: "New.phase4".to_owned(),
            }]))
        }
    }
    /// Helper function which returns true if a given message of the resharing protocol should be forwarded to a member of new committee
    pub fn is_message_to_committee(msg: &Message) -> bool {
        !matches!(msg, Message::Ack)
    }
}

#[cfg(test)]
mod tests {
    use crate::algorithms::zkp::ZkpSetup;
    use crate::ecdsa::keygen::MultiPartyInfo;
    use crate::ecdsa::messages::SecretShare;
    use crate::ecdsa::resharing::new_member::KeyResharingTraits;
    use crate::ecdsa::resharing::old_member::KeyResharingTraits as OldKeyResharingTraits;
    use crate::ecdsa::resharing::{InMsg, OutMsg};
    use crate::protocol::{Address, InputMessage, PartyIndex};
    use crate::state_machine::sync_channels::StateMachine;
    use crate::Parameters;
    use crossbeam_channel::{Receiver, Sender};
    use curv::cryptographic_primitives::secret_sharing::feldman_vss::VerifiableSS;
    use curv::elliptic::curves::traits::{ECPoint, ECScalar};
    use curv::{BigInt, FE, GE};
    use failure::bail;
    use std::path::Path;
    use std::{fs, thread};

    struct Node {
        party: PartyIndex,
        egress: Receiver<OutMsg>,
        ingress: Sender<InMsg>,
    }

    struct OutputMessageWithSource {
        msg: OutMsg,
        source: PartyIndex,
    }

    #[test]
    fn resharing() -> Result<(), failure::Error> {
        let _ = env_logger::builder().is_test(true).try_init();
        sharing_helper(false)
    }

    #[test]
    fn resharing_with_range_proofs() -> Result<(), failure::Error> {
        let _ = env_logger::builder().is_test(true).try_init();
        sharing_helper(true)
    }

    pub fn sharing_helper(use_range_proofs: bool) -> Result<(), failure::Error> {
        let _ = env_logger::builder().is_test(true).try_init();
        let old_params = Parameters {
            share_count: 3,
            threshold: 1,
        };

        let new_params = Parameters {
            share_count: 4,
            threshold: 1,
        };

        let mut old_nodes = Vec::new();
        let mut new_nodes = Vec::new();

        let old_committee: Vec<usize> = vec![0, 1, 2];
        let new_committee: Vec<usize> = vec![0, 1, 2, 3];

        let mut old_handles = Vec::new();
        let mut new_handles = Vec::new();

        // start new committee first
        for i in new_committee.clone() {
            let (ingress, rx) = crossbeam_channel::unbounded();
            let (tx, egress) = crossbeam_channel::unbounded();
            log::info!("starting party new_{}", i);
            let nc_clone = new_committee
                .iter()
                .map(|i| (*i).into())
                .collect::<Vec<PartyIndex>>();
            let oc_clone = old_committee
                .iter()
                .map(|i| (*i).into())
                .collect::<Vec<PartyIndex>>();

            let range_proof_setup = if use_range_proofs {
                // the setup from the bank of pre-generated ones
                let path = Path::new("tests/data/rp-setups.json");
                let zkp_setups: Vec<ZkpSetup> = serde_json::from_str(&fs::read_to_string(path)?)?;
                Some(zkp_setups[i].clone())
            } else {
                None
            };

            assert!(!use_range_proofs || range_proof_setup.is_some());

            let join_handle = thread::spawn(move || {
                let start_state = Box::new(super::new_member::Phase1::new(
                    &old_params,
                    &new_params,
                    &oc_clone,
                    &nc_clone,
                    i.into(),
                    range_proof_setup,
                    None,
                )?);

                let mut new_member_machine =
                    StateMachine::<KeyResharingTraits>::new(start_state, &rx, &tx);

                match new_member_machine.execute() {
                    Some(Ok(fs)) => {
                        log::trace!("new_{} success", i);
                        Ok(fs)
                    }
                    Some(Err(e)) => {
                        bail!("new_{} error {:?}", i, e);
                    }
                    None => {
                        bail!("new_{} error in the machine", i);
                    }
                }
            });

            new_nodes.push(Node {
                party: i.into(),
                egress,
                ingress,
            });
            new_handles.push(join_handle);
        }

        // and the start old committee
        for i in old_committee.clone() {
            let (ingress, rx) = crossbeam_channel::unbounded();
            let (tx, egress) = crossbeam_channel::unbounded();

            log::info!("starting party old_{}", i);
            let path = if use_range_proofs {
                format!("tests/data/zkrp-keys.{}.json", i)
            } else {
                format!("tests/data/keys.{}.json", i)
            };
            let path = Path::new(&path);
            let multi_party_shared_info: MultiPartyInfo =
                serde_json::from_str(&fs::read_to_string(path)?)?;

            assert!(!use_range_proofs || multi_party_shared_info.range_proof_setups.is_some());

            let oc_clone = old_committee
                .iter()
                .map(|i| (*i).into())
                .collect::<Vec<PartyIndex>>();

            let nc_clone = new_committee
                .iter()
                .map(|i| (*i).into())
                .collect::<Vec<PartyIndex>>();

            let join_handle = thread::spawn(move || {
                let start_state = Box::new(super::old_member::Phase1::new(
                    &multi_party_shared_info,
                    &new_params,
                    &oc_clone,
                    &nc_clone,
                    None,
                )?);

                let mut old_member_machine =
                    StateMachine::<OldKeyResharingTraits>::new(start_state, &rx, &tx);
                match old_member_machine.execute() {
                    Some(Ok(fs)) => {
                        log::trace!("old_{} success", i);
                        Ok(fs)
                    }
                    Some(Err(e)) => {
                        bail!("old_{} error {:?}", i, e);
                    }
                    None => {
                        bail!("old_{} error in the machine", i);
                    }
                }
            });

            old_nodes.push(Node {
                party: i.into(),
                egress,
                ingress,
            });
            old_handles.push(join_handle);
        }

        //run the multiplexor
        let _mx_thread = thread::spawn(move || {
            loop {
                let mut messages_to_new_nodes = Vec::new();
                let mut messages_to_old_nodes = Vec::new();

                // collect output from old odes
                for node in old_nodes.iter() {
                    if let Ok(out_msg) = node.egress.try_recv() {
                        if super::new_member::is_message_to_committee(&out_msg.body) {
                            messages_to_new_nodes.push(OutputMessageWithSource {
                                msg: out_msg,
                                source: node.party,
                            });
                        }
                    }
                }

                for node in new_nodes.iter() {
                    if let Ok(out_msg) = node.egress.try_recv() {
                        if super::old_member::is_message_to_committee(&out_msg.body) {
                            messages_to_old_nodes.push(OutputMessageWithSource {
                                msg: out_msg.clone(),
                                source: node.party,
                            });
                        }
                        if super::new_member::is_message_to_committee(&out_msg.body) {
                            messages_to_new_nodes.push(OutputMessageWithSource {
                                msg: out_msg.clone(),
                                source: node.party,
                            });
                        }
                    }
                }

                // forward collected messages
                messages_to_old_nodes
                    .iter()
                    .for_each(|mm| match &mm.msg.recipient {
                        Address::Broadcast => {
                            log::error!("the algorithm should not use broadcast");
                        }
                        Address::Peer(peer) => {
                            if let Some(node) = old_nodes.iter().find(|node| (*node).party == *peer)
                            {
                                match node.ingress.send(InputMessage {
                                    sender: mm.source,
                                    body: mm.msg.body.clone(),
                                }) {
                                    Ok(()) => {}
                                    Err(_) => log::warn!(
                                        "error sending msg to old node {}, msg: {}",
                                        peer,
                                        mm.msg.body
                                    ),
                                }
                            }
                        }
                    });

                messages_to_new_nodes
                    .iter()
                    .for_each(|mm| match &mm.msg.recipient {
                        Address::Broadcast => {
                            log::error!("the algorithm should not use broadcast");
                        }
                        Address::Peer(peer) => {
                            if let Some(node) = new_nodes.iter().find(|node| (*node).party == *peer)
                            {
                                match node.ingress.send(InputMessage {
                                    sender: mm.source,
                                    body: mm.msg.body.clone(),
                                }) {
                                    Ok(()) => {}
                                    Err(_) => log::warn!(
                                        "error sending msg to new node {}, msg:  {}",
                                        peer,
                                        mm.msg.body
                                    ),
                                };
                            }
                        }
                    });
            }
        });

        let old_committee_result = old_handles
            .into_iter()
            .map(|h| h.join())
            .collect::<Vec<_>>();

        let new_committee_result = new_handles
            .into_iter()
            .map(|h| h.join())
            .collect::<Vec<_>>();

        if old_committee_result
            .iter()
            .any(|r| r.is_err() || r.as_ref().unwrap().is_err())
        {
            old_committee_result.iter().for_each(|r| match r {
                Ok(result) => match result {
                    Ok(final_state) => log::error!("{:?}", final_state),
                    Err(e) => log::error!("{:?}", e),
                },
                Err(e) => log::error!("{:?}", e),
            });
            assert!(false, "Some state machines returned error");
        }

        if new_committee_result
            .iter()
            .any(|r| r.is_err() || r.as_ref().unwrap().is_err())
        {
            new_committee_result.iter().for_each(|r| match r {
                Ok(result) => match result {
                    Ok(final_state) => log::error!("{:?}", final_state),
                    Err(e) => log::error!("{:?}", e),
                },
                Err(e) => log::error!("{:?}", e),
            });
            assert!(false, "Some state machines returned error");
        }

        let new_final_states = new_committee_result
            .into_iter()
            .map(|x| x.unwrap().unwrap())
            .collect::<Vec<_>>();

        // reconstruct the secret
        let secret_shares = new_final_states
            .iter()
            .map(|fs| fs.info.secret_share)
            .collect::<Vec<_>>();

        let reconstructed_private_key = reconstruct(&secret_shares);
        let old_public_key = new_final_states
            .iter()
            .map(|s| s.info.public_key)
            .collect::<Vec<_>>()[0];
        let g: GE = ECPoint::generator();
        let new_public_key = g * reconstructed_private_key;

        assert_eq!(
            new_public_key.get_element(),
            old_public_key.get_element(),
            "new private key key does not match odl public key"
        );

        Ok(())
    }

    pub fn reconstruct(secret_shares: &[SecretShare]) -> FE {
        //assert!(shares.len() >= self.reconstruct_limit());
        let (points, shares): (Vec<FE>, Vec<FE>) = secret_shares
            .iter()
            .map(|(x, y)| {
                let index_bn = BigInt::from(*x as u64);
                let x_coord: FE = ECScalar::from(&index_bn);
                (x_coord, y)
            })
            .unzip();
        VerifiableSS::lagrange_interpolation_at_zero(&points, &shares)
    }
}
