use std::io::Write;

use sequoia_openpgp as openpgp;
use openpgp::{
    Packet,
    Profile,
    packet::{Any, Signature},
    parse::Parse,
    serialize::stream::{Message, Signer},
    crypto::{Decryptor, SessionKey, SymmetricAlgorithm},
};

use cryptoki::session::Session;
use cryptoki::session::UserType;
use cryptoki::types::AuthPin;

use sequoia_cryptoki::Key;
use sequoia_cryptoki::PublicKeyAlgorithm;
use sequoia_cryptoki::Token;
use sequoia_cryptoki::TokenSelector;

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

#[test]
fn key_generate_signing() -> Result<()> {
    type T = Key;

    let gen_v4_rsa = |sqc: &SqCryptoki, session: Session, token: Token, bits: usize| -> T {
        let key = sqc.ctx
            .generate_rsa_for_signing(
                &session, token.clone(), Profile::RFC4880,
                bits, None, None)
            .expect(&format!("Can generate a RSA key"));
        key
    };
    let gen_v6_rsa = |sqc: &SqCryptoki, session: Session, token: Token, bits: usize| -> T {
        let key = sqc.ctx
            .generate_rsa_for_signing(
                &session, token.clone(), Profile::RFC9580,
                bits, None, None)
            .expect(&format!("Can generate a RSA key"));
        key
    };
    let gen_v4_eddsa = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        let key = sqc.ctx
            .generate_eddsa(
                &session, token.clone(),
                None, None)
            .expect(&format!("Can generate a v4 EdDSA key"));
        key
    };
    let gen_v6_ed25519 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        let key = sqc.ctx
            .generate_ed25519(
                &session, token.clone(), Profile::RFC9580,
                None, None)
            .expect(&format!("Can generate a v6 Ed25519 key"));
        key
    };
    let gen_v4_ed25519 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        let key = sqc.ctx
            .generate_ed25519(
                &session, token.clone(), Profile::RFC4880,
                None, None)
            .expect(&format!("Can generate a v4 Ed25519 key"));
        key
    };
    let gen_v6_ed448 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        let key = sqc.ctx
            .generate_ed448(
                &session, token.clone(), Profile::RFC9580,
                None, None)
            .expect(&format!("Can generate a v6 Ed448 key"));
        key
    };
    let gen_v4_ed448 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        let key = sqc.ctx
            .generate_ed448(
                &session, token.clone(), Profile::RFC4880,
                None, None)
            .expect(&format!("Can generate a v4 Ed448 key"));
        key
    };
    let gen_mldsa65_ed25519 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        let key = sqc.ctx
            .generate_mldsa65_ed25519(
                &session, token.clone(), Profile::RFC9580,
                None, None)
            .expect(&format!("Can generate a v6 MLDSA65_Ed25519 key"));
        key
    };
    let gen_mldsa87_ed448 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        let key = sqc.ctx
            .generate_mldsa87_ed448(
                &session, token.clone(), Profile::RFC9580,
                None, None)
            .expect(&format!("Can generate a v6 MLDSA87_Ed448 key"));
        key
    };
    let gen_slhdsa128s = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        let key = sqc.ctx
            .generate_slh_dsa_128s(
                &session, token.clone(), Profile::RFC9580,
                None, None)
            .expect(&format!("Can generate a v6 SLHDSA128s key"));
        key
    };
    let gen_slhdsa128f = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        let key = sqc.ctx
            .generate_slh_dsa_128f(
                &session, token.clone(), Profile::RFC9580,
                None, None)
            .expect(&format!("Can generate a v6 SLHDSA128f key"));
        key
    };
    let gen_slhdsa256s = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        let key = sqc.ctx
            .generate_slh_dsa_256s(
                &session, token.clone(), Profile::RFC9580,
                None, None)
            .expect(&format!("Can generate a v6 SLHDSA256s key"));
        key
    };

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

    for (algo, expected_version, gen) in [
        (PublicKeyAlgorithm::RSA, 4,
         Box::new(|sqc, session, token| gen_v4_rsa(sqc, session, token, 3072))
            as Box<dyn Fn (_, _, _) -> T>),
        (PublicKeyAlgorithm::RSA, 6,
         Box::new(|sqc, session, token| gen_v6_rsa(sqc, session, token, 3072))),
        (PublicKeyAlgorithm::EdDSA, 4, Box::new(gen_v4_eddsa)),
        (PublicKeyAlgorithm::Ed25519, 6, Box::new(gen_v6_ed25519)),
        (PublicKeyAlgorithm::Ed25519, 4, Box::new(gen_v4_ed25519)),
        (PublicKeyAlgorithm::Ed448, 6, Box::new(gen_v6_ed448)),
        (PublicKeyAlgorithm::Ed448, 4, Box::new(gen_v4_ed448)),
        (PublicKeyAlgorithm::MLDSA65_Ed25519, 6, Box::new(gen_mldsa65_ed25519)),
        (PublicKeyAlgorithm::MLDSA87_Ed448, 6, Box::new(gen_mldsa87_ed448)),
        (PublicKeyAlgorithm::SLHDSA128s, 6, Box::new(gen_slhdsa128s)),
        (PublicKeyAlgorithm::SLHDSA128f, 6, Box::new(gen_slhdsa128f)),
        (PublicKeyAlgorithm::SLHDSA256s, 6, Box::new(gen_slhdsa256s)),
    ]
    {
        eprintln!("Checking {}.", algo);
        if ! SqCryptoki::algo_supported(&algo) {
            eprintln!("Token does not support {}, skipping test.", algo);
            continue;
        }

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

        let token = sqc.ctx.select_one_token(&TokenSelector {
            backend: None,
            slot: None,
            token_label: Some(token_label.to_string()),
            token_serial_number: None,
        })?;

        let session = token.open_rw_user_session("1234")?;

        let key = gen(&sqc, session, token);

        let pgp = key.to_openpgp_key()?;

        assert_eq!(pgp.version(), expected_version);

        assert_eq!(pgp.pk_algo(), algo.clone().into());

        let signer = sequoia_cryptoki::Signer::new(
            CancelPrompt::default(), &key, Some("1234".to_string().into()))?;

        let mut sink = vec![];
        let message = Message::new(&mut sink);
        let mut signer = Signer::new(message, signer)?.detached().build()?;
        signer.write_all(b"Make it so, number one!")?;
        signer.finalize()?;

        // Now check the signature.
        if openpgp::types::PublicKeyAlgorithm::from(algo.clone()).is_supported() {
            let sig: Signature =
                Packet::from_bytes(&sink)?.downcast().expect("a sig");

            sig.verify_message(&pgp, b"Make it so, number one!").unwrap();
        } else {
            eprintln!("Can't check signature: crypto backend does not support {}",
                      algo);
        }
    }

    Ok(())
}

#[test]
fn key_generate_encryption() -> Result<()> {
    type T = Key;

    let gen_v4_rsa = |sqc: &SqCryptoki, session: Session, token: Token, bits: usize| -> T {
        sqc.ctx
            .generate_rsa_for_encryption(
                &session, token.clone(), Profile::RFC4880,
                bits, None, None)
            .expect(&format!("Can generate a RSA key"))
    };
    let gen_v6_rsa = |sqc: &SqCryptoki, session: Session, token: Token, bits: usize| -> T {
        sqc.ctx
            .generate_rsa_for_encryption(
                &session, token.clone(), Profile::RFC9580,
                bits, None, None)
            .expect(&format!("Can generate a RSA key"))
    };
    let gen_v6_x25519 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        sqc.ctx
            .generate_x25519(
                &session, token.clone(), Profile::RFC9580,
                None, None)
            .expect(&format!("Can generate a v6 X25519 key"))
    };
    let gen_v4_x25519 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        sqc.ctx
            .generate_x25519(
                &session, token.clone(), Profile::RFC4880,
                None, None)
            .expect(&format!("Can generate a v4 X25519 key"))
    };
    let gen_v6_x448 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        sqc.ctx
            .generate_x448(
                &session, token.clone(), Profile::RFC9580,
                None, None)
            .expect(&format!("Can generate a v6 X448 key"))
    };
    let gen_v4_x448 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        sqc.ctx
            .generate_x448(
                &session, token.clone(), Profile::RFC4880,
                None, None)
            .expect(&format!("Can generate a v4 X448 key"))
    };
    let gen_v6_mlkem768_x25519 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        sqc.ctx
            .generate_mlkem768_x25519(
                &session, token.clone(), Profile::RFC9580,
                None, None)
            .expect(&format!("Can generate a v6 MLKEM768_X25519 key"))
    };
    let gen_v4_mlkem768_x25519 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        sqc.ctx
            .generate_mlkem768_x25519(
                &session, token.clone(), Profile::RFC4880,
                None, None)
            .expect(&format!("Can generate a v4 MLKEM768_X25519 key"))
    };
    let gen_v6_mlkem1024_x448 = |sqc: &SqCryptoki, session: Session, token: Token| -> T {
        sqc.ctx
            .generate_mlkem1024_x448(
                &session, token.clone(), Profile::RFC9580,
                None, None)
            .expect(&format!("Can generate a v6 MLKEM1024_X448 key"))
    };

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

    for (algo, expected_version, gen) in [
        (PublicKeyAlgorithm::RSA, 4,
         Box::new(|sqc, session, token| gen_v4_rsa(sqc, session, token, 3072))
            as Box<dyn Fn (_, _, _) -> T>),
        (PublicKeyAlgorithm::RSA, 6,
         Box::new(|sqc, session, token| gen_v6_rsa(sqc, session, token, 3072))),
        (PublicKeyAlgorithm::X25519, 6, Box::new(gen_v6_x25519)),
        (PublicKeyAlgorithm::X25519, 4, Box::new(gen_v4_x25519)),
        (PublicKeyAlgorithm::X448, 6, Box::new(gen_v6_x448)),
        (PublicKeyAlgorithm::X448, 4, Box::new(gen_v4_x448)),
        (PublicKeyAlgorithm::MLKEM768_X25519, 6, Box::new(gen_v6_mlkem768_x25519)),
        (PublicKeyAlgorithm::MLKEM768_X25519, 4, Box::new(gen_v4_mlkem768_x25519)),
        (PublicKeyAlgorithm::MLKEM1024_X448, 6, Box::new(gen_v6_mlkem1024_x448)),
    ]
    {
        eprintln!("Checking {}.", algo);
        if ! SqCryptoki::algo_supported(&algo) {
            eprintln!("Token does not support {}, skipping test.", algo);
            continue;
        }

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

        let token = sqc.ctx.select_one_token(&TokenSelector {
            backend: None,
            slot: None,
            token_label: Some(token_label.to_string()),
            token_serial_number: None,
        })?;

        let session = token.open_rw_user_session("1234")?;

        let key = gen(&sqc, session, token);

        let pgp = key.to_openpgp_key()?;

        assert_eq!(pgp.version(), expected_version);

        assert_eq!(pgp.pk_algo(), algo.clone().into());

        if openpgp::types::PublicKeyAlgorithm::from(algo.clone()).is_supported() {
            let algo = SymmetricAlgorithm::AES128;
            let sk = SessionKey::new(algo.key_size()?)?;
            let ciphertext = pgp.encrypt(&sk)?;

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

    Ok(())
}

#[test]
fn key_generate_legacy_cv25519() -> Result<()> {
    if ! SqCryptoki::algo_supported(&PublicKeyAlgorithm::Cv25519) {
        eprintln!("{} not supported, skipping", PublicKeyAlgorithm::Cv25519);
        return Ok(());
    }

    let sqc = SqCryptoki::new();

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

    use openpgp::types::{HashAlgorithm, SymmetricAlgorithm};

    for (ok, hash_algo, sym_algo) in &[
        (true, HashAlgorithm::SHA256, SymmetricAlgorithm::AES128),
        (true, HashAlgorithm::SHA384, SymmetricAlgorithm::AES128),
        (true, HashAlgorithm::SHA512, SymmetricAlgorithm::AES128),
        (true, HashAlgorithm::SHA256, SymmetricAlgorithm::AES192),
        (true, HashAlgorithm::SHA384, SymmetricAlgorithm::AES192),
        (true, HashAlgorithm::SHA512, SymmetricAlgorithm::AES192),
        (true, HashAlgorithm::SHA256, SymmetricAlgorithm::AES256),
        (true, HashAlgorithm::SHA384, SymmetricAlgorithm::AES256),
        (true, HashAlgorithm::SHA512, SymmetricAlgorithm::AES256),
        (false, HashAlgorithm::SHA224, SymmetricAlgorithm::AES256),
        (false, HashAlgorithm::SHA1, SymmetricAlgorithm::AES256),
    ] {
        sqc.token_reinit(&token_label, "123456", "1234");

        let result = sqc.try_legacy_cv25519_key_generate(
            &token_label, Some(hash_algo.clone()), Some(sym_algo.clone()),
            None, "1234");

        let key = match (ok, result) {
            (true, Ok(key)) => key,
            (false, Err(_)) => continue,
            (true, Err(err)) => {
                panic!("Should have succeeded, but failed to generate key \
                        using {} and {}: {}",
                       hash_algo, sym_algo, err);
            }
            (false, Ok(_)) => {
                panic!("Should have failed, but generated a key \
                        using {} and {}",
                       hash_algo, sym_algo);
            }
        };

        let pgp = key.to_openpgp_key()?;

        let sk = SessionKey::new(sym_algo.key_size()?)?;
        let ciphertext = pgp.encrypt(&sk)?;

        let mut decryptor = sequoia_cryptoki::Decryptor::new(
            CancelPrompt::default(), &key, Some("1234".to_string().into()))?;
        let sk_recovered = decryptor.decrypt(&ciphertext, None)?;
        assert_eq!(sk, sk_recovered,
                   "\nExpected: {}\n     Got: {}",
                   sk.iter()
                       .map(|k| format!("{:02X}", k))
                       .collect::<Vec<_>>()
                       .join(""),
                   sk_recovered.iter()
                       .map(|k| format!("{:02X}", k))
                       .collect::<Vec<_>>()
                       .join(""));
    }

    Ok(())
}
