use bls::{PublicKey, SecretKey, Signature};
use context_deserialize::context_deserialize;
use serde::{Deserialize, Serialize};
use ssz_derive::{Decode, Encode};
use superstruct::superstruct;
use tree_hash_derive::TreeHash;

use crate::{
    attestation::{
        Attestation, AttestationBase, AttestationElectra, AttestationRef, SelectionProof,
    },
    core::{ChainSpec, Domain, EthSpec, Hash256, SignedRoot},
    fork::{Fork, ForkName},
};

#[superstruct(
    variants(Base, Electra),
    variant_attributes(
        derive(
            Debug,
            Clone,
            PartialEq,
            Serialize,
            Deserialize,
            Encode,
            Decode,
            TreeHash,
        ),
        context_deserialize(ForkName),
        serde(bound = "E: EthSpec"),
        cfg_attr(
            feature = "arbitrary",
            derive(arbitrary::Arbitrary),
            arbitrary(bound = "E: EthSpec"),
        ),
    ),
    ref_attributes(
        derive(Debug, PartialEq, TreeHash, Serialize),
        serde(untagged, bound = "E: EthSpec"),
        tree_hash(enum_behaviour = "transparent")
    ),
    map_ref_into(AttestationRef)
)]
#[cfg_attr(
    feature = "arbitrary",
    derive(arbitrary::Arbitrary),
    arbitrary(bound = "E: EthSpec")
)]
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Encode, TreeHash)]
#[serde(untagged)]
#[tree_hash(enum_behaviour = "transparent")]
#[ssz(enum_behaviour = "transparent")]
#[serde(bound = "E: EthSpec", deny_unknown_fields)]
pub struct AggregateAndProof<E: EthSpec> {
    /// The index of the validator that created the attestation.
    #[serde(with = "serde_utils::quoted_u64")]
    #[superstruct(getter(copy))]
    pub aggregator_index: u64,
    /// The aggregate attestation.
    #[superstruct(flatten)]
    pub aggregate: Attestation<E>,
    /// A proof provided by the validator that permits them to publish on the
    /// `beacon_aggregate_and_proof` gossipsub topic.
    pub selection_proof: Signature,
}

impl<'a, E: EthSpec> AggregateAndProofRef<'a, E> {
    /// Returns `true` if `validator_pubkey` signed over `self.aggregate.data.slot`.
    pub fn aggregate(self) -> AttestationRef<'a, E> {
        map_aggregate_and_proof_ref_into_attestation_ref!(&'a _, self, |inner, cons| {
            cons(&inner.aggregate)
        })
    }
}
impl<E: EthSpec> AggregateAndProof<E> {
    /// Returns `true` if `validator_pubkey` signed over `self.aggregate.data.slot`.
    pub fn aggregate<'a>(&'a self) -> AttestationRef<'a, E> {
        map_aggregate_and_proof_ref_into_attestation_ref!(&'a _, self.to_ref(), |inner, cons| {
            cons(&inner.aggregate)
        })
    }
}

impl<E: EthSpec> AggregateAndProof<E> {
    /// Produces a new `AggregateAndProof` with a `selection_proof` generated by signing
    /// `aggregate.data.slot` with `secret_key`.
    ///
    /// If `selection_proof.is_none()` it will be computed locally.
    pub fn from_aggregate(
        aggregator_index: u64,
        aggregate: AttestationRef<'_, E>,
        selection_proof: Option<SelectionProof>,
        secret_key: &SecretKey,
        fork: &Fork,
        genesis_validators_root: Hash256,
        spec: &ChainSpec,
    ) -> Self {
        let selection_proof = selection_proof.unwrap_or_else(|| {
            SelectionProof::new::<E>(
                aggregate.data().slot,
                secret_key,
                fork,
                genesis_validators_root,
                spec,
            )
        });

        Self::from_attestation(
            aggregator_index,
            aggregate.clone_as_attestation(),
            selection_proof,
        )
    }

    /// Produces a new `AggregateAndProof` given a `selection_proof`
    pub fn from_attestation(
        aggregator_index: u64,
        aggregate: Attestation<E>,
        selection_proof: SelectionProof,
    ) -> Self {
        match aggregate {
            Attestation::Base(aggregate) => Self::Base(AggregateAndProofBase {
                aggregator_index,
                aggregate,
                selection_proof: selection_proof.into(),
            }),
            Attestation::Electra(aggregate) => Self::Electra(AggregateAndProofElectra {
                aggregator_index,
                aggregate,
                selection_proof: selection_proof.into(),
            }),
        }
    }

    /// Returns `true` if `validator_pubkey` signed over `self.aggregate.data.slot`.
    pub fn is_valid_selection_proof(
        &self,
        validator_pubkey: &PublicKey,
        fork: &Fork,
        genesis_validators_root: Hash256,
        spec: &ChainSpec,
    ) -> bool {
        let target_epoch = self.aggregate().data().slot.epoch(E::slots_per_epoch());
        let domain = spec.get_domain(
            target_epoch,
            Domain::SelectionProof,
            fork,
            genesis_validators_root,
        );
        let message = self.aggregate().data().slot.signing_root(domain);
        self.selection_proof().verify(validator_pubkey, message)
    }
}

impl<E: EthSpec> SignedRoot for AggregateAndProof<E> {}
impl<E: EthSpec> SignedRoot for AggregateAndProofRef<'_, E> {}
