//! WASM-safe cryptographic primitives for pairing and chat protocols.
//!
//! AES-128-GCM / AES-256-GCM encryption, HKDF-SHA256 key derivation,
//! HMAC-SHA256 topic derivation, P-256 ECDH key agreement,
//! SR25519→P-256 key derivation, and SCALE-encoded handshake construction.
//! Used by `pairing.rs` (AES-128 path) and `chat-v1` (AES-256 / P-256 path).

use aes_gcm::aead::Aead;
use aes_gcm::{Aes128Gcm, KeyInit, Nonce};
use hkdf::Hkdf;
use hmac::{Hmac, Mac};
use sha2::Sha256;

/// Generalized HKDF-SHA256 key derivation (RFC 5869).
///
/// - `ikm`: input keying material (e.g., raw ECDH shared secret)
/// - `salt`: optional salt (empty slice = zero-filled HashLen per RFC 5869)
/// - `info`: context string for domain separation
/// - `output`: filled with derived key material (max 255 * 32 = 8160 bytes)
///
/// Returns `Err` if `output` exceeds the HKDF-SHA256 maximum output length.
pub fn hkdf_sha256(ikm: &[u8], salt: &[u8], info: &[u8], output: &mut [u8]) -> Result<(), String> {
    let hk = Hkdf::<Sha256>::new(if salt.is_empty() { None } else { Some(salt) }, ikm);
    hk.expand(info, output).map_err(|_| {
        format!(
            "HKDF output length {} exceeds max {} (255 * 32)",
            output.len(),
            255 * 32
        )
    })
}

/// Derive the 16-byte AES-128 session key from a P-256 ECDH shared secret.
///
/// Uses HKDF-SHA256 with empty salt and info=b"polkadot-pairing".
///
/// # AES-128 vs AES-256
/// This function deliberately produces a 16-byte (AES-128) key for the
/// **pairing channel** — the short-lived, QR-code-initiated link between a
/// desktop host and a companion papp. AES-128 was chosen for backward
/// compatibility with host-papp, which negotiated this key size before AES-256
/// support was added. Both key sizes provide adequate security for the pairing
/// use case (single-session, ephemeral channel).
///
/// The **chat-v1 session layer** uses AES-256 (`host-chat-v1/src/crypto.rs`
/// `encrypt`/`decrypt`) because chat sessions are longer-lived and
/// high-value. There is no protocol-level inconsistency: the two channels have
/// separate HKDF derivations and are used for different purposes.
pub fn derive_aes_key(shared_secret: &[u8]) -> [u8; 16] {
    let mut key = [0u8; 16];
    hkdf_sha256(shared_secret, &[], b"polkadot-pairing", &mut key)
        .expect("16 bytes is within HKDF-SHA256 limit");
    key
}

/// Derive a 32-byte channel topic using HMAC-SHA256.
///
/// Key: `account_id` (32-byte sr25519 public key from //wallet//sso).
/// Message: `p256_pubkey_bytes || suffix`.
/// Returns first 32 bytes of the MAC output.
pub fn derive_topic(account_id: &[u8; 32], p256_pubkey: &[u8], suffix: &[u8]) -> [u8; 32] {
    let mut mac: Hmac<Sha256> =
        KeyInit::new_from_slice(account_id).expect("HMAC accepts any key length");
    Mac::update(&mut mac, p256_pubkey);
    Mac::update(&mut mac, suffix);
    let result = mac.finalize().into_bytes();
    let mut out = [0u8; 32];
    out.copy_from_slice(&result);
    out
}

/// Encrypt `plaintext` with AES-128-GCM using a random 12-byte nonce.
///
/// Returns `nonce (12B) || ciphertext`.
pub fn encrypt_aes_gcm(key: &[u8; 16], plaintext: &[u8]) -> Vec<u8> {
    use rand_core::{OsRng, RngCore};

    let cipher = Aes128Gcm::new(key.into());
    let mut nonce_bytes = [0u8; 12];
    OsRng.fill_bytes(&mut nonce_bytes);
    let nonce = Nonce::from_slice(&nonce_bytes);

    let ciphertext = cipher
        .encrypt(nonce, plaintext)
        .expect("AES-128-GCM encryption is infallible for valid key/nonce");

    let mut out = Vec::with_capacity(12 + ciphertext.len());
    out.extend_from_slice(&nonce_bytes);
    out.extend_from_slice(&ciphertext);
    out
}

/// Encrypt `plaintext` with AES-128-GCM using a caller-provided nonce.
///
/// Returns `nonce (12B) || ciphertext`. Useful for deterministic test vectors.
pub fn encrypt_aes_gcm_with_nonce(
    key: &[u8; 16],
    nonce_bytes: &[u8; 12],
    plaintext: &[u8],
) -> Vec<u8> {
    let cipher = Aes128Gcm::new(key.into());
    let nonce = Nonce::from_slice(nonce_bytes);

    let ciphertext = cipher
        .encrypt(nonce, plaintext)
        .expect("AES-128-GCM encryption is infallible for valid key/nonce");

    let mut out = Vec::with_capacity(12 + ciphertext.len());
    out.extend_from_slice(nonce_bytes);
    out.extend_from_slice(&ciphertext);
    out
}

/// Decrypt `nonce (12B) || ciphertext` with AES-128-GCM.
///
/// Returns the plaintext, or an error string on authentication failure.
pub fn decrypt_aes_gcm(key: &[u8; 16], data: &[u8]) -> Result<Vec<u8>, String> {
    if data.len() < 12 {
        return Err(format!(
            "Ciphertext too short: {} bytes (minimum 12)",
            data.len()
        ));
    }
    let (nonce_bytes, ciphertext) = data.split_at(12);
    let cipher = Aes128Gcm::new(key.into());
    let nonce = Nonce::from_slice(nonce_bytes);

    cipher
        .decrypt(nonce, ciphertext)
        .map_err(|_| "AES-128-GCM decryption failed (authentication tag mismatch)".to_string())
}

/// AES-256-GCM encrypt. Returns `nonce (12B) || ciphertext || tag (16B)`.
///
/// Nonce is generated from `OsRng` (cryptographically secure).
///
/// # Nonce reuse bound
/// Random 96-bit nonces have negligible collision probability (< 2^-32)
/// for up to ~2^32 messages per key. The chat protocol must rotate
/// session keys before this limit.
pub fn encrypt_aes256_gcm(key: &[u8; 32], plaintext: &[u8]) -> Vec<u8> {
    use aes_gcm::{aead::KeyInit as _, Aes256Gcm};
    use rand_core::{OsRng, RngCore};

    let cipher = Aes256Gcm::new(key.into());
    let mut nonce_bytes = [0u8; 12];
    OsRng.fill_bytes(&mut nonce_bytes);
    let nonce = aes_gcm::Nonce::from_slice(&nonce_bytes);

    let ciphertext = cipher
        .encrypt(nonce, plaintext)
        .expect("AES-256-GCM encryption cannot fail for valid key/nonce");

    let mut result = Vec::with_capacity(12 + ciphertext.len());
    result.extend_from_slice(&nonce_bytes);
    result.extend_from_slice(&ciphertext);
    result
}

/// AES-256-GCM decrypt. Input: `nonce (12B) || ciphertext || tag (16B)`.
pub fn decrypt_aes256_gcm(key: &[u8; 32], data: &[u8]) -> Result<Vec<u8>, String> {
    use aes_gcm::{aead::KeyInit as _, Aes256Gcm};

    if data.len() < 28 {
        return Err(format!(
            "ciphertext too short: {} bytes (need at least 12B nonce + 16B tag)",
            data.len()
        ));
    }
    let cipher = Aes256Gcm::new(key.into());
    let nonce = aes_gcm::Nonce::from_slice(&data[..12]);
    cipher
        .decrypt(nonce, &data[12..])
        .map_err(|_| "AES-256-GCM decryption failed (authentication error)".into())
}

/// Derive a P-256 private key from 32 bytes of SR25519 key material.
///
/// Accepts the first 32 bytes (scalar half) of a schnorrkel expanded
/// `SecretKey` as returned by `SecretKey::to_bytes()[..32]`.
/// This is the convention used by the chat protocol on both iOS and desktop.
///
/// The seed is hashed with unkeyed Blake2b-256 and the result is
/// interpreted as a P-256 private scalar.
///
/// Returns `Err` if the hash output is zero or >= the P-256 group order
/// (probability ≈ 2^-128 — operationally negligible). Callers must
/// propagate this error; do not retry with different input, as that would
/// break deterministic derivation.
///
/// # Security
/// This derivation links the SR25519 and P-256 key pairs: compromising
/// the SR25519 key reveals the derived P-256 key. This is intentional —
/// both keys belong to the same identity in the chat protocol.
pub fn p256_from_sr25519_scalar(scalar_bytes: &[u8; 32]) -> Result<p256::SecretKey, String> {
    use blake2::digest::consts::U32;
    use blake2::{Blake2b, Digest};

    let hash = Blake2b::<U32>::digest(scalar_bytes);
    p256::SecretKey::from_slice(hash.as_slice())
        .map_err(|e| format!("P-256 scalar derivation failed: {e}"))
}

/// P-256 ECDH key agreement.
///
/// Returns the raw 32-byte x-coordinate of the shared point, wrapped in
/// [`zeroize::Zeroizing`] so the secret material is overwritten when dropped.
///
/// The output is NOT uniformly random. Callers MUST pass this through
/// `hkdf_sha256` (or equivalent KDF) before using it as symmetric key material.
pub fn p256_ecdh(
    own_secret: &p256::SecretKey,
    peer_public: &p256::PublicKey,
) -> zeroize::Zeroizing<[u8; 32]> {
    let shared =
        p256::ecdh::diffie_hellman(own_secret.to_nonzero_scalar(), peer_public.as_affine());
    let mut out = zeroize::Zeroizing::new([0u8; 32]);
    out.copy_from_slice(shared.raw_secret_bytes());
    out
}

/// SCALE-encode the pairing handshake tuple for embedding in the QR URI.
///
/// Format:
/// - `sr25519_pubkey`: 32 raw bytes (fixed-size array, no length prefix)
/// - `p256_pubkey_uncompressed`: 65 raw bytes (fixed-size array, no length prefix)
/// - `metadata_url`: compact-length-prefixed UTF-8 string
pub fn encode_handshake(
    sr25519_pubkey: &[u8; 32],
    p256_pubkey_uncompressed: &[u8; 65],
    metadata_url: &str,
) -> Vec<u8> {
    let url_bytes = metadata_url.as_bytes();
    let url_len = url_bytes.len();

    let mut out = Vec::with_capacity(97 + 4 + url_len);
    out.extend_from_slice(sr25519_pubkey);
    out.extend_from_slice(p256_pubkey_uncompressed);
    out.extend_from_slice(&scale_compact_encode(url_len));
    out.extend_from_slice(url_bytes);
    out
}

/// Encode a `usize` length as a SCALE compact integer.
pub fn scale_compact_encode(n: usize) -> Vec<u8> {
    if n <= 63 {
        vec![(n << 2) as u8]
    } else if n <= 16383 {
        let encoded = ((n << 2) | 0b01) as u16;
        encoded.to_le_bytes().to_vec()
    } else if n <= 1_073_741_823 {
        let encoded = ((n << 2) | 0b10) as u32;
        encoded.to_le_bytes().to_vec()
    } else {
        let bytes = n.to_le_bytes();
        let byte_count = bytes
            .iter()
            .rposition(|&b| b != 0)
            .map(|i| i + 1)
            .unwrap_or(1);
        let mut out = vec![(((byte_count - 4) << 2) | 0b11) as u8];
        out.extend_from_slice(&bytes[..byte_count]);
        out
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // -----------------------------------------------------------------------
    // hkdf_sha256
    // -----------------------------------------------------------------------

    #[test]
    fn hkdf_sha256_produces_deterministic_output() {
        let mut out1 = [0u8; 32];
        let mut out2 = [0u8; 32];
        hkdf_sha256(b"ikm", b"salt", b"info", &mut out1).unwrap();
        hkdf_sha256(b"ikm", b"salt", b"info", &mut out2).unwrap();
        assert_eq!(out1, out2);
        assert_ne!(out1, [0u8; 32]);
    }

    #[test]
    fn hkdf_sha256_empty_salt_differs_from_non_empty_salt() {
        let mut out_empty = [0u8; 32];
        let mut out_salted = [0u8; 32];
        hkdf_sha256(b"ikm", &[], b"info", &mut out_empty).unwrap();
        hkdf_sha256(b"ikm", b"salt", b"info", &mut out_salted).unwrap();
        assert_ne!(out_empty, out_salted);
    }

    #[test]
    fn hkdf_sha256_output_exceeds_limit_returns_error() {
        // HKDF-SHA256 max output = 255 * 32 = 8160 bytes
        let mut out = vec![0u8; 8161];
        let result = hkdf_sha256(b"ikm", &[], b"info", &mut out);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("exceeds max 8160"));
    }

    #[test]
    fn hkdf_sha256_max_output_length_succeeds() {
        let mut out = vec![0u8; 8160];
        assert!(hkdf_sha256(b"ikm", &[], b"info", &mut out).is_ok());
    }

    /// Verify that refactoring `derive_aes_key` to call `hkdf_sha256` did not
    /// change its output — the pinned vector from the original implementation
    /// must still match.
    #[test]
    fn derive_aes_key_still_matches_pinned_vector_after_refactor() {
        // Pinned in tests/vectors.rs as HKDF_KEY.
        let expected: [u8; 16] = [
            0x03, 0xcc, 0x29, 0x7a, 0xbb, 0xd2, 0x59, 0xd9, 0xe5, 0xa7, 0xe3, 0xfd, 0x3d, 0x70,
            0x10, 0xc7,
        ];
        assert_eq!(derive_aes_key(&[0xABu8; 32]), expected);
    }

    // -----------------------------------------------------------------------
    // encrypt_aes256_gcm / decrypt_aes256_gcm
    // -----------------------------------------------------------------------

    #[test]
    fn aes256_gcm_roundtrip() {
        let key = [0x42u8; 32];
        let plaintext = b"hello aes-256-gcm";
        let ciphertext = encrypt_aes256_gcm(&key, plaintext);
        // 12-byte nonce + plaintext + 16-byte tag
        assert_eq!(ciphertext.len(), 12 + plaintext.len() + 16);
        let decrypted = decrypt_aes256_gcm(&key, &ciphertext).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn aes256_gcm_wrong_key_fails() {
        let key1 = [0x42u8; 32];
        let key2 = [0x43u8; 32];
        let ciphertext = encrypt_aes256_gcm(&key1, b"secret");
        assert!(decrypt_aes256_gcm(&key2, &ciphertext).is_err());
    }

    #[test]
    fn aes256_gcm_too_short_returns_error() {
        // Minimum valid input is 12 (nonce) + 16 (tag) = 28 bytes
        let key = [0u8; 32];
        let result = decrypt_aes256_gcm(&key, &[0u8; 27]);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("ciphertext too short"));
    }

    #[test]
    fn aes256_gcm_tampered_ciphertext_fails() {
        let key = [0x11u8; 32];
        let mut ciphertext = encrypt_aes256_gcm(&key, b"do not tamper");
        // Flip a byte in the ciphertext body (after the nonce)
        ciphertext[13] ^= 0xff;
        assert!(decrypt_aes256_gcm(&key, &ciphertext).is_err());
    }

    #[test]
    fn aes256_gcm_nonce_is_random_per_call() {
        let key = [0x55u8; 32];
        let ct1 = encrypt_aes256_gcm(&key, b"same plaintext");
        let ct2 = encrypt_aes256_gcm(&key, b"same plaintext");
        // Different nonces mean different ciphertexts (with overwhelming probability)
        assert_ne!(ct1, ct2, "random nonces must produce different ciphertexts");
    }

    #[test]
    fn aes256_gcm_empty_plaintext_roundtrip() {
        let key = [0u8; 32];
        let ciphertext = encrypt_aes256_gcm(&key, b"");
        // 12-byte nonce + 0-byte plaintext + 16-byte tag
        assert_eq!(ciphertext.len(), 28);
        let decrypted = decrypt_aes256_gcm(&key, &ciphertext).unwrap();
        assert_eq!(decrypted, b"");
    }

    // -----------------------------------------------------------------------
    // p256_from_sr25519_scalar
    // -----------------------------------------------------------------------

    #[test]
    fn p256_from_sr25519_scalar_produces_valid_key() {
        let scalar = [0x42u8; 32];
        let result = p256_from_sr25519_scalar(&scalar);
        assert!(result.is_ok(), "valid scalar must produce a P-256 key");
    }

    #[test]
    fn p256_from_sr25519_scalar_is_deterministic() {
        let scalar = [0x7fu8; 32];
        let k1 = p256_from_sr25519_scalar(&scalar).unwrap();
        let k2 = p256_from_sr25519_scalar(&scalar).unwrap();
        assert_eq!(
            k1.to_bytes().as_slice(),
            k2.to_bytes().as_slice(),
            "same scalar must produce same P-256 key"
        );
    }

    #[test]
    fn p256_from_sr25519_scalar_different_inputs_produce_different_keys() {
        let scalar_a = [0x01u8; 32];
        let scalar_b = [0x02u8; 32];
        let ka = p256_from_sr25519_scalar(&scalar_a).unwrap();
        let kb = p256_from_sr25519_scalar(&scalar_b).unwrap();
        assert_ne!(ka.to_bytes().as_slice(), kb.to_bytes().as_slice());
    }

    // -----------------------------------------------------------------------
    // p256_ecdh
    // -----------------------------------------------------------------------

    #[test]
    fn p256_ecdh_is_commutative() {
        let mut scalar_a = [0u8; 32];
        scalar_a[31] = 0x01;
        let mut scalar_b = [0u8; 32];
        scalar_b[31] = 0x02;

        let sk_a = p256::SecretKey::from_bytes((&scalar_a).into()).unwrap();
        let sk_b = p256::SecretKey::from_bytes((&scalar_b).into()).unwrap();

        let shared_a = p256_ecdh(&sk_a, &sk_b.public_key());
        let shared_b = p256_ecdh(&sk_b, &sk_a.public_key());

        assert_eq!(*shared_a, *shared_b, "P-256 ECDH must be commutative");
    }

    /// Pinned vector: uses the same scalars as the existing `p256_ecdh_pinned_vector`
    /// test in `pairing.rs`. The raw x-coordinate of the shared point is:
    ///   7cf27b188d034f7e8a52380304b51ac3c08969e277f21b35a60b48fc47669978
    #[test]
    fn p256_ecdh_pinned_shared_secret() {
        let mut scalar_a = [0u8; 32];
        scalar_a[31] = 0x01;
        let mut scalar_b = [0u8; 32];
        scalar_b[31] = 0x02;

        let sk_a = p256::SecretKey::from_bytes((&scalar_a).into()).unwrap();
        let sk_b = p256::SecretKey::from_bytes((&scalar_b).into()).unwrap();

        let shared = p256_ecdh(&sk_a, &sk_b.public_key());

        let expected: [u8; 32] = [
            0x7c, 0xf2, 0x7b, 0x18, 0x8d, 0x03, 0x4f, 0x7e, 0x8a, 0x52, 0x38, 0x03, 0x04, 0xb5,
            0x1a, 0xc3, 0xc0, 0x89, 0x69, 0xe2, 0x77, 0xf2, 0x1b, 0x35, 0xa6, 0x0b, 0x48, 0xfc,
            0x47, 0x66, 0x99, 0x78,
        ];
        assert_eq!(*shared, expected, "shared secret must match pinned vector");
    }

    #[test]
    fn p256_ecdh_then_hkdf_matches_derive_aes_key_vector() {
        // Verify that p256_ecdh + hkdf_sha256 produces the same result as the
        // existing derive_aes_key on the same shared secret.
        let mut scalar_a = [0u8; 32];
        scalar_a[31] = 0x01;
        let mut scalar_b = [0u8; 32];
        scalar_b[31] = 0x02;

        let sk_a = p256::SecretKey::from_bytes((&scalar_a).into()).unwrap();
        let sk_b = p256::SecretKey::from_bytes((&scalar_b).into()).unwrap();

        let shared = p256_ecdh(&sk_a, &sk_b.public_key());

        let mut aes_key = [0u8; 16];
        hkdf_sha256(&*shared, &[], b"polkadot-pairing", &mut aes_key).unwrap();

        // Must match the P256_ECDH_AES_KEY pinned vector from tests/vectors.rs.
        let expected: [u8; 16] = [
            0x02, 0x0e, 0x75, 0x45, 0xa8, 0x9f, 0x35, 0xc5, 0x53, 0x47, 0x4f, 0x69, 0xb9, 0x18,
            0xe6, 0xd7,
        ];
        assert_eq!(
            aes_key, expected,
            "HKDF output must match pinned AES key vector"
        );
    }
}
