pub use signature::Signer;

use generic_array::GenericArray;

use super::{ParameterSet, SignatureLen, SigningKeyLen, VerifyingKey, VerifyingKeyLen};

use super::utils;
use utils::transcoding;
use utils::typenum::Unsigned;

use crate::ffi::{SUCCESS, c_int};
use crate::{FFIError, prepare_domain_separation_prefix};

/// Secret key for signature generation.
#[derive(Debug, PartialEq)]
#[repr(transparent)]
pub struct SigningKey<P: ParameterSet> {
    sk: GenericArray<u8, <P as SigningKeyLen>::LEN>,
}

/// A seed for a ML-DSA signing key.
///
/// A `&MlDsaSeed` can be created from a `&[u8]` via the
/// `TryInto::<&MlDsaSeed>::try_into` method. The length of the seed is the same
/// for all parameter sets.
pub type MlDsaSeed = GenericArray<u8, generic_array::typenum::U32>;

pub type SecretKey<P> = SigningKey<P>;

#[cfg(feature = "rand")]
/// Generate a keypair.
pub fn keygen<P: ParameterSet>() -> Result<(SigningKey<P>, VerifyingKey<P>), FFIError> {
    let mut pk: GenericArray<u8, <P as VerifyingKeyLen>::LEN> = GenericArray::default();
    let mut sk: GenericArray<u8, <P as SigningKeyLen>::LEN> = GenericArray::default();
    let seed: MlDsaSeed = utils::rand::random_generic_byte_array();
    match unsafe { P::KEYGEN_FROM_SEED_FN(pk.as_mut_ptr(), sk.as_mut_ptr(), seed.as_ptr()) } {
        SUCCESS => Ok((SigningKey { sk }, VerifyingKey { pk })),
        error_code => Err(FFIError { code: error_code }),
    }
}

/// Generate a keypair from the given seed.
pub fn keygen_from_seed<P: ParameterSet>(
    seed: &MlDsaSeed,
) -> Result<(SigningKey<P>, VerifyingKey<P>), FFIError> {
    let mut pk: GenericArray<u8, <P as VerifyingKeyLen>::LEN> = GenericArray::default();
    let mut sk: GenericArray<u8, <P as SigningKeyLen>::LEN> = GenericArray::default();
    match unsafe { P::KEYGEN_FROM_SEED_FN(pk.as_mut_ptr(), sk.as_mut_ptr(), seed.as_ptr()) } {
        SUCCESS => Ok((SigningKey { sk }, VerifyingKey { pk })),
        error_code => Err(FFIError { code: error_code }),
    }
}

pub(super) fn pk_from_sk<P: ParameterSet>(sk: &SigningKey<P>) -> Result<VerifyingKey<P>, FFIError> {
    let mut pk: GenericArray<u8, <P as VerifyingKeyLen>::LEN> = GenericArray::default();
    match unsafe { P::PK_FROM_SK_FN(pk.as_mut_ptr(), sk.sk.as_ptr()) } {
        SUCCESS => Ok(VerifyingKey { pk }),
        error_code => Err(FFIError { code: error_code }),
    }
}

impl<P: ParameterSet> SigningKey<P> {
    #[cfg(feature = "rand")]
    /// Generate a new signing key.
    pub fn new() -> Result<Self, FFIError> {
        let (sk, _) = keygen::<P>()?;
        Ok(sk)
    }

    /// Generate a new signing key from the given seed.
    pub fn new_from_seed(seed: &MlDsaSeed) -> Result<Self, FFIError> {
        let (sk, _) = keygen_from_seed::<P>(seed)?;
        Ok(sk)
    }

    /// Attempt to use [`Self`] to sign the given `message` bytestring under the
    /// associated `context` bytestring, using the seed `seed`, returning
    /// a digital signature on success or a [`signature::Error`] if something
    /// went wrong.
    ///
    /// # Errors
    ///
    /// The main intended use case for signing errors is when communicating
    /// with external signers, e.g. cloud KMS, HSMs, or other hardware tokens.
    ///
    /// This method returns a [`signature::Error`] if the underlying FFI
    /// signature generation fails.
    pub fn try_sign_with_ctx_seeded(
        &self,
        message: &[u8],
        context: &[u8],
        seed: &MlDsaSeed,
    ) -> Result<super::Signature<P>, signature::Error> {
        type Siglen<P> = <P as SignatureLen>::LEN;
        let mut sig: GenericArray<u8, Siglen<P>> = GenericArray::default();

        let mut siglen: usize = 0;
        let ret: c_int = {
            let sk = self.sk.as_ptr();
            let prefix = prepare_domain_separation_prefix::<P>(context)?;

            unsafe {
                P::SIGN_WITH_SEED_FN(
                    sig.as_mut_ptr(),
                    &mut siglen as *mut usize,
                    message.as_ptr(),
                    message.len(),
                    prefix.as_ptr(),
                    prefix.len(),
                    seed.as_ptr(),
                    sk,
                    0, // do not use "external mu" mode
                )
            }
        };
        if ret != SUCCESS || siglen != <Siglen<P>>::USIZE {
            return Err(signature::Error::new());
        }

        // SAFETY: We assume the backend fully initialized all bytes of the
        // array, if it set the expected siglen and didn't return an error code.
        let s = super::Signature::<P> { sig };

        Ok(s)
    }
}

pub(super) const EMPTY_CTX: &[u8; 0] = &[];

#[cfg(feature = "rand")]
impl<P: ParameterSet> signature::Signer<super::Signature<P>> for SigningKey<P> {
    fn try_sign(&self, msg: &[u8]) -> Result<super::Signature<P>, signature::Error> {
        let seed = utils::rand::random_generic_byte_array();
        self.try_sign_with_ctx_seeded(msg, EMPTY_CTX, &seed)
    }
}

#[cfg(feature = "rand")]
impl<P: ParameterSet> ContextSigner<super::Signature<P>> for SigningKey<P> {
    fn try_sign_with_ctx(
        &self,
        msg: &[u8],
        ctx: &[u8],
    ) -> Result<super::Signature<P>, signature::Error> {
        let seed = utils::rand::random_generic_byte_array();
        self.try_sign_with_ctx_seeded(msg, ctx, &seed)
    }
}

impl<P: ParameterSet> SeededContextSigner<super::Signature<P>> for SigningKey<P> {
    fn try_sign_with_ctx_and_seed(
        &self,
        seed: &[u8],
        msg: &[u8],
        ctx: &[u8],
    ) -> Result<super::Signature<P>, signature::Error> {
        let seed = GenericArray::try_from_slice(seed).map_err(|_| signature::Error::new())?;
        self.try_sign_with_ctx_seeded(msg, ctx, seed)
    }
}

impl<P: ParameterSet> SeededSigner<super::Signature<P>> for SigningKey<P> {
    fn try_sign_with_seed(
        &self,
        seed: &[u8],
        msg: &[u8],
    ) -> Result<super::Signature<P>, signature::Error> {
        let seed = GenericArray::try_from_slice(seed).map_err(|_| signature::Error::new())?;
        self.try_sign_with_ctx_seeded(msg, EMPTY_CTX, seed)
    }
}

impl<P: ParameterSet> From<SigningKey<P>> for GenericArray<u8, <P as crate::SigningKeyLen>::LEN> {
    fn from(sk: SigningKey<P>) -> Self {
        sk.sk
    }
}

impl<P: ParameterSet> TryFrom<&[u8]> for SigningKey<P> {
    type Error = transcoding::TranscodingError;

    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
        // Type inference for arr due to its later use in constructing a Self (i.e. a SigningKey<P>)
        // ensures the byte slice is the correct length to be a signing key in this parameter set.
        let arr = GenericArray::try_from_slice(bytes)?;

        // In mldsa-native, the validity of a (purported) signing key is only checked when
        // attempting to derive the corresponding public key. Here we provisionally assume the bytes
        // represent a valid secret key in order to perform that derivation, but we discard the
        // result if successful.
        let provisional_sk = Self { sk: arr.clone() };
        let _ = pk_from_sk(&provisional_sk)?;

        Ok(provisional_sk)
    }
}

impl<P: ParameterSet> AsRef<[u8]> for SigningKey<P> {
    fn as_ref(&self) -> &[u8] {
        self.sk.as_ref()
    }
}

/// Sign the given message, using the provided seed for the signing process.
pub trait SeededSigner<S> {
    /// Sign the given message and return a digital signature
    fn sign_with_seed(&self, seed: &[u8], msg: &[u8]) -> S {
        self.try_sign_with_seed(seed, msg)
            .expect("signature operation failed")
    }

    /// Attempt to sign the given message, returning a digital signature on
    /// success, or an error if something went wrong.
    ///
    /// The main intended use case for signing errors is when communicating
    /// with external signers, e.g. cloud KMS, HSMs, or other hardware tokens.
    fn try_sign_with_seed(&self, seed: &[u8], msg: &[u8]) -> Result<S, signature::Error>;
}

/// Sign the given message, using the provided context string.
pub trait ContextSigner<S> {
    /// Sign the given message and return a digital signature
    fn sign_with_context(&self, msg: &[u8], ctx: &[u8]) -> S {
        self.try_sign_with_ctx(msg, ctx)
            .expect("signature operation failed")
    }

    /// Attempt to sign the given message, returning a digital signature on
    /// success, or an error if something went wrong.
    ///
    /// The main intended use case for signing errors is when communicating
    /// with external signers, e.g. cloud KMS, HSMs, or other hardware tokens.
    fn try_sign_with_ctx(&self, msg: &[u8], ctx: &[u8]) -> Result<S, signature::Error>;
}

/// Sign the given message, using the provided context string and the provided
/// seed for the signing process.
pub trait SeededContextSigner<S> {
    /// Sign the given message and return a digital signature
    fn sign_with_context_and_seed(&self, seed: &[u8], msg: &[u8], ctx: &[u8]) -> S {
        self.try_sign_with_ctx_and_seed(seed, msg, ctx)
            .expect("signature operation failed")
    }

    /// Attempt to sign the given message, returning a digital signature on
    /// success, or an error if something went wrong.
    ///
    /// The main intended use case for signing errors is when communicating
    /// with external signers, e.g. cloud KMS, HSMs, or other hardware tokens.
    fn try_sign_with_ctx_and_seed(
        &self,
        seed: &[u8],
        msg: &[u8],
        ctx: &[u8],
    ) -> Result<S, signature::Error>;
}
