//! X3DH (Extended Triple Diffie-Hellman) key agreement.
//! Follows Signal's X3DH specification.

use x25519_dalek::{PublicKey, StaticSecret};
use zeroize::Zeroize;

use crate::crypto::hkdf_derive;
use crate::errors::ProtocolError;
use crate::identity::IdentityKeyPair;
use crate::prekey::PreKeyBundle;

/// Result of initiating X3DH (Alice's side).
pub struct X3DHInitResult {
    /// The shared secret (32 bytes), used to initialize the Double Ratchet.
    pub shared_secret: [u8; 32],
    /// Alice's ephemeral public key (sent to Bob).
    pub ephemeral_public: PublicKey,
    /// Which one-time pre-key was used (if any).
    pub used_one_time_pre_key_id: Option<u32>,
}

/// Initiate X3DH key exchange (Alice's side).
///
/// Alice fetches Bob's pre-key bundle from the server, performs four DH
/// operations, and derives a shared secret.
pub fn initiate(
    our_identity: &IdentityKeyPair,
    their_bundle: &PreKeyBundle,
) -> Result<X3DHInitResult, ProtocolError> {
    // Verify the signed pre-key signature
    let their_identity = ed25519_dalek::VerifyingKey::from_bytes(
        &their_bundle.identity_key,
    )
    .map_err(|_| ProtocolError::X3DHFailed("invalid identity key".into()))?;
    their_bundle
        .signed_pre_key
        .verify(&their_identity)
        .map_err(|_| ProtocolError::X3DHFailed("signed pre-key verification failed".into()))?;

    let ephemeral_secret = StaticSecret::random_from_rng(rand::rngs::OsRng);
    let ephemeral_public = PublicKey::from(&ephemeral_secret);

    let their_spk = PublicKey::from(their_bundle.signed_pre_key.public_key);
    let their_identity_x25519 = PublicKey::from(their_bundle.identity_encryption_key);

    // DH1: our_identity_x25519 * their_signed_pre_key
    let dh1 = our_identity.encryption.diffie_hellman(&their_spk);

    // DH2: our_ephemeral * their_identity_x25519
    let dh2 = ephemeral_secret.diffie_hellman(&their_identity_x25519);

    // DH3: our_ephemeral * their_signed_pre_key
    let dh3 = ephemeral_secret.diffie_hellman(&their_spk);

    // DH4: our_ephemeral * their_one_time_pre_key (if available)
    let mut dh_concat = Vec::with_capacity(128);
    dh_concat.extend_from_slice(dh1.as_bytes());
    dh_concat.extend_from_slice(dh2.as_bytes());
    dh_concat.extend_from_slice(dh3.as_bytes());

    let used_otpk_id = if let Some(ref otpk) = their_bundle.one_time_pre_key {
        let their_otpk = PublicKey::from(otpk.public_key);
        let dh4 = ephemeral_secret.diffie_hellman(&their_otpk);
        dh_concat.extend_from_slice(dh4.as_bytes());
        Some(otpk.id)
    } else {
        None
    };

    // KDF: derive 32-byte shared secret
    let mut shared_secret = [0u8; 32];
    let derived = hkdf_derive(&dh_concat, b"", b"warzone-x3dh", 32);
    shared_secret.copy_from_slice(&derived);
    dh_concat.zeroize();

    Ok(X3DHInitResult {
        shared_secret,
        ephemeral_public,
        used_one_time_pre_key_id: used_otpk_id,
    })
}

/// Respond to X3DH key exchange (Bob's side).
///
/// Bob receives Alice's ephemeral public key and performs the same DH
/// operations to derive the same shared secret.
pub fn respond(
    our_identity: &IdentityKeyPair,
    our_signed_pre_key_secret: &StaticSecret,
    our_one_time_pre_key_secret: Option<&StaticSecret>,
    their_identity_x25519: &PublicKey,
    their_ephemeral_public: &PublicKey,
) -> Result<[u8; 32], ProtocolError> {
    let their_eph = *their_ephemeral_public;

    // DH1: our_signed_pre_key * their_identity_x25519
    let dh1 = our_signed_pre_key_secret.diffie_hellman(their_identity_x25519);

    // DH2: our_identity_x25519 * their_ephemeral
    let dh2 = our_identity.encryption.diffie_hellman(&their_eph);

    // DH3: their_ephemeral * our_signed_pre_key
    let dh3 = our_signed_pre_key_secret.diffie_hellman(&their_eph);

    let mut dh_concat = Vec::with_capacity(128);
    dh_concat.extend_from_slice(dh1.as_bytes());
    dh_concat.extend_from_slice(dh2.as_bytes());
    dh_concat.extend_from_slice(dh3.as_bytes());

    if let Some(otpk) = our_one_time_pre_key_secret {
        let dh4 = otpk.diffie_hellman(&their_eph);
        dh_concat.extend_from_slice(dh4.as_bytes());
    }

    let mut shared_secret = [0u8; 32];
    let derived = hkdf_derive(&dh_concat, b"", b"warzone-x3dh", 32);
    shared_secret.copy_from_slice(&derived);
    dh_concat.zeroize();

    Ok(shared_secret)
}


#[cfg(test)]
mod tests {
    use super::*;
    use crate::identity::Seed;
    use crate::prekey::{generate_one_time_pre_keys, generate_signed_pre_key};

    #[test]
    fn x3dh_shared_secret_matches() {
        let alice_seed = Seed::generate();
        let alice_id = alice_seed.derive_identity();

        let bob_seed = Seed::generate();
        let bob_id = bob_seed.derive_identity();

        let (bob_spk_secret, bob_spk) = generate_signed_pre_key(&bob_id, 1);
        let bob_otpks = generate_one_time_pre_keys(0, 1);
        let bob_pub = bob_id.public_identity();

        let alice_pub = alice_id.public_identity();

        let bundle = PreKeyBundle {
            identity_key: *bob_pub.signing.as_bytes(),
            identity_encryption_key: *bob_pub.encryption.as_bytes(),
            signed_pre_key: bob_spk,
            one_time_pre_key: Some(crate::prekey::OneTimePreKeyPublic {
                id: bob_otpks[0].id,
                public_key: *bob_otpks[0].public.as_bytes(),
            }),
        };

        let alice_result = initiate(&alice_id, &bundle).unwrap();
        let bob_secret = respond(
            &bob_id,
            &bob_spk_secret,
            Some(&bob_otpks[0].secret),
            &alice_pub.encryption,
            &alice_result.ephemeral_public,
        )
        .unwrap();

        assert_eq!(alice_result.shared_secret, bob_secret);
    }
}