use crate::metrics;
use crate::network::NetworkTaskChannel;
use crate::network::computation::MpcLeaderCentricComputation;
use crate::primitives::UniqueId;
use crate::protocol::NamedProtocol;
use crate::providers::ecdsa::presign::PresignOutputWithParticipants;
use crate::providers::ecdsa::{
    EcdsaSignatureProvider, EcdsaTaskId, KeygenOutput, PresignatureStorage,
};
use crate::types::{SignatureId, SignatureRequest};
use anyhow::Context;
use k256::Scalar;
use k256::elliptic_curve::PrimeField;
use near_mpc_contract_interface::types::Tweak;
use std::sync::Arc;
use std::time::Duration;
use threshold_signatures::ParticipantList;
use threshold_signatures::ReconstructionThreshold;
use threshold_signatures::ecdsa::ot_based_ecdsa::{PresignOutput, RerandomizedPresignOutput};
use threshold_signatures::ecdsa::{RerandomizationArguments, Signature, SignatureOption};
use threshold_signatures::frost_secp256k1::VerifyingKey;
use threshold_signatures::participants::Participant;
use tokio::time::timeout;

impl EcdsaSignatureProvider {
    pub(crate) async fn make_signature_leader_given_parameters(
        &self,
        sign_request: SignatureRequest,
        presignature: PresignOutputWithParticipants,
        channel: NetworkTaskChannel,
    ) -> anyhow::Result<(Signature, VerifyingKey)> {
        let keyshare = self.keyshare(sign_request.domain)?;
        let participants = presignature.participants.clone();
        let reconstruction_threshold: usize =
            keyshare.reconstruction_threshold.inner().try_into()?;
        let reconstruction_threshold = ReconstructionThreshold::from(reconstruction_threshold);

        let (signature, public_key) = SignComputation {
            keygen_out: keyshare.keygen_output,
            reconstruction_threshold,
            presign_out: presignature.presignature,
            msg_hash: *sign_request
                .payload
                .as_ecdsa()
                .ok_or_else(|| anyhow::anyhow!("Payload is not an ECDSA payload"))?,
            tweak: sign_request.tweak,
            entropy: sign_request.entropy,
        }
        .perform_leader_centric_computation(
            channel,
            Duration::from_secs(self.config.signature.timeout_sec),
        )
        .await
        .inspect_err(|_| {
            participants.iter().for_each(|id| {
                metrics::PARTICIPANT_TOTAL_TIMES_SEEN_IN_FAILED_SIGNATURE_COMPUTATION_LEADER
                    .with_label_values(&[&id.raw().to_string()])
                    .inc();
            });
        })?;

        Ok((
            signature.context("Leader should obtain a signature")?,
            public_key,
        ))
    }

    pub(super) async fn make_signature_leader(
        &self,
        id: SignatureId,
    ) -> anyhow::Result<(Signature, VerifyingKey)> {
        let sign_request = self.sign_request_store.get(id).await?;
        let keyshare = self.keyshare(sign_request.domain)?;
        let (presignature_id, presignature) = keyshare.presignature_store.take_owned().await;
        let participants = presignature.participants.clone();
        let channel = self.new_channel_for_task(
            EcdsaTaskId::Signature {
                id,
                presignature_id,
            },
            participants,
        )?;
        self.make_signature_leader_given_parameters(sign_request, presignature, channel)
            .await
    }

    pub(crate) async fn make_signature_follower_given_request(
        &self,
        channel: NetworkTaskChannel,
        presignature_id: UniqueId,
        sign_request: SignatureRequest,
    ) -> anyhow::Result<()> {
        // The presignature must be owned by the leader, never one of ours.
        presignature_id.validate_owned_by(channel.sender().get_leader())?;
        let keyshare = self.keyshare(sign_request.domain)?;
        let reconstruction_threshold: usize =
            keyshare.reconstruction_threshold.inner().try_into()?;
        let reconstruction_threshold = ReconstructionThreshold::from(reconstruction_threshold);

        let participants = channel.participants().to_vec();
        FollowerSignComputation {
            keygen_out: keyshare.keygen_output,
            reconstruction_threshold,
            presignature_store: keyshare.presignature_store.clone(),
            presignature_id,
            msg_hash: *sign_request
                .payload
                .as_ecdsa()
                .ok_or_else(|| anyhow::anyhow!("Payload is not an ECDSA payload"))?,
            tweak: sign_request.tweak,
            entropy: sign_request.entropy,
        }
        .perform_leader_centric_computation(
            channel,
            Duration::from_secs(self.config.signature.timeout_sec),
        )
        .await
        .inspect_err(|_| {
            participants.iter().for_each(|id| {
                metrics::PARTICIPANT_TOTAL_TIMES_SEEN_IN_FAILED_SIGNATURE_COMPUTATION_FOLLOWER
                    .with_label_values(&[&id.raw().to_string()])
                    .inc();
            });
        })?;

        Ok(())
    }

    pub(crate) async fn make_signature_follower(
        &self,
        channel: NetworkTaskChannel,
        id: SignatureId,
        presignature_id: UniqueId,
    ) -> anyhow::Result<()> {
        metrics::MPC_NUM_PASSIVE_SIGN_REQUESTS_RECEIVED.inc();
        let sign_request = timeout(
            Duration::from_secs(self.config.signature.timeout_sec),
            self.sign_request_store.get(id),
        )
        .await??;
        metrics::MPC_NUM_PASSIVE_SIGN_REQUESTS_LOOKUP_SUCCEEDED.inc();

        self.make_signature_follower_given_request(channel, presignature_id, sign_request)
            .await
    }
}

/// Performs an MPC signature operation. This is the same for the initiator
/// and for passive participants.
/// The entropy is used to rerandomize the presignature (inspired by
/// \[[GS21](https://eprint.iacr.org/2021/1330.pdf)\])
/// The tweak allows key derivation
pub struct SignComputation {
    pub keygen_out: KeygenOutput,
    pub reconstruction_threshold: ReconstructionThreshold,
    pub presign_out: PresignOutput,
    pub msg_hash: [u8; 32],
    pub tweak: Tweak,
    pub entropy: [u8; 32],
}

impl NamedProtocol for SignComputation {
    const NAME: &'static str = "sign cait-sith";
}

#[async_trait::async_trait]
impl MpcLeaderCentricComputation<(SignatureOption, VerifyingKey)> for SignComputation {
    async fn compute(
        self,
        channel: &mut NetworkTaskChannel,
    ) -> anyhow::Result<(SignatureOption, VerifyingKey)> {
        let cs_participants = channel
            .participants()
            .iter()
            .copied()
            .map(Participant::from)
            .collect::<Vec<_>>();

        let tweak = Scalar::from_repr(self.tweak.as_bytes().into())
            .into_option()
            .context("Couldn't construct k256 point")?;
        let tweak = threshold_signatures::Tweak::new(tweak);

        let msg_hash = Scalar::from_repr(self.msg_hash.into())
            .into_option()
            .context("Couldn't construct k256 point")?;

        let derived_public_key = tweak
            .derive_verifying_key(&self.keygen_out.public_key)
            .to_element()
            .to_affine();
        let participants = ParticipantList::new(&cs_participants).unwrap();

        let rerand_args = RerandomizationArguments::new(
            self.keygen_out.public_key.to_element().to_affine(),
            tweak,
            self.msg_hash,
            self.presign_out.big_r,
            participants,
            self.entropy,
        );
        let rerandomized_presignature =
            RerandomizedPresignOutput::rerandomize_presign(&self.presign_out, &rerand_args)?;

        let protocol = threshold_signatures::ecdsa::ot_based_ecdsa::sign::sign(
            &cs_participants,
            channel.sender().get_leader().into(),
            self.reconstruction_threshold,
            channel.my_participant_id().into(),
            derived_public_key,
            rerandomized_presignature,
            msg_hash,
        )?;
        let _timer = metrics::MPC_SIGNATURE_TIME_ELAPSED.start_timer();
        let signature = Self::run(channel, protocol).await?;
        Ok((signature, VerifyingKey::new(derived_public_key.into())))
    }

    fn leader_waits_for_success(&self) -> bool {
        false
    }
}

/// Performs an MPC signature operation as a follower.
/// The difference is that the follower needs to look up the presignature, which may fail.
pub struct FollowerSignComputation {
    pub keygen_out: KeygenOutput,
    pub reconstruction_threshold: ReconstructionThreshold,
    pub presignature_id: UniqueId,
    pub presignature_store: Arc<PresignatureStorage>,
    pub msg_hash: [u8; 32],
    pub tweak: Tweak,
    pub entropy: [u8; 32],
}

#[async_trait::async_trait]
impl MpcLeaderCentricComputation<()> for FollowerSignComputation {
    async fn compute(self, channel: &mut NetworkTaskChannel) -> anyhow::Result<()> {
        let presign_out = self
            .presignature_store
            .take_unowned(self.presignature_id)?
            .presignature;
        SignComputation {
            keygen_out: self.keygen_out,
            reconstruction_threshold: self.reconstruction_threshold,
            presign_out,
            msg_hash: self.msg_hash,
            tweak: self.tweak,
            entropy: self.entropy,
        }
        .compute(channel)
        .await?;
        Ok(())
    }

    fn leader_waits_for_success(&self) -> bool {
        false
    }
}