use sequoia_openpgp as openpgp;
use openpgp::{
    crypto::{Decryptor, SessionKey, Signer, SymmetricAlgorithm},
    packet,
    types::Curve,
    types::HashAlgorithm,
    types::PublicKeyAlgorithm,
};

use sequoia_cryptoki::KeyCapability;

use crate::integration::common::*;

fn check(algo: PublicKeyAlgorithm) -> Option<PublicKeyAlgorithm> {
    if algo.is_supported() {
        Some(algo)
    } else {
        None
    }
}

#[test]
fn key_import_encryption() -> Result<()> {
    type T = Option<
            (PublicKeyAlgorithm, u8,
             packet::Key<packet::key::SecretParts, packet::key::UnspecifiedRole>)>;

    let gen_v4_rsa = |bits: usize| -> T {
        (check(PublicKeyAlgorithm::RSAEncryptSign)?, 4,
         packet::key::Key4::generate_rsa(bits)
             .expect("Can generate key").into()).into()
    };
    let gen_v6_rsa = |bits: usize| -> T {
        (check(PublicKeyAlgorithm::RSAEncryptSign)?, 6,
         packet::key::Key6::generate_rsa(bits)
             .expect("Can generate key").into()).into()
    };
    let gen_cv25519 = || -> T {
        (check(PublicKeyAlgorithm::ECDH)?, 4,
         packet::key::Key4::generate_ecc(false, Curve::Cv25519)
             .expect("Can generate key").into()).into()
    };
    let gen_v6_x25519 = || -> T {
        (check(PublicKeyAlgorithm::X25519)?, 6,
         packet::key::Key6::generate_x25519()
             .expect("Can generate key").into()).into()
    };
    let gen_v4_x25519 = || -> T {
        (check(PublicKeyAlgorithm::X25519)?, 4,
         packet::key::Key4::generate_x25519()
             .expect("Can generate key").into()).into()
    };
    let gen_v6_x448 = || -> T {
        (check(PublicKeyAlgorithm::X448)?, 6,
         packet::key::Key6::generate_x448()
             .expect("Can generate key").into()).into()
    };
    let gen_v4_x448 = || -> T {
        (check(PublicKeyAlgorithm::X448)?, 4,
         packet::key::Key4::generate_x448()
             .expect("Can generate key").into()).into()
    };
    let gen_v6_mlkem768 = || -> T {
        (check(PublicKeyAlgorithm::MLKEM768_X25519)?, 6,
         packet::key::Key6::generate_mlkem768_x25519()
             .expect("Can generate key").into()).into()
    };
    let gen_v4_mlkem768 = || -> T {
        (check(PublicKeyAlgorithm::MLKEM768_X25519)?, 4,
         packet::key::Key4::generate_mlkem768_x25519()
             .expect("Can generate key").into()).into()
    };
    let gen_v6_mlkem1024 = || -> T {
        (check(PublicKeyAlgorithm::MLKEM1024_X448)?, 6,
         packet::key::Key6::generate_mlkem1024_x448()
            .expect("Can generate key").into()).into()
    };

    let sqc = SqCryptoki::new();
    let token_label = sqc.token_init(
        sqc.allocate_slot(), None, "123456", "1234");

    for gen in &[
        Box::new(|| gen_v4_rsa(2028)) as Box<dyn Fn () -> T>,
        Box::new(|| gen_v4_rsa(3072)),
        Box::new(|| gen_v4_rsa(4096)),
        Box::new(|| gen_v6_rsa(2028)),
        Box::new(|| gen_v6_rsa(3072)),
        Box::new(|| gen_v6_rsa(4096)),
        Box::new(gen_cv25519),
        Box::new(gen_v4_x25519),
        Box::new(gen_v6_x25519),
        Box::new(gen_v6_x448),
        Box::new(gen_v4_x448),
        Box::new(gen_v6_mlkem768),
        Box::new(gen_v4_mlkem768),
        Box::new(gen_v6_mlkem1024),
    ] {
        let Some((algo, version, key)) = gen() else { continue; };

        let cryptoki_algo = sequoia_cryptoki::PublicKeyAlgorithm::try_from(&key)
            .expect("Known combination");
        if ! SqCryptoki::algo_supported(&cryptoki_algo) {
            eprintln!("{} not supported by backend", cryptoki_algo);
            continue;
        }

        eprintln!("Importing a {} key", algo);

        // Check that what we generated is what we expect...
        assert_eq!(key.pk_algo(), algo);
        assert_eq!(key.version(), version);

        sqc.token_reinit(&token_label, "123456", "1234");

        let key_handle = sqc.import(
            &token_label, &key, KeyCapability::Encryption, None, "1234");

        let exported_key = key_handle.to_openpgp_key()?;
        assert_eq!(key.parts_as_public(), exported_key.role_as_unspecified());

        // Check that we can round trip some random data.
        let algo = SymmetricAlgorithm::AES128;
        let sk = SessionKey::new(algo.key_size()?)?;
        let ciphertext = exported_key.encrypt(&sk)?;

        let mut decryptor = sequoia_cryptoki::Decryptor::new(
            CancelPrompt::default(), &key_handle, Some("1234".to_string().into()))?;
        let sk_recovered = decryptor.decrypt(&ciphertext, None)?;
        assert_eq!(sk, sk_recovered);
    }

    Ok(())
}

#[test]
fn key_import_signature() -> Result<()> {
    type T = Option<
            (PublicKeyAlgorithm, u8,
             packet::Key<packet::key::SecretParts, packet::key::UnspecifiedRole>)>;

    let gen_v4_rsa = |bits: usize| -> T {
        (check(PublicKeyAlgorithm::RSAEncryptSign)?, 4,
         packet::key::Key4::generate_rsa(bits)
             .expect("Can generate key").into()).into()
    };
    let gen_v6_rsa = |bits: usize| -> T {
        (check(PublicKeyAlgorithm::RSAEncryptSign)?, 6,
         packet::key::Key6::generate_rsa(bits)
             .expect("Can generate key").into()).into()
    };
    let gen_legecy_ed25519 = || -> T {
        (check(PublicKeyAlgorithm::EdDSA)?, 4,
         packet::key::Key4::generate_ecc(true, Curve::Ed25519)
             .expect("Can generate key").into()).into()
    };
    let gen_v6_ed25519 = || -> T {
        (check(PublicKeyAlgorithm::Ed25519)?, 6,
         packet::key::Key6::generate_ed25519()
             .expect("Can generate key").into()).into()
    };
    let gen_v4_ed25519 = || -> T {
        (check(PublicKeyAlgorithm::Ed25519)?, 4,
         packet::key::Key4::generate_ed25519()
             .expect("Can generate key").into()).into()
    };
    let gen_v6_ed448 = || -> T {
        (check(PublicKeyAlgorithm::Ed448)?, 6,
         packet::key::Key6::generate_ed448()
             .expect("Can generate key").into()).into()
    };
    let gen_v4_ed448 = || -> T {
        (check(PublicKeyAlgorithm::Ed448)?, 4,
         packet::key::Key4::generate_ed448()
             .expect("Can generate key").into()).into()
    };
    let gen_mldsa65 = || -> T {
        (check(PublicKeyAlgorithm::MLDSA65_Ed25519)?, 6,
         packet::key::Key6::generate_mldsa65_ed25519()
             .expect("Can generate key").into()).into()
    };
    let gen_mldsa87 = || -> T {
        (check(PublicKeyAlgorithm::MLDSA87_Ed448)?, 6,
         packet::key::Key6::generate_mldsa87_ed448()
             .expect("Can generate key").into()).into()
    };
    let gen_slh_dsa_128s = || -> T {
        (check(PublicKeyAlgorithm::SLHDSA128s)?, 6,
         packet::key::Key6::generate_slhdsa128s()
             .expect("Can generate key").into()).into()
    };
    let gen_slh_dsa_128f = || -> T {
        (check(PublicKeyAlgorithm::SLHDSA128f)?, 6,
         packet::key::Key6::generate_slhdsa128f()
             .expect("Can generate key").into()).into()
    };
    let gen_slh_dsa_256s = || -> T {
        (check(PublicKeyAlgorithm::SLHDSA256s)?, 6,
         packet::key::Key6::generate_slhdsa256s()
             .expect("Can generate key").into()).into()
    };

    let sqc = SqCryptoki::new();
    let token_label = sqc.token_init(
        sqc.allocate_slot(), None, "123456", "1234");

    for gen in [
        Box::new(|| gen_v4_rsa(2028)) as Box<dyn Fn () -> T>,
        Box::new(|| gen_v4_rsa(3072)),
        Box::new(|| gen_v4_rsa(4096)),
        Box::new(|| gen_v6_rsa(2028)),
        Box::new(|| gen_v6_rsa(3072)),
        Box::new(|| gen_v6_rsa(4096)),
        Box::new(gen_legecy_ed25519),
        Box::new(gen_v6_ed25519),
        Box::new(gen_v4_ed25519),
        Box::new(gen_v6_ed448),
        Box::new(gen_v4_ed448),
        Box::new(gen_slh_dsa_128s),
        Box::new(gen_slh_dsa_128f),
        Box::new(gen_slh_dsa_256s),
        Box::new(gen_mldsa65),
        Box::new(gen_mldsa87),
    ] {
        let Some((algo, version, key)) = gen() else { continue; };

        let cryptoki_algo = sequoia_cryptoki::PublicKeyAlgorithm::try_from(&key)
            .expect("Known combination");
        if ! SqCryptoki::algo_supported(&cryptoki_algo) {
            eprintln!("{} not supported by backend", cryptoki_algo);
            continue;
        }

        eprintln!("Importing a v{} {} key", version, algo);
        assert_eq!(key.pk_algo(), algo);
        assert_eq!(key.version(), version);

        sqc.token_reinit(&token_label, "123456", "1234");

        // Import the key.
        let imported_key_handle = sqc.import(
            &token_label, &key, KeyCapability::Signing, Some("a-label"), "1234");

        // Make sure the returned public key matches the public key that
        // we imported.
        let exported_key = imported_key_handle.to_openpgp_key()
            .expect("can export cert");
        assert_eq!(exported_key, key.clone().into());

        // Sign and verify a message to check that the imported private
        // key is correct.
        let mut signer = sequoia_cryptoki::Signer::new(
            CancelPrompt::default(), &imported_key_handle,
            Some("1234".to_string().into()))?;

        // Our "message".
        let hash_algo = HashAlgorithm::SHA256;
        let digest = vec![0; hash_algo.digest_size().expect("algo is supported")];

        let sig = signer.sign(hash_algo, &digest)
            .expect("can sign");

        key.verify(&sig, hash_algo, &digest)
            .expect("can verify signature");
    }

    Ok(())
}
