//! MLS-Lite: Lightweight symmetric encryption for constrained mesh links.
//!
//! MLS-Lite provides group encryption without the overhead of full MLS:
//! - Pre-shared group secret (exchanged out-of-band: QR code, NFC, voice)
//! - ChaCha20-Poly1305 symmetric encryption (same as MLS application messages)
//! - Per-message nonce derived from epoch + sequence
//! - Replay protection via sequence numbers
//! - Optional Ed25519 signatures for sender authentication
//!
//! # Security Properties
//!
//! - **Confidentiality**: ChaCha20-Poly1305 (256-bit key)
//! - **Integrity**: Poly1305 MAC
//! - **Replay protection**: Sequence numbers
//! - **Sender authentication (optional)**: Ed25519 signatures
//!
//! # NOT Provided (vs full MLS)
//!
//! - Automatic post-compromise security (requires manual key rotation)
//! - Automatic forward secrecy (only per-epoch, not per-message)
//! - Key agreement (keys are pre-shared)
//!
//! # Wire Format
//!
//! See [`MlsLiteEnvelope`] for the compact envelope structure.

use chacha20poly1305::{
    aead::{Aead, KeyInit},
    ChaCha20Poly1305, Nonce,
};
use hkdf::Hkdf;
use rand::RngCore;
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use std::collections::HashMap;

use crate::address::MeshAddress;
use crate::identity::MeshIdentity;

/// Maximum replay window size (track last N sequence numbers per sender).
const REPLAY_WINDOW_SIZE: usize = 64;

/// MLS-Lite group state.
pub struct MlsLiteGroup {
    /// 8-byte group identifier.
    group_id: [u8; 8],
    /// Current epoch (incremented on key rotation).
    epoch: u16,
    /// 32-byte symmetric encryption key (derived from group_secret + epoch).
    encryption_key: [u8; 32],
    /// 7-byte nonce prefix (derived from group_secret).
    nonce_prefix: [u8; 7],
    /// Next sequence number for sending.
    next_seq: u32,
    /// Replay protection: track seen (sender_addr, seq) pairs.
    replay_window: HashMap<MeshAddress, ReplayWindow>,
}

/// Sliding window for replay detection.
struct ReplayWindow {
    /// Highest sequence number seen.
    max_seq: u32,
    /// Bitmap of seen sequence numbers in window.
    seen: u64,
}

impl ReplayWindow {
    fn new() -> Self {
        Self { max_seq: 0, seen: 0 }
    }

    /// Check if sequence number is valid (not replayed).
    /// Returns true if valid, false if replayed or too old.
    fn check_and_update(&mut self, seq: u32) -> bool {
        if seq == 0 {
            // Seq 0 is always allowed once (first message)
            if self.max_seq == 0 && self.seen == 0 {
                self.seen = 1;
                return true;
            }
        }

        if seq > self.max_seq {
            // New highest sequence
            let shift = (seq - self.max_seq).min(64);
            self.seen = self.seen.checked_shl(shift as u32).unwrap_or(0);
            self.seen |= 1; // Mark current as seen
            self.max_seq = seq;
            true
        } else if self.max_seq - seq >= REPLAY_WINDOW_SIZE as u32 {
            // Too old
            false
        } else {
            // Within window — check bitmap
            let idx = (self.max_seq - seq) as u32;
            let bit = 1u64 << idx;
            if self.seen & bit != 0 {
                false // Already seen
            } else {
                self.seen |= bit;
                true
            }
        }
    }
}

/// Result of decryption.
#[derive(Debug)]
pub enum DecryptResult {
    /// Successfully decrypted plaintext.
    Success(Vec<u8>),
    /// Decryption failed (wrong key, corrupted, etc).
    DecryptionFailed,
    /// Replay detected (sequence number already seen).
    ReplayDetected,
    /// Signature verification failed.
    SignatureFailed,
}

impl MlsLiteGroup {
    /// Create a new MLS-Lite group from a pre-shared secret.
    ///
    /// The `group_secret` should be at least 32 bytes of high-entropy data.
    /// It can be:
    /// - Randomly generated and shared via QR code
    /// - Derived from a password via Argon2id
    /// - Exported from a full MLS group's epoch secret
    pub fn new(group_id: [u8; 8], group_secret: &[u8], epoch: u16) -> Self {
        let (encryption_key, nonce_prefix) = Self::derive_keys(group_secret, &group_id, epoch);

        Self {
            group_id,
            epoch,
            encryption_key,
            nonce_prefix,
            next_seq: 0,
            replay_window: HashMap::new(),
        }
    }

    /// Derive encryption key and nonce prefix from group secret and epoch.
    fn derive_keys(group_secret: &[u8], group_id: &[u8; 8], epoch: u16) -> ([u8; 32], [u8; 7]) {
        let salt = b"quicprochat-mls-lite-v1";
        let hk = Hkdf::<Sha256>::new(Some(salt), group_secret);

        // Include epoch in the info to get different keys per epoch
        let mut info = Vec::with_capacity(10);
        info.extend_from_slice(group_id);
        info.extend_from_slice(&epoch.to_be_bytes());

        let mut okm = [0u8; 39]; // 32 bytes key + 7 bytes nonce prefix
        hk.expand(&info, &mut okm)
            .expect("HKDF expand should not fail with valid length");

        let mut key = [0u8; 32];
        let mut prefix = [0u8; 7];
        key.copy_from_slice(&okm[..32]);
        prefix.copy_from_slice(&okm[32..39]);

        (key, prefix)
    }

    /// Rotate to a new epoch with a new group secret.
    pub fn rotate(&mut self, new_secret: &[u8], new_epoch: u16) {
        let (key, prefix) = Self::derive_keys(new_secret, &self.group_id, new_epoch);
        self.encryption_key = key;
        self.nonce_prefix = prefix;
        self.epoch = new_epoch;
        self.next_seq = 0;
        self.replay_window.clear();
    }

    /// Encrypt a plaintext payload.
    ///
    /// Returns `(ciphertext, nonce_suffix, seq)`.
    /// The ciphertext includes the 16-byte Poly1305 tag.
    pub fn encrypt(&mut self, plaintext: &[u8]) -> anyhow::Result<(Vec<u8>, [u8; 5], u32)> {
        let seq = self.next_seq;
        self.next_seq = self.next_seq.wrapping_add(1);

        // Build nonce: 7-byte prefix + 5-byte suffix (1 byte random + 4 byte seq)
        let mut nonce_suffix = [0u8; 5];
        rand::thread_rng().fill_bytes(&mut nonce_suffix[..1]);
        nonce_suffix[1..].copy_from_slice(&seq.to_be_bytes());

        let mut nonce_bytes = [0u8; 12];
        nonce_bytes[..7].copy_from_slice(&self.nonce_prefix);
        nonce_bytes[7..].copy_from_slice(&nonce_suffix);
        let nonce = Nonce::from_slice(&nonce_bytes);

        let cipher = ChaCha20Poly1305::new_from_slice(&self.encryption_key)
            .expect("key length is 32 bytes");

        let ciphertext = cipher
            .encrypt(nonce, plaintext)
            .map_err(|e| anyhow::anyhow!("encryption failed: {e}"))?;

        Ok((ciphertext, nonce_suffix, seq))
    }

    /// Decrypt a ciphertext.
    ///
    /// `sender_addr` is used for replay detection.
    pub fn decrypt(
        &mut self,
        ciphertext: &[u8],
        nonce_suffix: &[u8; 5],
        sender_addr: MeshAddress,
    ) -> DecryptResult {
        // Extract sequence number from nonce suffix
        let seq = u32::from_be_bytes([
            nonce_suffix[1],
            nonce_suffix[2],
            nonce_suffix[3],
            nonce_suffix[4],
        ]);

        // Replay check
        let window = self.replay_window.entry(sender_addr).or_insert_with(ReplayWindow::new);
        if !window.check_and_update(seq) {
            return DecryptResult::ReplayDetected;
        }

        // Build nonce
        let mut nonce_bytes = [0u8; 12];
        nonce_bytes[..7].copy_from_slice(&self.nonce_prefix);
        nonce_bytes[7..].copy_from_slice(nonce_suffix);
        let nonce = Nonce::from_slice(&nonce_bytes);

        let cipher = ChaCha20Poly1305::new_from_slice(&self.encryption_key)
            .expect("key length is 32 bytes");

        match cipher.decrypt(nonce, ciphertext) {
            Ok(plaintext) => DecryptResult::Success(plaintext),
            Err(_) => DecryptResult::DecryptionFailed,
        }
    }

    /// Current epoch.
    pub fn epoch(&self) -> u16 {
        self.epoch
    }

    /// Group ID.
    pub fn group_id(&self) -> &[u8; 8] {
        &self.group_id
    }
}

/// Compact MLS-Lite envelope for constrained links.
///
/// # Wire overhead (approximate)
///
/// - Version: 1 byte
/// - Flags: 1 byte
/// - Group ID: 8 bytes
/// - Sender addr: 4 bytes (truncated further for constrained)
/// - Seq: 4 bytes
/// - Epoch: 2 bytes
/// - Nonce suffix: 5 bytes
/// - Ciphertext: variable (payload + 16 byte tag)
/// - Signature (optional): 64 bytes
///
/// **Minimum overhead without signature: ~41 bytes**
/// **Minimum overhead with signature: ~105 bytes**
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct MlsLiteEnvelope {
    /// Format version (0x03 for MLS-Lite).
    pub version: u8,
    /// Flags: bit 0 = has_signature, bits 1-2 = priority.
    pub flags: u8,
    /// 8-byte group identifier.
    pub group_id: [u8; 8],
    /// 4-byte truncated sender address (first 4 bytes of MeshAddress).
    pub sender_addr: [u8; 4],
    /// Sequence number.
    pub seq: u32,
    /// Key epoch.
    pub epoch: u16,
    /// 5-byte nonce suffix.
    pub nonce: [u8; 5],
    /// Encrypted payload (includes 16-byte Poly1305 tag).
    pub ciphertext: Vec<u8>,
    /// Optional Ed25519 signature (64 bytes, stored as Vec for serde).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub signature: Option<Vec<u8>>,
}

/// MLS-Lite envelope version byte.
const MLS_LITE_VERSION: u8 = 0x03;

impl MlsLiteEnvelope {
    /// Create a new MLS-Lite envelope (without signature).
    pub fn new(
        identity: &MeshIdentity,
        group: &mut MlsLiteGroup,
        plaintext: &[u8],
        sign: bool,
    ) -> anyhow::Result<Self> {
        let (ciphertext, nonce, seq) = group.encrypt(plaintext)?;

        let sender_full = MeshAddress::from_public_key(&identity.public_key());
        let mut sender_addr = [0u8; 4];
        sender_addr.copy_from_slice(&sender_full.as_bytes()[..4]);

        let flags = if sign { 0x01 } else { 0x00 };

        let mut envelope = Self {
            version: MLS_LITE_VERSION,
            flags,
            group_id: *group.group_id(),
            sender_addr,
            seq,
            epoch: group.epoch(),
            nonce,
            ciphertext,
            signature: None,
        };

        if sign {
            let signable = envelope.signable_bytes();
            let sig = identity.sign(&signable);
            envelope.signature = Some(sig.to_vec());
        }

        Ok(envelope)
    }

    /// Bytes to sign (everything except signature).
    fn signable_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(32 + self.ciphertext.len());
        buf.push(self.version);
        buf.push(self.flags);
        buf.extend_from_slice(&self.group_id);
        buf.extend_from_slice(&self.sender_addr);
        buf.extend_from_slice(&self.seq.to_le_bytes());
        buf.extend_from_slice(&self.epoch.to_le_bytes());
        buf.extend_from_slice(&self.nonce);
        buf.extend_from_slice(&self.ciphertext);
        buf
    }

    /// Verify signature (if present) using sender's full public key.
    pub fn verify_signature(&self, sender_public_key: &[u8; 32]) -> bool {
        match &self.signature {
            None => true, // No signature to verify
            Some(sig_vec) => {
                // Signature must be exactly 64 bytes
                let sig: [u8; 64] = match sig_vec.as_slice().try_into() {
                    Ok(s) => s,
                    Err(_) => return false,
                };
                let signable = self.signable_bytes();
                quicprochat_core::IdentityKeypair::verify_raw(sender_public_key, &signable, &sig)
                    .is_ok()
            }
        }
    }

    /// Whether this envelope has a signature.
    pub fn has_signature(&self) -> bool {
        self.flags & 0x01 != 0
    }

    /// Serialize to CBOR.
    pub fn to_wire(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        ciborium::into_writer(self, &mut buf).expect("CBOR serialization should not fail");
        buf
    }

    /// Deserialize from CBOR.
    pub fn from_wire(bytes: &[u8]) -> anyhow::Result<Self> {
        let env: Self = ciborium::from_reader(bytes)?;
        if env.version != MLS_LITE_VERSION {
            anyhow::bail!("unexpected MLS-Lite version: {}", env.version);
        }
        Ok(env)
    }
}

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

    fn test_identity() -> MeshIdentity {
        MeshIdentity::generate()
    }

    #[test]
    fn encrypt_decrypt_roundtrip() {
        let secret = b"super secret group key material!";
        let group_id = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08];

        let mut alice_group = MlsLiteGroup::new(group_id, secret, 0);
        let mut bob_group = MlsLiteGroup::new(group_id, secret, 0);

        let plaintext = b"hello from alice";
        let (ciphertext, nonce, _seq) = alice_group.encrypt(plaintext).expect("encrypt");

        let alice_addr = MeshAddress::from_bytes([0xAA; 16]);
        match bob_group.decrypt(&ciphertext, &nonce, alice_addr) {
            DecryptResult::Success(pt) => assert_eq!(pt, plaintext),
            other => panic!("expected Success, got {other:?}"),
        }
    }

    #[test]
    fn replay_detection() {
        let secret = b"replay test key material here!!!";
        let group_id = [0x11; 8];

        let mut alice_group = MlsLiteGroup::new(group_id, secret, 0);
        let mut bob_group = MlsLiteGroup::new(group_id, secret, 0);

        let (ciphertext, nonce, _seq) = alice_group.encrypt(b"msg1").expect("encrypt");
        let alice_addr = MeshAddress::from_bytes([0xAA; 16]);

        // First decrypt succeeds
        match bob_group.decrypt(&ciphertext, &nonce, alice_addr) {
            DecryptResult::Success(_) => {}
            other => panic!("first decrypt should succeed, got {other:?}"),
        }

        // Replay attempt fails
        match bob_group.decrypt(&ciphertext, &nonce, alice_addr) {
            DecryptResult::ReplayDetected => {}
            other => panic!("replay should be detected, got {other:?}"),
        }
    }

    #[test]
    fn different_epochs_different_keys() {
        let secret = b"epoch rotation test material!!!";
        let group_id = [0x22; 8];

        let mut group_e0 = MlsLiteGroup::new(group_id, secret, 0);
        let mut group_e1 = MlsLiteGroup::new(group_id, secret, 1);

        let (ciphertext_e0, nonce_e0, _) = group_e0.encrypt(b"epoch 0").expect("encrypt");

        // Decrypt with wrong epoch should fail
        let sender = MeshAddress::from_bytes([0xBB; 16]);
        match group_e1.decrypt(&ciphertext_e0, &nonce_e0, sender) {
            DecryptResult::DecryptionFailed => {}
            other => panic!("wrong epoch should fail decryption, got {other:?}"),
        }
    }

    #[test]
    fn envelope_with_signature() {
        let id = test_identity();
        let secret = b"envelope signature test material";
        let group_id = [0x33; 8];

        let mut group = MlsLiteGroup::new(group_id, secret, 0);

        let envelope = MlsLiteEnvelope::new(&id, &mut group, b"signed message", true)
            .expect("create envelope");

        assert!(envelope.has_signature());
        assert!(envelope.verify_signature(&id.public_key()));

        // Wrong key should fail
        let wrong_key = [0x42u8; 32];
        assert!(!envelope.verify_signature(&wrong_key));
    }

    #[test]
    fn envelope_without_signature() {
        let id = test_identity();
        let secret = b"unsigned envelope test material!";
        let group_id = [0x44; 8];

        let mut group = MlsLiteGroup::new(group_id, secret, 0);

        let envelope = MlsLiteEnvelope::new(&id, &mut group, b"no sig", false)
            .expect("create envelope");

        assert!(!envelope.has_signature());
        assert!(envelope.signature.is_none());
    }

    #[test]
    fn envelope_cbor_roundtrip() {
        let id = test_identity();
        let secret = b"cbor roundtrip test material!!!!";
        let group_id = [0x55; 8];

        let mut group = MlsLiteGroup::new(group_id, secret, 0);

        let envelope = MlsLiteEnvelope::new(&id, &mut group, b"roundtrip", true)
            .expect("create envelope");

        let wire = envelope.to_wire();
        let restored = MlsLiteEnvelope::from_wire(&wire).expect("deserialize");

        assert_eq!(envelope.version, restored.version);
        assert_eq!(envelope.flags, restored.flags);
        assert_eq!(envelope.group_id, restored.group_id);
        assert_eq!(envelope.sender_addr, restored.sender_addr);
        assert_eq!(envelope.seq, restored.seq);
        assert_eq!(envelope.epoch, restored.epoch);
        assert_eq!(envelope.nonce, restored.nonce);
        assert_eq!(envelope.ciphertext, restored.ciphertext);
        assert_eq!(envelope.signature, restored.signature);
    }

    #[test]
    fn measure_mls_lite_overhead() {
        let id = test_identity();
        let secret = b"overhead measurement test secret";
        let group_id = [0x66; 8];

        let mut group = MlsLiteGroup::new(group_id, secret, 0);

        println!("=== MLS-Lite Wire Overhead (CBOR) ===");

        // Without signature
        let env_no_sig = MlsLiteEnvelope::new(&id, &mut group, b"", false)
            .expect("create");
        let wire_no_sig = env_no_sig.to_wire();
        // Overhead = wire - payload - 16 byte tag
        let overhead_no_sig = wire_no_sig.len() - 16; // tag is in ciphertext
        println!("No signature, 0B payload: {} bytes (overhead: {})", wire_no_sig.len(), overhead_no_sig);

        // With signature
        let env_sig = MlsLiteEnvelope::new(&id, &mut group, b"", true)
            .expect("create");
        let wire_sig = env_sig.to_wire();
        let overhead_sig = wire_sig.len() - 16;
        println!("With signature, 0B payload: {} bytes (overhead: {})", wire_sig.len(), overhead_sig);

        // 10-byte payload without sig
        let env_10 = MlsLiteEnvelope::new(&id, &mut group, b"hello mesh", false)
            .expect("create");
        let wire_10 = env_10.to_wire();
        println!("No signature, 10B payload: {} bytes", wire_10.len());

        // Compare to MeshEnvelope V1
        let v1_env = crate::envelope::MeshEnvelope::new(
            &id,
            &[0x77; 32],
            b"hello mesh".to_vec(),
            3600,
            5,
        );
        let v1_wire = v1_env.to_wire();
        println!("MeshEnvelope V1, 10B payload: {} bytes", v1_wire.len());
        println!("MLS-Lite savings (no sig): {} bytes", v1_wire.len() as i32 - wire_10.len() as i32);

        // MLS-Lite overhead is higher than raw struct due to CBOR encoding
        // but still much less than full MLS or MeshEnvelope
        assert!(overhead_no_sig < 150, "MLS-Lite overhead without sig should be under 150 bytes");
        assert!(overhead_sig < 300, "MLS-Lite overhead with sig should be under 300 bytes");
        // Key assertion: MLS-Lite should be significantly smaller than V1
        assert!(
            wire_10.len() < v1_wire.len() / 2,
            "MLS-Lite should be at least 2x smaller than MeshEnvelope V1"
        );
    }
}