use super::participants::{ParticipantId, Participants};
use crate::crypto_shared::types::PublicKeyExtended;
use crate::errors::{DomainError, Error, InvalidState};
use near_account_id::AccountId;
use near_mpc_contract_interface::types::DomainId;
use near_sdk::{env, near};

pub use mpc_primitives::{AttemptId, EpochId, KeyEventId};

/// The identification of a specific distributed key, based on which a node would know exactly what
/// keyshare it has corresponds to this distributed key. (A distributed key refers to a specific set
/// of keyshares that nodes have which can be pieced together to form the secret key.)
//
// This is the contract-internal storage type, distinct from the DTO
// [`near_mpc_contract_interface::types::KeyForDomain`] used over the wire.
// They are kept separate because the contract stores public keys as
// [`PublicKeyExtended`] (`near_sdk::PublicKey` plus a decompressed Edwards
// point), while the DTO uses the JSON-friendly string/byte form.
#[near(serializers=[borsh, json])]
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct KeyForDomain {
    /// Identifies the domain this key is intended for.
    pub domain_id: DomainId,
    /// Identifies the public key. Although technically redundant given that we have the AttemptId,
    /// we keep it here in the contract so that it can be verified against and queried.
    pub key: PublicKeyExtended,
    /// The attempt ID that generated (initially or as a result of resharing) this distributed key.
    /// Nodes may have made multiple attempts to generate the distributed key, and this uniquely
    /// identifies which one should ultimately be used.
    pub attempt: AttemptId,
}

/// Represents a key for every domain in a specific epoch.
//
// Contract-internal counterpart to [`near_mpc_contract_interface::types::Keyset`].
#[near(serializers=[borsh, json])]
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Keyset {
    pub epoch_id: EpochId,
    pub domains: Vec<KeyForDomain>,
}

impl Keyset {
    pub fn new(epoch_id: EpochId, domains: Vec<KeyForDomain>) -> Self {
        Keyset { epoch_id, domains }
    }

    pub fn public_key(&self, domain_id: DomainId) -> Result<PublicKeyExtended, Error> {
        Ok(self
            .domains
            .iter()
            .find(|k| k.domain_id == domain_id)
            .ok_or(DomainError::NoSuchDomain)?
            .key
            .clone())
    }

    #[cfg(feature = "dev-utils")]
    pub fn get_domain_ids(&self) -> Vec<DomainId> {
        self.domains.iter().map(|domain| domain.domain_id).collect()
    }
}
/// This struct is supposed to contain the participant id associated to the account `env::signer_account_id()`,
/// but is only constructible given a set of participants that includes the signer, thus acting as
/// a type system-based enforcement mechanism (albeit a best-effort one) for authenticating the
/// signer.
#[near(serializers=[borsh, json])]
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct AuthenticatedParticipantId(ParticipantId);
impl AuthenticatedParticipantId {
    pub fn get(&self) -> ParticipantId {
        self.0
    }

    pub fn new(participants: &Participants) -> Result<Self, Error> {
        let signer = env::signer_account_id();
        participants
            .participants()
            .iter()
            .find(|(a_id, _, _)| *a_id == signer)
            .map(|(_, p_id, _)| AuthenticatedParticipantId(*p_id))
            .ok_or_else(|| InvalidState::NotParticipant { account_id: signer }.into())
    }
}

/// This struct contains the account `env::signer_account_id()`, but is only constructible given a
/// set of participants that include the signer, thus acting as a typesystem-based enforcement
/// mechanism (albeit a best-effort one) for authenticating the signer.
#[near(serializers=[borsh, json])]
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AuthenticatedAccountId(AccountId);
impl AuthenticatedAccountId {
    pub fn get(&self) -> &AccountId {
        &self.0
    }

    pub fn new(participants: &Participants) -> Result<Self, Error> {
        let signer = env::signer_account_id();
        if participants
            .participants()
            .iter()
            .any(|(a_id, _, _)| *a_id == signer)
        {
            Ok(AuthenticatedAccountId(signer))
        } else {
            Err(InvalidState::NotParticipant { account_id: signer }.into())
        }
    }
}

#[cfg(test)]
pub mod tests {
    use crate::primitives::{
        key_state::{
            AttemptId, AuthenticatedAccountId, AuthenticatedParticipantId, EpochId, KeyForDomain,
            Keyset,
        },
        test_utils::{bogus_ed25519_public_key_extended, gen_account_id, gen_threshold_params},
    };
    use near_mpc_contract_interface::types::DomainId;
    use near_sdk::{test_utils::VMContextBuilder, testing_env};
    use rand::Rng;

    const MAX_N: usize = 900;

    #[test]
    fn test_epoch_id() {
        let id = rand::thread_rng().r#gen();
        let epoch_id = EpochId::new(id);
        assert_eq!(epoch_id.get(), id);
        assert_eq!(epoch_id.next().get(), id + 1);
    }

    #[test]
    fn test_attempt_id() {
        let attempt_id = AttemptId::new();
        assert_eq!(attempt_id.get(), 0);
        assert_eq!(attempt_id.next().get(), 1);
    }

    #[test]
    fn test_keyset() {
        let domain_id0 = DomainId(0);
        let domain_id1 = DomainId(3);
        let key0 = bogus_ed25519_public_key_extended();
        let key1 = bogus_ed25519_public_key_extended();
        let keyset = Keyset::new(
            EpochId::new(5),
            vec![
                KeyForDomain {
                    domain_id: domain_id0,
                    key: key0.clone(),
                    attempt: AttemptId::new(),
                },
                KeyForDomain {
                    domain_id: domain_id1,
                    key: key1.clone(),
                    attempt: AttemptId::new(),
                },
            ],
        );
        assert_eq!(keyset.public_key(domain_id0).unwrap(), key0);
        assert_eq!(keyset.public_key(domain_id1).unwrap(), key1);
        let _ = keyset.public_key(DomainId(1)).unwrap_err();
    }

    #[test]
    fn test_authenticated_participant_id() {
        let proposed_parameters = gen_threshold_params(MAX_N);
        proposed_parameters
            .validate()
            .expect("Proposed parameters should validate");
        for (account_id, _, _) in proposed_parameters.participants().participants() {
            let mut context = VMContextBuilder::new();
            context.signer_account_id(account_id.clone());
            testing_env!(context.build());
            let _ = AuthenticatedParticipantId::new(proposed_parameters.participants()).unwrap();
            let mut context = VMContextBuilder::new();
            context.signer_account_id(gen_account_id());
            testing_env!(context.build());
            let _ =
                AuthenticatedParticipantId::new(proposed_parameters.participants()).unwrap_err();
        }
    }

    #[test]
    fn test_authenticated_account_id() {
        let proposed_parameters = gen_threshold_params(MAX_N);
        proposed_parameters
            .validate()
            .expect("Proposed parameters should validate");
        for (account_id, _, _) in proposed_parameters.participants().participants() {
            let mut context = VMContextBuilder::new();
            context.signer_account_id(account_id.clone());
            testing_env!(context.build());
            let _ = AuthenticatedAccountId::new(proposed_parameters.participants()).unwrap();
            let mut context = VMContextBuilder::new();
            context.signer_account_id(gen_account_id());
            testing_env!(context.build());
            let _ = AuthenticatedAccountId::new(proposed_parameters.participants()).unwrap_err();
        }
    }
}