use sp_std::vec::Vec;
use scale_info::TypeInfo;
use sp_runtime::RuntimeDebug;
use codec::{Encode, Decode};

use frame_support::{BoundedVec, traits::ConstU32};
use crate::{AuthIndex, ExodusSession, ExodusHash, ExodusError};

pub trait PackageMetadata<NetworkCurve>
where
	NetworkCurve: Copy,
{
    fn context(&self) -> &PackageContext<NetworkCurve>;

    fn get_authority_index(&self) -> AuthIndex {
        self.context().get_authority_index()
    }
    fn get_network_curve(&self) -> NetworkCurve {
        self.context().get_network_curve()
    }
}

#[derive(Encode, Decode, Clone, RuntimeDebug, TypeInfo, PartialEq, Eq)]
#[scale_info(skip_type_params(T))]
pub enum ExodusPackage<NetworkCurve, T>
where
	T: crate::pallet::Config,
{
    NonceCommitment(ExodusNonceCommitment<NetworkCurve>),
    GroupCommitment(ExodusGroupCommitment<NetworkCurve>),
    SignatureShare(ExodusSignatureShare<NetworkCurve, T>),
}

impl<NetworkCurve, T> ExodusPackage<NetworkCurve, T>
where
	T: crate::pallet::Config,
{
    pub fn bytes_len(&self) -> u32 {
        let length = match self {
            Self::GroupCommitment(package) => package.group_commitment.len(),
            Self::SignatureShare(package) => {
                (package.binding_factors_proof.len() * ExodusHash::len_bytes())
                    + package.self_binding_factor.len()
                    + package.signature_share.len()
            },
            Self::NonceCommitment(package) => {
                package.hiding_commitment.len()
                    + package.binding_commitment.len()
            },
        };

        length as u32
    }

    pub fn get_exodus_session(&self) -> ExodusSession {
        match &self {
            Self::SignatureShare(package) => package.session,
            Self::NonceCommitment(package) => package.session,
            Self::GroupCommitment(package) => package.session,
        }
    }
}

#[derive(Encode, Decode, Clone, RuntimeDebug, TypeInfo, PartialEq, Eq)]
#[scale_info(skip_type_params(T))]
pub enum DkgPackage<NetworkCurve, T>
where
	T: crate::pallet::Config,
{
    Round0(Round0Package<NetworkCurve>),
    Round1(Round1Package<NetworkCurve, T>),
    Round2(Round2Package<NetworkCurve, T>),
    Round3(Round3Package<NetworkCurve, T>),
    Round4(Round4Package<NetworkCurve, T>),
    Round5(Round5Package<NetworkCurve>),
    Round6(Round6Package<NetworkCurve, T>),
}

impl<NetworkCurve, T> DkgPackage<NetworkCurve, T>
where
	T: crate::pallet::Config,
{
    pub fn bytes_len(&self) -> u32 {
        let length = match self {
            Self::Round0(_)      => ExodusHash::len_bytes(),
            Self::Round1(round1) => round1.package.len(),
            Self::Round2(round2) => round2.bundle.bitmask.len(),
            Self::Round3(round3) => round3.indices.len(),
            Self::Round4(round4) => round4.bundle.bitmask.len(),
            Self::Round5(round5) => {
                round5.verifying_key.len()
                    .saturating_add(ExodusHash::len_bytes())
            },
            Self::Round6(round6) => round6.merkle_proof.len(),
        };

        length as u32
    }
}

#[derive(Default, RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
pub struct PackageContext<NetworkCurve> {
    authority_index: AuthIndex,
    network_curve: NetworkCurve
}

impl<NetworkCurve> PackageContext<NetworkCurve>
where
	NetworkCurve: Copy,
{
    pub fn get_authority_index(&self) -> AuthIndex {
        self.authority_index
    }

    pub fn get_network_curve(&self) -> NetworkCurve {
        self.network_curve
    }

    pub fn with_network_curve(mut self, network_curve: NetworkCurve) -> Self {
        self.network_curve = network_curve;
        self
    }

    pub fn with_authority_index(mut self, authority_index: AuthIndex) -> Self {
        self.authority_index = authority_index;
        self
    }
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
pub struct BundledPackages {
    pub blob: Vec<u8>,
    pub bitmask: Vec<u8>,
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
#[scale_info(skip_type_params(T))]
pub struct EncryptionBundledPackages<T>
where
	T: crate::pallet::Config,
{
    pub blob: BoundedVec<u8, crate::EncryptedRound2BlobMaxBytes<T>>,
    pub bitmask: BoundedVec<u8, crate::MaxAuthoritiesBitmaskSize<T>>
}

impl<T: crate::pallet::Config> EncryptionBundledPackages<T> {
    pub fn len(&self) -> usize {
        let blob_len = self.blob.len();
        let bitmask_len = self.bitmask.len();

        bitmask_len.saturating_add(blob_len)
    }
}

impl<T: crate::pallet::Config> TryFrom<BundledPackages> for EncryptionBundledPackages<T> {
    type Error = ExodusError;

    fn try_from(source: BundledPackages) -> Result<Self, Self::Error> {
        let blob =
            BoundedVec::<u8, crate::EncryptedRound2BlobMaxBytes<T>>::try_from(source.blob)
                .map_err(|_| ExodusError::DeserializationError)?;

        let bitmask = BoundedVec::<u8, crate::MaxAuthoritiesBitmaskSize<T>>::try_from(source.bitmask)
            .map_err(|_| ExodusError::DeserializationError)?;

        Ok(Self { blob, bitmask })
    }
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
#[scale_info(skip_type_params(T))]
pub struct JustificationsBundledPackages<T>
where
	T: crate::pallet::Config,
{
    pub blob: BoundedVec<u8, crate::Round2BlobMaxBytes<T>>,
    pub bitmask: BoundedVec<u8, crate::MaxAuthoritiesBitmaskSize<T>>
}

impl<T: crate::pallet::Config> JustificationsBundledPackages<T> {
    pub fn len(&self) -> usize {
        let blob_len = self.blob.len();
        let bitmask_len = self.bitmask.len();

        bitmask_len.saturating_add(blob_len)
    }
}

impl<T: crate::pallet::Config> TryFrom<BundledPackages> for JustificationsBundledPackages<T> {
    type Error = ExodusError;

    fn try_from(source: BundledPackages) -> Result<Self, Self::Error> {
        let blob =
            BoundedVec::<u8, crate::Round2BlobMaxBytes<T>>::try_from(source.blob)
                .map_err(|_| ExodusError::DeserializationError)?;

        let bitmask = BoundedVec::<u8, crate::MaxAuthoritiesBitmaskSize<T>>::try_from(source.bitmask)
            .map_err(|_| ExodusError::DeserializationError)?;

        Ok(Self { blob, bitmask })
    }
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
#[scale_info(skip_type_params(T))]
pub struct ExodusNonceCommitment<NetworkCurve> {
    context: PackageContext<NetworkCurve>,
    pub hiding_commitment: BoundedVec<u8, ConstU32<{ crate::ELEMENT_MAX_BYTES }>>,
    pub binding_commitment: BoundedVec<u8, ConstU32<{ crate::ELEMENT_MAX_BYTES }>>,
    pub session: ExodusSession,
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
pub struct ExodusGroupCommitment<NetworkCurve> {
    context: PackageContext<NetworkCurve>,
    pub group_commitment: BoundedVec<u8, ConstU32<{ crate::ELEMENT_MAX_BYTES }>>,
    pub binding_factors_root: ExodusHash,
    pub session: ExodusSession,
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
#[scale_info(skip_type_params(T))]
pub struct ExodusSignatureShare<NetworkCurve, T>
where
	T: crate::pallet::Config,
{
    context: PackageContext<NetworkCurve>,
    pub binding_factors_proof: BoundedVec<ExodusHash, crate::BindingFactorsProof<T>>,
    pub signature_share: BoundedVec<u8, ConstU32<{ crate::SCALAR_MAX_BYTES }>>,
    pub self_binding_factor: BoundedVec<u8, ConstU32<{ crate::SCALAR_MAX_BYTES }>>,
    pub session: ExodusSession,
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
pub struct Round0Package<NetworkCurve> {
    context: PackageContext<NetworkCurve>,
    pub package_hash: ExodusHash,
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
#[scale_info(skip_type_params(T))]
pub struct Round1Package<NetworkCurve, T>
where
	T: crate::pallet::Config,
{
    context: PackageContext<NetworkCurve>,
    pub package: BoundedVec<u8, crate::Round1MaxBytes<T>>,
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
#[scale_info(skip_type_params(T))]
pub struct Round2Package<NetworkCurve, T>
where
	T: crate::pallet::Config,
{
    context: PackageContext<NetworkCurve>,
    pub bundle: EncryptionBundledPackages<T>,
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
#[scale_info(skip_type_params(T))]
pub struct Round3Package<NetworkCurve, T>
where
	T: crate::pallet::Config,
{
    context: PackageContext<NetworkCurve>,
    pub indices: BoundedVec<u8, crate::AccusedIndicesMaxBytes<T>>,
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
#[scale_info(skip_type_params(T))]
pub struct Round4Package<NetworkCurve, T>
where
	T: crate::pallet::Config,
{
    context: PackageContext<NetworkCurve>,
    pub bundle: JustificationsBundledPackages<T>,
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
pub struct Round5Package<NetworkCurve> {
    context: PackageContext<NetworkCurve>,
    pub verifying_key: BoundedVec<u8, ConstU32<{ crate::ELEMENT_MAX_BYTES }>>,
    pub merkle_root: ExodusHash,
}

#[derive(RuntimeDebug, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
#[scale_info(skip_type_params(T))]
pub struct Round6Package<NetworkCurve, T>
where
	T: crate::pallet::Config,
{
    context: PackageContext<NetworkCurve>,
    pub verifying_share: BoundedVec<u8, ConstU32<{ crate::ELEMENT_MAX_BYTES }>>,
    pub merkle_proof: BoundedVec<ExodusHash, crate::MerkleProofMaxSize<T>>,
}

impl<NetworkCurve> PackageContext<NetworkCurve> {
    pub fn with_exodus_nonce_commitment(
        self,
        session: ExodusSession,
        hiding_commitment_vec: Vec<u8>,
        binding_commitment_vec: Vec<u8>,
    ) -> Result<ExodusNonceCommitment<NetworkCurve>, ExodusError> {
        let hiding_commitment =
            BoundedVec::<u8, ConstU32<{ crate::ELEMENT_MAX_BYTES }>>::try_from(hiding_commitment_vec)
                .map_err(|_| ExodusError::DeserializationError)?;

        let binding_commitment =
            BoundedVec::<u8, ConstU32<{ crate::ELEMENT_MAX_BYTES }>>::try_from(binding_commitment_vec)
                .map_err(|_| ExodusError::DeserializationError)?;

        Ok(ExodusNonceCommitment { context: self, hiding_commitment, binding_commitment, session })
    }

    pub fn with_exodus_group_commitment(
        self,
        session: ExodusSession,
        commitment_vec: Vec<u8>,
        binding_factors_root: ExodusHash,
    ) -> Result<ExodusGroupCommitment<NetworkCurve>, ExodusError> {
        let group_commitment =
            BoundedVec::<u8, ConstU32<{ crate::ELEMENT_MAX_BYTES }>>::try_from(commitment_vec)
                .map_err(|_| ExodusError::DeserializationError)?;

        Ok(ExodusGroupCommitment {
            context: self,
            group_commitment,
            binding_factors_root,
            session,
        })
    }

    pub fn with_exodus_signature_share<T: crate::pallet::Config>(
        self,
        session: ExodusSession,
        signature_share_vec: Vec<u8>,
        binding_factors_proof_vec: Vec<ExodusHash>,
        self_binding_factor_vec: Vec<u8>,
    ) -> Result<ExodusSignatureShare<NetworkCurve, T>, ExodusError> {
        let signature_share =
            BoundedVec::<u8, ConstU32<{ crate::SCALAR_MAX_BYTES }>>::try_from(signature_share_vec)
                .map_err(|_| ExodusError::DeserializationError)?;

        let binding_factors_proof =
            BoundedVec::<ExodusHash, crate::BindingFactorsProof<T>>::try_from(binding_factors_proof_vec)
                .map_err(|_| ExodusError::DeserializationError)?;

        let self_binding_factor =
            BoundedVec::<u8, ConstU32<{ crate::SCALAR_MAX_BYTES }>>::try_from(self_binding_factor_vec)
                .map_err(|_| ExodusError::DeserializationError)?;

        Ok(ExodusSignatureShare {
            context: self,
            binding_factors_proof,
            self_binding_factor,
            signature_share,
            session,
        })
    }

    pub fn with_dkg_round0(
        self,
        package_hash: ExodusHash,
    ) -> Round0Package<NetworkCurve> {
        Round0Package { context: self, package_hash }
    }

    pub fn with_dkg_round1<T: crate::pallet::Config>(
        self,
        package_vec: Vec<u8>,
    ) -> Result<Round1Package<NetworkCurve, T>, ExodusError> {
        let package =
            BoundedVec::<u8, crate::Round1MaxBytes<T>>::try_from(package_vec)
                .map_err(|_| ExodusError::DeserializationError)?;

        Ok(Round1Package { context: self, package })
    }

    pub fn with_dkg_round2<T: crate::pallet::Config>(
        self,
        bundle_raw: BundledPackages,
    ) -> Result<Round2Package<NetworkCurve, T>, ExodusError> {
        let bundle: EncryptionBundledPackages<T> = bundle_raw.try_into()?;
        Ok(Round2Package { context: self, bundle })
    }

    pub fn with_dkg_round3<T: crate::pallet::Config>(
        self,
        indices_vec: Vec<u8>
    ) -> Result<Round3Package<NetworkCurve, T>, ExodusError> {
        let indices =
            BoundedVec::<u8, crate::AccusedIndicesMaxBytes<T>>::try_from(indices_vec)
                .map_err(|_| ExodusError::DeserializationError)?;

        Ok(Round3Package { context: self, indices })
    }

    pub fn with_dkg_round4<T: crate::pallet::Config>(
        self,
        bundle_raw: BundledPackages,
    ) -> Result<Round4Package<NetworkCurve, T>, ExodusError> {
        let bundle: JustificationsBundledPackages<T> = bundle_raw.try_into()?;
        Ok(Round4Package { context: self, bundle })
    }

    pub fn with_dkg_round5(
        self,
        verifying_key_vec: Vec<u8>,
        merkle_root: ExodusHash,
    ) -> Result<Round5Package<NetworkCurve>, ExodusError> {
        let verifying_key =
            BoundedVec::<u8, ConstU32<{ crate::ELEMENT_MAX_BYTES }>>::try_from(verifying_key_vec)
                .map_err(|_| ExodusError::DeserializationError)?;

        Ok(Round5Package { context: self, verifying_key, merkle_root })
    }

    pub fn with_dkg_round6<T: crate::pallet::Config>(
        self,
        verifying_share_vec: Vec<u8>,
        merkle_proof_vec: Vec<ExodusHash>,
    ) -> Result<Round6Package<NetworkCurve, T>, ExodusError> {
        let verifying_share =
            BoundedVec::<u8, ConstU32<{ crate::ELEMENT_MAX_BYTES }>>::try_from(verifying_share_vec)
                .map_err(|_| ExodusError::DeserializationError)?;

        let merkle_proof =
            BoundedVec::<ExodusHash, crate::MerkleProofMaxSize<T>>::try_from(merkle_proof_vec)
                .map_err(|_| ExodusError::DeserializationError)?;

        Ok(Round6Package { context: self, verifying_share, merkle_proof })
    }
}

impl<NC, T> PackageMetadata<NC> for ExodusPackage<NC, T>
where
	NC: Copy,
	T: crate::pallet::Config,
{
    fn context(&self) -> &PackageContext<NC> {
        match self {
            Self::NonceCommitment(package) => &package.context,
            Self::GroupCommitment(package) => &package.context,
            Self::SignatureShare(package) => &package.context,
        }
    }
}

impl<NC, T> PackageMetadata<NC> for DkgPackage<NC, T>
where
	NC: Copy,
	T: crate::pallet::Config,
{
    fn context(&self) -> &PackageContext<NC> {
        match self {
            Self::Round0(package) => &package.context,
            Self::Round1(package) => &package.context,
            Self::Round2(package) => &package.context,
            Self::Round3(package) => &package.context,
            Self::Round4(package) => &package.context,
            Self::Round5(package) => &package.context,
            Self::Round6(package) => &package.context,
        }
    }
}
