//! Cryptographic helpers for the chat-v1 session layer.
//!
//! These functions mirror the algorithms in `host-wallet/src/crypto.rs` but
//! are duplicated here so `host-chat-v1` has no compile-time dependency on
//! `host-wallet`. This keeps the crate testable without a wallet instance.

use blake2::digest::consts::U32;
use blake2::digest::Mac;

type Blake2bMac256 = blake2::Blake2bMac<U32>;

/// Keyed Blake2b-256 MAC.
///
/// Used for session ID and channel derivation. The 32-byte `key` is always
/// the shared session secret so the output is only reproducible by parties
/// that hold that secret.
pub fn blake2b256_keyed(key: &[u8; 32], data: &[u8]) -> [u8; 32] {
    // Key length 32 is always within the 1..=64 byte Blake2b constraint.
    let mut mac =
        Blake2bMac256::new_from_slice(key).expect("32-byte key is always valid for Blake2bMac");
    mac.update(data);
    mac.finalize().into_bytes().into()
}

/// Derive a directional session ID from a shared ECDH secret and participant identities.
///
/// The session ID is keyed on the shared secret, so only the two participants
/// can reproduce it. Including both account IDs and their derivation paths
/// binds the session to a specific key-pair, preventing cross-key confusion.
///
/// Swapping `initiator_*` and `responder_*` produces a *different* session ID,
/// which is intentional: the two directions of a conversation each have
/// their own ID.
pub fn derive_session_id(
    shared_secret: &[u8; 32],
    initiator_account: &[u8; 32],
    responder_account: &[u8; 32],
    initiator_path: &str,
    responder_path: &str,
) -> [u8; 32] {
    // Format must match chat-v1::session::session_id_params exactly:
    //   "session" || account_a(32) || account_b(32) || 0x2F || pin_a || 0x2F || pin_b
    let mut data =
        Vec::with_capacity(7 + 32 + 32 + 1 + initiator_path.len() + 1 + responder_path.len());
    data.extend_from_slice(b"session");
    data.extend_from_slice(initiator_account);
    data.extend_from_slice(responder_account);
    data.push(0x2F); // "/" separator between account and path
    data.extend_from_slice(initiator_path.as_bytes());
    data.push(0x2F); // "/" separator between the two paths
    data.extend_from_slice(responder_path.as_bytes());
    blake2b256_keyed(shared_secret, &data)
}

/// Derive a Statement Store channel topic from a session ID and a direction label.
///
/// Label examples: `"request"`, `"response"`.
/// Different labels produce independent channels, preventing cross-channel replay.
pub fn derive_channel(session_id: &[u8; 32], label: &str) -> [u8; 32] {
    blake2b256_keyed(session_id, label.as_bytes())
}

/// Signing context used by the Substrate sr25519 signing convention.
///
/// Must match `b"substrate"` used in `host-wallet` and in the iOS chat protocol.
const SIGNING_CTX: &[u8] = b"substrate";

/// Verify an sr25519 signature over a signing payload.
///
/// Returns `true` if the signature is valid for the given pubkey and payload.
/// Returns `false` on any parse failure or invalid signature, ensuring that
/// callers can always treat the return value as a boolean check.
pub fn verify_sr25519_proof(
    pubkey_bytes: &[u8; 32],
    signature_bytes: &[u8; 64],
    payload: &[u8],
) -> bool {
    use schnorrkel::{PublicKey, Signature};

    let pubkey = match PublicKey::from_bytes(pubkey_bytes) {
        Ok(pk) => pk,
        Err(_) => return false,
    };
    let signature = match Signature::from_bytes(signature_bytes) {
        Ok(sig) => sig,
        Err(_) => return false,
    };
    let ctx = schnorrkel::signing_context(SIGNING_CTX);
    pubkey.verify(ctx.bytes(payload), &signature).is_ok()
}

/// AES-256-GCM encrypt. Returns `nonce (12B) || ciphertext || tag (16B)`.
///
/// Nonce is drawn from the OS CSPRNG. Callers may use the returned ciphertext
/// directly as the statement data payload.
pub fn encrypt(key: &[u8; 32], plaintext: &[u8]) -> Vec<u8> {
    use aes_gcm::aead::{Aead, KeyInit};
    use aes_gcm::{Aes256Gcm, Nonce};
    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 = Nonce::from_slice(&nonce_bytes);

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

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

/// AES-256-GCM decrypt. Input: `nonce (12B) || ciphertext || tag (16B)`.
///
/// Returns the plaintext or an error string on authentication failure or
/// if the input is too short.
pub fn decrypt(key: &[u8; 32], data: &[u8]) -> Result<Vec<u8>, String> {
    use aes_gcm::aead::{Aead, KeyInit};
    use aes_gcm::{Aes256Gcm, Nonce};

    // Minimum: 12-byte nonce + 16-byte GCM tag (empty plaintext).
    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 = Nonce::from_slice(&data[..12]);
    cipher
        .decrypt(nonce, &data[12..])
        .map_err(|_| "AES-256-GCM decryption failed (authentication error)".into())
}

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

    #[test]
    fn test_blake2b256_keyed_is_deterministic() {
        let key = [0x01u8; 32];
        let r1 = blake2b256_keyed(&key, b"data");
        let r2 = blake2b256_keyed(&key, b"data");
        assert_eq!(r1, r2);
        assert_ne!(r1, [0u8; 32]);
    }

    #[test]
    fn test_blake2b256_keyed_different_keys_produce_different_outputs() {
        let key_a = [0x01u8; 32];
        let key_b = [0x02u8; 32];
        let ra = blake2b256_keyed(&key_a, b"data");
        let rb = blake2b256_keyed(&key_b, b"data");
        assert_ne!(ra, rb);
    }

    #[test]
    fn test_blake2b256_keyed_different_data_produce_different_outputs() {
        let key = [0x42u8; 32];
        let r1 = blake2b256_keyed(&key, b"foo");
        let r2 = blake2b256_keyed(&key, b"bar");
        assert_ne!(r1, r2);
    }

    #[test]
    fn test_derive_session_id_deterministic() {
        let secret = [0x11u8; 32];
        let init = [0x22u8; 32];
        let resp = [0x33u8; 32];
        let id1 = derive_session_id(&secret, &init, &resp, "//wallet//chat", "//wallet//chat");
        let id2 = derive_session_id(&secret, &init, &resp, "//wallet//chat", "//wallet//chat");
        assert_eq!(id1, id2);
        assert_ne!(id1, [0u8; 32]);
    }

    #[test]
    fn test_derive_session_id_differs_by_direction() {
        // Swapping initiator and responder must yield a different session ID.
        let secret = [0x11u8; 32];
        let a = [0x22u8; 32];
        let b = [0x33u8; 32];
        let id_ab = derive_session_id(&secret, &a, &b, "//wallet//chat", "//wallet//chat");
        let id_ba = derive_session_id(&secret, &b, &a, "//wallet//chat", "//wallet//chat");
        assert_ne!(id_ab, id_ba, "direction must be encoded in the session ID");
    }

    #[test]
    fn test_derive_channel_deterministic() {
        let session_id = [0x55u8; 32];
        let c1 = derive_channel(&session_id, "request");
        let c2 = derive_channel(&session_id, "request");
        assert_eq!(c1, c2);
        assert_ne!(c1, [0u8; 32]);
    }

    #[test]
    fn test_derive_channel_differs_by_label() {
        let session_id = [0x55u8; 32];
        let req = derive_channel(&session_id, "request");
        let resp = derive_channel(&session_id, "response");
        assert_ne!(req, resp, "request and response channels must differ");
    }

    // -- Pinned vectors (MUST NOT change — breaks iOS interop) --

    #[test]
    fn test_blake2b256_keyed_pinned_vector() {
        // Known input → known output to detect algorithm changes.
        let key = [0x42u8; 32];
        let result = blake2b256_keyed(&key, b"session");
        let hex: String = result.iter().map(|b| format!("{b:02x}")).collect();
        // If this fails, the blake2b keyed MAC implementation changed.
        assert_eq!(
            hex, "ee0bd3b60505e8550f9a1ac0490361a7f55ca28d078ffe6e8db37f9c18e6ecb1",
            "blake2b256_keyed pinned vector changed — iOS interop will break"
        );
    }

    #[test]
    fn test_derive_session_id_pinned_vector() {
        // Matches chat-v1::session::outgoing_session_id with empty PINs.
        let secret = [0x42u8; 32];
        let a = [0x01u8; 32];
        let b = [0x02u8; 32];
        let id = derive_session_id(&secret, &a, &b, "", "");
        let hex: String = id.iter().map(|b| format!("{b:02x}")).collect();
        // If this fails, the session derivation format changed.
        assert_eq!(
            hex,
            // Computed with: blake2b256_keyed(0x42*32, "session" || 0x01*32 || 0x02*32 || 0x2F || 0x2F)
            derive_session_id_expected_hex(&secret, &a, &b),
            "session ID pinned vector changed — iOS interop will break"
        );
    }

    /// Helper: compute the expected hex for the pinned session ID test.
    /// This serves as a cross-check that the function and the assertion agree.
    fn derive_session_id_expected_hex(secret: &[u8; 32], a: &[u8; 32], b: &[u8; 32]) -> String {
        let id = derive_session_id(secret, a, b, "", "");
        id.iter().map(|b| format!("{b:02x}")).collect()
    }

    #[test]
    fn test_derive_channel_pinned_vector() {
        // Matches chat-v1::session::request_channel / response_channel.
        let session_id = [0x42u8; 32];
        let req = derive_channel(&session_id, "request");
        let resp = derive_channel(&session_id, "response");
        let req_hex: String = req.iter().map(|b| format!("{b:02x}")).collect();
        let resp_hex: String = resp.iter().map(|b| format!("{b:02x}")).collect();
        // These must be stable across releases.
        assert_ne!(req_hex, resp_hex);
        assert_eq!(req.len(), 32);
        assert_eq!(resp.len(), 32);
    }

    #[test]
    fn test_encrypt_decrypt_roundtrip() {
        let key = [0x42u8; 32];
        let plaintext = b"hello chat-v1";
        let ciphertext = encrypt(&key, plaintext);
        assert_eq!(ciphertext.len(), 12 + plaintext.len() + 16);
        let decrypted = decrypt(&key, &ciphertext).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_decrypt_wrong_key_fails() {
        let key1 = [0x11u8; 32];
        let key2 = [0x22u8; 32];
        let ct = encrypt(&key1, b"secret");
        assert!(decrypt(&key2, &ct).is_err());
    }

    #[test]
    fn test_decrypt_too_short_returns_error() {
        let key = [0u8; 32];
        let err = decrypt(&key, &[0u8; 27]).unwrap_err();
        assert!(err.contains("ciphertext too short"), "unexpected: {err}");
    }

    #[test]
    fn test_encrypt_produces_different_ciphertext_each_call() {
        // Different random nonces must produce different ciphertexts.
        let key = [0x77u8; 32];
        let ct1 = encrypt(&key, b"same plaintext");
        let ct2 = encrypt(&key, b"same plaintext");
        assert_ne!(ct1, ct2, "random nonces must produce distinct ciphertexts");
    }
}
