use frost_core::{Ciphersuite, Group, Identifier};
use rand_chacha::rand_core::{CryptoRng, RngCore};
use sp_std::{
    collections::btree_map::BTreeMap,
    marker::PhantomData,
    vec::Vec,
};

pub trait IdentifierConverter<A, E> {
    fn non_zero_index(index: A) -> Result<A, E>;
    fn convert_identifier_to_index(bytes: &[u8]) -> Result<A, E>;
}

pub trait FlexibleRoundOptimizedSchnorrThresholdSignature<A, E>: MerkleTreeBuilder<A, E> {
    type Packages<'a>;
    type NoncePackages<'a>;

    fn header_bytes_len(&self) -> usize;
    fn element_bytes_len(&self) -> usize;
    fn scalar_bytes_len(&self) -> usize;

    fn verify_signature_share(
        &self,
        i: A,
        is: impl Iterator<Item = A>,
        ss: &[u8],
        bf: &[u8],
        nh: &[u8],
        bh: &[u8],
        gc: &[u8],
        vs: &[u8],
        vk: &[u8],
        m: &[u8],
    ) -> Result<(), E>;
    fn is_signature_valid(&self, vk: &[u8], r: &[u8], z: &[u8], m: &[u8]) -> Result<(), E>;
    fn accumulate_signature_scalar(&self, ss: &[u8], sh: &[u8]) -> Result<Vec<u8>, E>;

    fn generate_nonce_commitment<R>(&self, sk: &[u8], r: R) -> Result<(Vec<u8>, Vec<u8>, Vec<u8>), E>
    where
        R: CryptoRng + RngCore;

    fn generate_group_commitment(
        &self,
        a: A,
        sc: &Self::NoncePackages<'_>,
        vk: &[u8],
        m: &[u8],
    ) -> Result<(Vec<u8>, Vec<u8>, Vec<Self::Hash>, Self::Hash), E>;
    fn init_signing_package(&self, n: &Self::NoncePackages<'_>, m: &[u8]) -> Result<Vec<u8>, E>;
    fn partial_sign_message(&self, sp: &[u8], sn: &[u8], ss: &[u8]) -> Result<Vec<u8>, E>;
}

pub trait EllipticCurveDiffieHellman<A, C, D, K, E> {
    fn ecdh_encrypt_package<R>(&self, c: &C, aad: &[u8], m: &[u8], r: R) -> Result<D, E>
    where
        R: RngCore + CryptoRng;

    fn ecdh_decrypt_package(&self, c: &C, d: &D, aad: &[u8]) -> Result<Vec<u8>, E>;
    fn ecdh_prepare_additional_info(&self, si: A, ri: A, di: K, _m: PhantomData<D>) -> Vec<u8>;
    fn ecdh_get_cipher_from_keys(&self, rp: &[u8], sp: &[u8], i: &[u8], _m: PhantomData<D>) -> Result<C, E>;
}

pub trait DistributedKeyGeneration<A, E>: MerkleTreeBuilder<A, E> {
    type RoundPackages;
    type Part1Result;
    type Part2Result;
    type Part3Result;

    fn dkg_narrow_round_max_signers(&self, s: &[u8], m: A, r: u8) -> Result<Vec<u8>, E>;
    fn dkg_verify_proof_of_knowledge(&self, i: A, r1: &[u8]) -> Result<A, E>;
    fn dkg_part1<R>(&self, i: A, max: A, min: A, rng: R) -> Self::Part1Result
    where
        R: RngCore + CryptoRng;

    fn dkg_part2(&self, r1_s: &[u8], r1_pm: &Self::RoundPackages) -> Self::Part2Result;
    fn dkg_verify_private_package(&self, i: A, r1: &[u8], r2: &[u8]) -> Result<(), E>;

    fn dkg_part3(
        &self,
        i: A,
        r2_s: &[u8],
        r1_pm: &Self::RoundPackages,
        r2_pm: &Self::RoundPackages,
    ) -> Self::Part3Result;

    fn dkg_derive_verifying_share(&self, sk: &[u8]) -> Result<Vec<u8>, E>;
    fn dkg_decompress_key(&self, vk: &[u8]) -> Result<Vec<u8>, E>;
}

pub trait MerkleTreeBuilder<A, E>: IdentifierConverter<A, E> {
    type Hash;

    fn build_merkle_tree<C, V, F>(
        verifying_shares: &BTreeMap<Identifier<C>, V>,
        serialize_fn: F,
    ) -> Result<Vec<Self::Hash>, E>
    where
        C: frost_core::Ciphersuite,
        <<C as Ciphersuite>::Group as Group>::Field: frost_core::Field,
		F: Fn(&V) -> Result<Vec<u8>, E>;

    fn generate_merkle_proof_from_tree<C, V, F>(
        verifying_shares: &BTreeMap<Identifier<C>, V>,
        merkle_tree: &[Self::Hash],
        authority_index: A,
        serialize_fn: F,
    ) -> Result<(Vec<u8>, Vec<Self::Hash>), E>
    where
        C: frost_core::Ciphersuite,
        <<C as Ciphersuite>::Group as Group>::Field: frost_core::Field,
		F: Fn(&V) -> Result<Vec<u8>, E>;

    fn verify_merkle_proof(
        &self,
        verifying_share: &[u8],
        merkle_proof: &[Self::Hash],
        merkle_root: Self::Hash,
        authority_index: A,
    ) -> Result<(), E>;
}
