pub use signature;

use crate::ffi;

/// Parameter sets and related constants.
pub mod parameter_sets;

mod signature_encoding;
mod signing_key;
mod verifying_key;

pub use signature_encoding::*;
pub use signing_key::*;
pub use verifying_key::*;

pub(crate) mod utils;
pub use transcoding::*;
pub use utils::transcoding;

// Trait that holds the C function pointers for the underlying implementation.
trait FFIFunctionPointers {
    const KEYGEN_FROM_SEED_FN: unsafe extern "C" fn(
        pk: *mut u8,
        sk: *mut u8,
        seed: *const u8,
    ) -> ffi::c_int;
    const PK_FROM_SK_FN: unsafe extern "C" fn(pk: *mut u8, sk: *const u8) -> ffi::c_int;
    const SIGN_WITH_SEED_FN: unsafe extern "C" fn(
        sig: *mut u8,
        siglen: *mut usize,
        m: *const u8,
        mlen: usize,
        pre: *const u8,
        prelen: usize,
        rnd: *const u8,
        sk: *const u8,
        externalmu: ffi::c_int,
    ) -> ffi::c_int;
    const VERIFY_FN: unsafe extern "C" fn(
        sig: *const u8,
        siglen: usize,
        m: *const u8,
        mlen: usize,
        ctx: *const u8,
        ctxlen: usize,
        pk: *const u8,
    ) -> ffi::c_int;
    const PREPARE_DOMAIN_SEPARATION_PREFIX_FN: unsafe extern "C" fn(
        prefix: *mut u8,
        ph: *const u8,
        phlen: usize,
        ctx: *const u8,
        ctxlen: usize,
        hashalg: ffi::c_int,
    ) -> usize;
}

/// Trait for types that provide the length of the secret (signing) key.
pub trait SigningKeyLen {
    /// Length of the secret (signing) key in bytes.
    ///
    /// Same as `Self::SECRET_KEY_LEN`.
    const SIGNING_KEY_LEN: usize;

    /// Length of the secret (signing) key in bytes.
    ///
    /// Same as `Self::SIGNING_KEY_LEN`.
    const SECRET_KEY_LEN: usize = Self::SIGNING_KEY_LEN;

    /// `Self::SIGNING_KEY_LEN` as a type.
    type LEN: generic_array::ArrayLength;
}

/// Trait for types that provide the length of the public (verifying) key.
pub trait VerifyingKeyLen {
    /// Length of the public (verifying) key in bytes.
    ///
    /// Same as `Self::PUBLIC_KEY_LEN`.
    const VERIFYING_KEY_LEN: usize;

    /// Length of the public (verifying) key in bytes.
    ///
    /// Same as `Self::VERIFYING_KEY_LEN`.
    const PUBLIC_KEY_LEN: usize = Self::VERIFYING_KEY_LEN;

    /// `Self::VERIFYING_KEY_LEN` as a type.
    type LEN: generic_array::ArrayLength;
}

/// Trait for types that provide the length of the signature.
pub trait SignatureLen {
    /// Length of the signature in bytes.
    const SIGNATURE_LEN: usize;

    /// `Self::SIGNATURE_LEN` as a type.
    type LEN: generic_array::ArrayLength;
}

/// Trait for types that provide the length of the seed.
pub trait SeedLen {
    /// Length of the seed in bytes.
    const SEED_LEN: usize;

    /// `Self::SEED_LEN` as a type.
    type LEN: generic_array::ArrayLength;
}

/// Trait implemented by each ML-DSA parameter set.
#[allow(private_bounds)]
pub trait ParameterSet:
    FFIFunctionPointers
    + SigningKeyLen
    + VerifyingKeyLen
    + SignatureLen
    + SeedLen
    + PartialEq
    + Eq
    + Clone
    + core::fmt::Debug
{
    /// Human-readable name for parameter set, matching the FIPS-204 designations
    const NAME: &'static str;
}

/// Error returned by the backend implementation.
#[derive(Debug)]
pub struct FFIError {
    /// Error code.
    pub code: ffi::c_int,
}

// needed in order to impl Error
impl std::fmt::Display for FFIError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{self:?}")
    }
}

impl std::error::Error for FFIError {}

use signature::Keypair;
impl<P: ParameterSet> Keypair for SigningKey<P> {
    type VerifyingKey = VerifyingKey<P>;

    fn verifying_key(&self) -> Self::VerifyingKey {
        match pk_from_sk(self) {
            Ok(vk) => vk,
            Err(_) => unreachable!(),
        }
    }
}

#[cfg(test)]
mod tests {
    use signature::Signer;

    use super::*;

    fn test_sign_verify_tampered_generic<P: ParameterSet>() {
        let seed = utils::rand::random_generic_byte_array();
        let (sk, vk) = signing_key::keygen_from_seed::<P>(&seed).expect("Keygen failed");
        std::println!("{sk:?}");
        std::println!("{vk:?}");
        let msg = b"Hello world!";
        let sig = sk.try_sign(msg).unwrap();
        let other_msg = b"Hello Earth!";
        assert!(vk.verify(other_msg, &sig).is_err());
    }

    fn test_sign_verify_generic<P: ParameterSet>() {
        let seed = utils::rand::random_generic_byte_array();
        let (sk, vk) = signing_key::keygen_from_seed::<P>(&seed).expect("Keygen failed");
        std::println!("{sk:?}");
        std::println!("{vk:?}");
        let msg = b"Hello world!";
        let sig = sk.try_sign(msg).unwrap();
        vk.verify(msg, &sig).unwrap();
    }

    fn test_recover_verifying_key_generic<P: ParameterSet>() {
        let seed = utils::rand::random_generic_byte_array();
        let (sk, vk) = signing_key::keygen_from_seed::<P>(&seed).expect("Keygen failed");
        let vk2 = sk.verifying_key();
        assert_eq!(vk, vk2);
    }

    fn test_keygen_from_seed_generic<P: ParameterSet>() {
        let seed = utils::rand::random_generic_byte_array();
        let (_sk, _vk) = signing_key::keygen_from_seed::<P>(&seed).expect("Keygen failed");
    }

    #[test]
    fn test_sign_verify() {
        test_sign_verify_generic::<parameter_sets::ML_DSA_44>();
        test_sign_verify_generic::<parameter_sets::ML_DSA_65>();
        test_sign_verify_generic::<parameter_sets::ML_DSA_87>();
    }

    #[test]
    fn test_sign_verify_tampered() {
        test_sign_verify_tampered_generic::<parameter_sets::ML_DSA_44>();
        test_sign_verify_tampered_generic::<parameter_sets::ML_DSA_65>();
        test_sign_verify_tampered_generic::<parameter_sets::ML_DSA_87>();
    }

    #[test]
    fn test_recover_verifying_key() {
        test_recover_verifying_key_generic::<parameter_sets::ML_DSA_44>();
        test_recover_verifying_key_generic::<parameter_sets::ML_DSA_65>();
        test_recover_verifying_key_generic::<parameter_sets::ML_DSA_87>();
    }

    #[test]
    fn test_keygen_from_seed() {
        test_keygen_from_seed_generic::<parameter_sets::ML_DSA_44>();
        test_keygen_from_seed_generic::<parameter_sets::ML_DSA_65>();
        test_keygen_from_seed_generic::<parameter_sets::ML_DSA_87>();
    }
}
