//! Post-quantum cryptography example — ML-KEM, ML-DSA, and SLH-DSA.
//!
//! Demonstrates the three NIST PQC algorithm families available in OpenSSL 3.5+:
//!
//! - **ML-KEM** (FIPS 203, formerly CRYSTALS-Kyber) — key encapsulation
//! - **ML-DSA** (FIPS 204, formerly CRYSTALS-Dilithium) — lattice signatures
//! - **SLH-DSA** (FIPS 205, formerly SPHINCS+) — hash-based signatures
//!
//! Availability is detected at build time via cfg flags set by the build system:
//! `ossl_mlkem`, `ossl_mldsa`, `ossl_slhdsa`.  The example compiles and runs
//! correctly on OpenSSL builds that have any subset of these algorithms.
//!
//! Run with: cargo run --example pqc -p native-ossl

fn main() -> Result<(), Box<dyn std::error::Error>> {
    #[cfg(all(ossl320, ossl_mlkem))]
    ml_kem()?;

    #[cfg(ossl_mldsa)]
    ml_dsa()?;

    #[cfg(ossl_slhdsa)]
    slh_dsa()?;

    // Printed only when none of the three algorithm families is present.
    #[cfg(not(any(all(ossl320, ossl_mlkem), ossl_mldsa, ossl_slhdsa)))]
    {
        println!("No PQC algorithms available in this OpenSSL build.");
        println!("Requires OpenSSL 3.5+ compiled with ML-KEM, ML-DSA, or SLH-DSA support.");
    }

    Ok(())
}

// ── ML-KEM — Key Encapsulation Mechanism (FIPS 203) ──────────────────────────
//
// ML-KEM-768 provides ~128 bits of post-quantum security (Category 3).
// The encapsulate/decapsulate cycle yields a shared secret that both parties
// use as key material without ever transmitting it.

#[cfg(all(ossl320, ossl_mlkem))]
fn ml_kem() -> Result<(), Box<dyn std::error::Error>> {
    use native_ossl::pkey::{DecapCtx, EncapCtx, KeygenCtx, Pkey, Public};

    println!("── ML-KEM-768 (FIPS 203) ──────────────────────────────────────");

    // Generate a key pair (recipient side).
    let mut kgen = KeygenCtx::new(c"ML-KEM-768")?;
    let priv_key = kgen.generate()?;
    let pub_key = Pkey::<Public>::from(priv_key.clone());
    println!("Key size: {} bits", priv_key.bits());

    // Sender: encapsulate against the recipient's public key.
    // Produces a wrapped key (ciphertext) and a shared secret.
    let mut enc = EncapCtx::new(&pub_key)?;
    let result = enc.encapsulate()?;
    println!("Wrapped key:   {} bytes", result.wrapped_key.len());
    println!("Shared secret: {} bytes", result.shared_secret.len());
    println!(
        "Sender's ss:   {}",
        native_ossl::util::hex_encode(&result.shared_secret)
    );

    // Recipient: decapsulate using the private key and the received wrapped key.
    let mut dec = DecapCtx::new(&priv_key)?;
    let shared = dec.decapsulate(&result.wrapped_key)?;
    println!("Receiver's ss: {}", native_ossl::util::hex_encode(&shared));

    assert_eq!(result.shared_secret, shared, "shared secrets must match");
    println!("ML-KEM-768 key agreement: OK\n");

    Ok(())
}

// ── ML-DSA — Module Lattice Signature (FIPS 204) ─────────────────────────────
//
// ML-DSA-65 provides ~128 bits of post-quantum security (Category 3).
// In OpenSSL 3.5 the built-in implementation uses the DigestSign path
// (Signer with digest=None); MessageSigner will be used once the provider
// exposes OSSL_FUNC_SIGNATURE_SIGN_MESSAGE_FINAL.

#[cfg(ossl_mldsa)]
fn ml_dsa() -> Result<(), Box<dyn std::error::Error>> {
    use native_ossl::pkey::{KeygenCtx, Pkey, Public, SignInit, Signer, Verifier};

    println!("── ML-DSA-65 (FIPS 204) ───────────────────────────────────────");

    // Generate a key pair.
    let mut kgen = KeygenCtx::new(c"ML-DSA-65")?;
    let priv_key = kgen.generate()?;
    let pub_key = Pkey::<Public>::from(priv_key.clone());
    println!("Key size: {} bits", priv_key.bits());

    let message = b"Post-quantum signature over this message";

    // Sign — ML-DSA does not use an external hash (digest = None).
    let init = SignInit {
        digest: None,
        params: None,
    };
    let mut signer = Signer::new(&priv_key, &init)?;
    let sig = signer.sign_oneshot(message)?;
    println!("Signature: {} bytes", sig.len());

    // Verify.
    let mut verifier = Verifier::new(&pub_key, &init)?;
    let valid = verifier.verify_oneshot(message, &sig)?;
    assert!(valid, "signature must verify");
    println!("ML-DSA-65 sign/verify: OK");

    // Tampered message must fail.
    let mut bad = message.to_vec();
    bad[0] ^= 1;
    let mut verifier2 = Verifier::new(&pub_key, &init)?;
    let bad_result = verifier2.verify_oneshot(&bad, &sig)?;
    assert!(!bad_result, "tampered message must not verify");
    println!("ML-DSA-65 tampered message rejected: OK\n");

    Ok(())
}

// ── SLH-DSA — Stateless Hash-Based Signature (FIPS 205) ──────────────────────
//
// SLH-DSA-SHA2-128s is the smallest NIST parameter set: ~128 bits of
// post-quantum security, smallest signatures (7856 bytes), slowest keygen.
// Use SLH-DSA-SHA2-128f for faster signing at the cost of larger signatures.

#[cfg(ossl_slhdsa)]
fn slh_dsa() -> Result<(), Box<dyn std::error::Error>> {
    use native_ossl::pkey::{KeygenCtx, Pkey, Public, SignInit, Signer, Verifier};

    println!("── SLH-DSA-SHA2-128s (FIPS 205) ───────────────────────────────");

    // Generate a key pair.
    let mut kgen = KeygenCtx::new(c"SLH-DSA-SHA2-128s")?;
    let priv_key = kgen.generate()?;
    let pub_key = Pkey::<Public>::from(priv_key.clone());
    println!("Key size: {} bits", priv_key.bits());

    let message = b"Hash-based post-quantum signature";

    // Sign — SLH-DSA is also a one-shot algorithm with no external digest.
    let init = SignInit {
        digest: None,
        params: None,
    };
    let mut signer = Signer::new(&priv_key, &init)?;
    let sig = signer.sign_oneshot(message)?;
    println!("Signature: {} bytes", sig.len());

    // Verify.
    let mut verifier = Verifier::new(&pub_key, &init)?;
    let valid = verifier.verify_oneshot(message, &sig)?;
    assert!(valid, "signature must verify");
    println!("SLH-DSA-SHA2-128s sign/verify: OK");

    // Tampered message must fail.
    let mut bad = message.to_vec();
    bad[0] ^= 1;
    let mut verifier2 = Verifier::new(&pub_key, &init)?;
    let bad_result = verifier2.verify_oneshot(&bad, &sig)?;
    assert!(!bad_result, "tampered message must not verify");
    println!("SLH-DSA-SHA2-128s tampered message rejected: OK\n");

    Ok(())
}
