mod key_generation;
mod key_resharing;
mod sign;

use crate::config::{MpcConfig, ParticipantsConfig};
use crate::metrics::tokio_task_metrics::EDDSA_TASK_MONITORS;
use crate::network::{MeshNetworkClient, NetworkTaskChannel};
use crate::primitives::MpcTaskId;
#[cfg(test)]
use crate::providers::PublicKeyConversion;
use crate::providers::SignatureProvider;
use crate::storage::SignRequestStorage;
use crate::types::SignatureId;
#[cfg(test)]
use anyhow::Context;
use borsh::{BorshDeserialize, BorshSerialize};
use mpc_node_config::ConfigFile;
use mpc_primitives::domain::DomainId;
#[cfg(test)]
use near_mpc_contract_interface::types::Ed25519PublicKey;
use near_mpc_contract_interface::types::KeyEventId;
use std::collections::HashMap;
use std::sync::Arc;
use threshold_signatures::ReconstructionThreshold;
use threshold_signatures::frost::eddsa::KeygenOutput;
use threshold_signatures::frost_ed25519::keys::SigningShare;
use threshold_signatures::frost_ed25519::{Signature, VerifyingKey};

#[derive(Clone)]
pub struct EddsaSignatureProvider {
    config: Arc<ConfigFile>,
    mpc_config: Arc<MpcConfig>,
    client: Arc<MeshNetworkClient>,
    sign_request_store: Arc<SignRequestStorage>,
    keyshares: HashMap<DomainId, KeygenOutput>,
}

impl EddsaSignatureProvider {
    pub fn new(
        config: Arc<ConfigFile>,
        mpc_config: Arc<MpcConfig>,
        client: Arc<MeshNetworkClient>,
        sign_request_store: Arc<SignRequestStorage>,
        keyshares: HashMap<DomainId, KeygenOutput>,
    ) -> Self {
        Self {
            config,
            mpc_config,
            client,
            sign_request_store,
            keyshares,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, BorshSerialize, BorshDeserialize)]
pub enum EddsaTaskId {
    KeyGeneration { key_event: KeyEventId },
    KeyResharing { key_event: KeyEventId },
    Signature { id: SignatureId },
}

impl From<EddsaTaskId> for MpcTaskId {
    fn from(value: EddsaTaskId) -> Self {
        MpcTaskId::EddsaTaskId(value)
    }
}

impl SignatureProvider for EddsaSignatureProvider {
    type PublicKey = VerifyingKey;
    type SecretShare = SigningShare;
    type KeygenOutput = KeygenOutput;
    type Signature = Signature;
    type TaskId = EddsaTaskId;

    async fn make_signature(
        &self,
        id: SignatureId,
    ) -> anyhow::Result<(Self::Signature, Self::PublicKey)> {
        EDDSA_TASK_MONITORS
            .make_signature_leader
            .instrument(self.make_signature_leader(id))
            .await
    }

    async fn run_key_generation_client(
        threshold: ReconstructionThreshold,
        channel: NetworkTaskChannel,
    ) -> anyhow::Result<Self::KeygenOutput> {
        Self::run_key_generation_client_internal(threshold, channel).await
    }

    async fn run_key_resharing_client(
        new_threshold: ReconstructionThreshold,
        key_share: Option<SigningShare>,
        public_key: VerifyingKey,
        old_participants: &ParticipantsConfig,
        channel: NetworkTaskChannel,
    ) -> anyhow::Result<Self::KeygenOutput> {
        Self::run_key_resharing_client_internal(
            new_threshold,
            key_share,
            public_key,
            old_participants,
            channel,
        )
        .await
    }

    async fn process_channel(&self, channel: NetworkTaskChannel) -> anyhow::Result<()> {
        match channel.task_id() {
            MpcTaskId::EddsaTaskId(task) => match task {
                EddsaTaskId::KeyGeneration { .. } => {
                    anyhow::bail!("Key generation rejected in normal node operation");
                }
                EddsaTaskId::KeyResharing { .. } => {
                    anyhow::bail!("Key resharing rejected in normal node operation");
                }
                EddsaTaskId::Signature { id } => {
                    EDDSA_TASK_MONITORS
                        .make_signature_follower
                        .instrument(self.make_signature_follower(channel, id))
                        .await?;
                }
            },
            _ => anyhow::bail!(
                "eddsa task handler: received unexpected task id: {:?}",
                channel.task_id()
            ),
        }

        Ok(())
    }

    async fn spawn_background_tasks(self: Arc<Self>) -> anyhow::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
impl PublicKeyConversion for VerifyingKey {
    fn to_near_sdk_public_key(&self) -> anyhow::Result<near_sdk::PublicKey> {
        let data = self.serialize()?;
        let data: [u8; 32] = data
            .try_into()
            .or_else(|_| anyhow::bail!("Serialized public key is not 32 bytes."))?;
        Ok(near_sdk::PublicKey::from(Ed25519PublicKey::from(data)))
    }

    fn from_near_sdk_public_key(public_key: &near_sdk::PublicKey) -> anyhow::Result<Self> {
        let ed25519_pk = Ed25519PublicKey::try_from(public_key)
            .map_err(|_| anyhow::anyhow!("Not an ed25519 public key"))?;
        VerifyingKey::deserialize(ed25519_pk.as_ref())
            .context("Failed to convert SDK public key to frost_ed25519::VerifyingKey")
    }
}

#[cfg(test)]
impl PublicKeyConversion for ed25519_dalek::VerifyingKey {
    fn to_near_sdk_public_key(&self) -> anyhow::Result<near_sdk::PublicKey> {
        Ok(near_sdk::PublicKey::from(Ed25519PublicKey::from(self)))
    }

    fn from_near_sdk_public_key(public_key: &near_sdk::PublicKey) -> anyhow::Result<Self> {
        let ed25519_pk = Ed25519PublicKey::try_from(public_key)
            .map_err(|_| anyhow::anyhow!("Not an ed25519 public key"))?;
        ed25519_dalek::VerifyingKey::try_from(&ed25519_pk)
            .map_err(|_| anyhow::anyhow!("Failed to convert to ed25519_dalek::VerifyingKey"))
    }
}

#[cfg(test)]
mod tests {
    use std::str::FromStr;

    use rand::SeedableRng as _;
    use threshold_signatures::frost_ed25519::{Ed25519Sha512, VerifyingKey};
    use threshold_signatures::test_utils::{generate_participants_with_random_ids, run_keygen};

    use crate::providers::PublicKeyConversion;
    #[test]
    fn check_pubkey_conversion_to_sdk() -> anyhow::Result<()> {
        let mut rng = rand::rngs::StdRng::from_seed([1u8; 32]);
        let x = run_keygen::<Ed25519Sha512, _>(
            &generate_participants_with_random_ids(4, &mut rng),
            3,
            &mut rng,
        )
        .into_iter()
        .next()
        .unwrap()
        .1;
        let _ = x.public_key.to_near_sdk_public_key().unwrap();
        Ok(())
    }

    #[test]
    fn check_pubkey_conversion_from_sdk() -> anyhow::Result<()> {
        let near_sdk =
            near_sdk::PublicKey::from_str("ed25519:6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp")?;
        let _ = VerifyingKey::from_near_sdk_public_key(&near_sdk)?;
        Ok(())
    }
}