//! Lightweight encrypted mesh link for constrained transports.
//!
//! On high-bandwidth transports (QUIC/TCP), we use TLS 1.3. On constrained
//! transports (LoRa, Serial), the full TLS handshake is too expensive
//! (~2-4 KB). This module provides a minimal 3-packet handshake that
//! establishes a ChaCha20-Poly1305 encrypted session in ~240 bytes total.
//!
//! # Handshake Protocol
//!
//! ```text
//! Packet 1: Initiator -> Responder  (80 bytes)
//!   [initiator_addr: 16][eph_x25519_pub: 32][nonce: 24][flags: 8]
//!
//! Packet 2: Responder -> Initiator  (96 bytes)
//!   [responder_addr: 16][eph_x25519_pub: 32][encrypted_proof: 32][tag: 16]
//!
//! Packet 3: Initiator -> Responder  (48 bytes)
//!   [encrypted_proof: 32][tag: 16]
//!
//! Total: 224 bytes
//!
//! Shared secret: HKDF-SHA256(ikm = X25519(eph_a, eph_b), info = "qpc-mesh-link-v1")
//! ```

use chacha20poly1305::aead::{Aead, KeyInit};
use chacha20poly1305::{ChaCha20Poly1305, Nonce};
use hkdf::Hkdf;
use rand::rngs::OsRng;
use rand::RngCore;
use sha2::Sha256;
use x25519_dalek::{EphemeralSecret, PublicKey as X25519Public};
use zeroize::Zeroize;

use crate::address::MeshAddress;

/// Errors that can occur during link handshake or encryption.
#[derive(Debug, thiserror::Error)]
pub enum LinkError {
    /// Received packet has wrong length.
    #[error("invalid packet length: expected {expected}, got {got}")]
    InvalidLength { expected: usize, got: usize },

    /// AEAD decryption failed (wrong key or tampered data).
    #[error("decryption failed: invalid ciphertext or authentication tag")]
    DecryptionFailed,

    /// The proof inside a handshake packet did not match the expected address.
    #[error("handshake proof mismatch: peer address does not match encrypted proof")]
    ProofMismatch,
}

/// Packet sizes for the 3-packet handshake.
pub const PACKET1_LEN: usize = 80; // 16 + 32 + 24 + 8
pub const PACKET2_LEN: usize = 96; // 16 + 32 + 16 + 16 + 16 (addr + pub + encrypted_addr + tag)
pub const PACKET3_LEN: usize = 48; // 16 + 16 + 16 (encrypted_addr + tag)

/// Derive a 32-byte session key from a shared secret and nonce via HKDF-SHA256.
fn derive_session_key(shared_secret: &[u8], salt: &[u8]) -> [u8; 32] {
    let hk = Hkdf::<Sha256>::new(Some(salt), shared_secret);
    let mut key = [0u8; 32];
    hk.expand(b"qpc-mesh-link-v1", &mut key)
        .expect("HKDF expand to 32 bytes should never fail");
    key
}

/// Build a ChaCha20Poly1305 nonce from a u64 counter (zero-padded, little-endian).
fn counter_nonce(counter: u64) -> Nonce {
    let mut nonce_bytes = [0u8; 12];
    nonce_bytes[..8].copy_from_slice(&counter.to_le_bytes());
    *Nonce::from_slice(&nonce_bytes)
}

/// An established encrypted mesh link session.
pub struct MeshLink {
    /// Derived symmetric key for ChaCha20-Poly1305.
    session_key: [u8; 32],
    /// Remote peer's mesh address.
    remote_address: MeshAddress,
    /// Message counter for nonce derivation (send direction).
    send_counter: u64,
    /// Message counter for nonce derivation (receive direction).
    recv_counter: u64,
}

impl Drop for MeshLink {
    fn drop(&mut self) {
        self.session_key.zeroize();
    }
}

impl MeshLink {
    /// Encrypt a message using the session key.
    ///
    /// Returns the ciphertext (plaintext + 16-byte Poly1305 tag).
    pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, LinkError> {
        // Nonces for encrypt start at offset 256 to avoid collision with handshake nonces.
        let nonce = counter_nonce(256 + self.send_counter);
        let cipher = ChaCha20Poly1305::new((&self.session_key).into());
        let ciphertext = cipher
            .encrypt(&nonce, plaintext)
            .map_err(|_| LinkError::DecryptionFailed)?;
        self.send_counter += 1;
        Ok(ciphertext)
    }

    /// Decrypt a message using the session key.
    pub fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>, LinkError> {
        let nonce = counter_nonce(256 + self.recv_counter);
        let cipher = ChaCha20Poly1305::new((&self.session_key).into());
        let plaintext = cipher
            .decrypt(&nonce, ciphertext)
            .map_err(|_| LinkError::DecryptionFailed)?;
        self.recv_counter += 1;
        Ok(plaintext)
    }

    /// Remote peer's address.
    pub fn remote_address(&self) -> MeshAddress {
        self.remote_address
    }

    /// Number of messages sent on this link.
    pub fn messages_sent(&self) -> u64 {
        self.send_counter
    }

    /// Number of messages received on this link.
    pub fn messages_received(&self) -> u64 {
        self.recv_counter
    }

    /// Access the session key (for testing only).
    #[cfg(test)]
    fn session_key(&self) -> &[u8; 32] {
        &self.session_key
    }
}

/// Handshake state for the initiator side of a mesh link.
pub struct LinkInitiator {
    local_address: MeshAddress,
    eph_secret: EphemeralSecret,
    nonce: [u8; 24],
}

/// Handshake state for the responder side of a mesh link.
pub struct LinkResponder {
    remote_address: MeshAddress,
    session_key: [u8; 32],
}

impl Drop for LinkResponder {
    fn drop(&mut self) {
        self.session_key.zeroize();
    }
}

impl LinkInitiator {
    /// Create initiator state and generate Packet 1.
    ///
    /// Packet 1 layout (80 bytes):
    /// `[initiator_addr: 16][eph_pub: 32][nonce: 24][flags: 8]`
    pub fn new(local_address: MeshAddress) -> (Self, Vec<u8>) {
        let eph_secret = EphemeralSecret::random_from_rng(OsRng);
        let eph_public = X25519Public::from(&eph_secret);

        let mut nonce = [0u8; 24];
        OsRng.fill_bytes(&mut nonce);

        let mut packet = Vec::with_capacity(PACKET1_LEN);
        packet.extend_from_slice(local_address.as_bytes());
        packet.extend_from_slice(eph_public.as_bytes());
        packet.extend_from_slice(&nonce);
        packet.extend_from_slice(&[0u8; 8]); // flags: reserved

        let initiator = Self {
            local_address,
            eph_secret,
            nonce,
        };

        (initiator, packet)
    }

    /// Process Packet 2 from responder, generate Packet 3, return completed link.
    ///
    /// Packet 2 layout (96 bytes):
    /// `[responder_addr: 16][eph_pub: 32][encrypted_responder_addr: 16+16]`
    ///
    /// Packet 3 layout (48 bytes):
    /// `[encrypted_initiator_addr: 16+16][padding: 16]`
    pub fn process_response(self, packet2: &[u8]) -> Result<(MeshLink, Vec<u8>), LinkError> {
        if packet2.len() != PACKET2_LEN {
            return Err(LinkError::InvalidLength {
                expected: PACKET2_LEN,
                got: packet2.len(),
            });
        }

        // Parse Packet 2.
        let mut responder_addr_bytes = [0u8; 16];
        responder_addr_bytes.copy_from_slice(&packet2[..16]);
        let responder_address = MeshAddress::from_bytes(responder_addr_bytes);

        let mut responder_eph_pub_bytes = [0u8; 32];
        responder_eph_pub_bytes.copy_from_slice(&packet2[16..48]);
        let responder_eph_pub = X25519Public::from(responder_eph_pub_bytes);

        let encrypted_proof = &packet2[48..80]; // 16-byte ciphertext + 16-byte Poly1305 tag = 32 bytes

        // Compute shared secret (consumes eph_secret).
        let shared_secret = self.eph_secret.diffie_hellman(&responder_eph_pub);

        // Derive session key.
        let session_key = derive_session_key(shared_secret.as_bytes(), &self.nonce);

        // Verify responder's proof: decrypt and check it matches responder_addr.
        let cipher = ChaCha20Poly1305::new((&session_key).into());
        let proof_nonce = counter_nonce(0);
        let decrypted_proof = cipher
            .decrypt(&proof_nonce, encrypted_proof)
            .map_err(|_| LinkError::DecryptionFailed)?;

        if decrypted_proof.as_slice() != responder_addr_bytes.as_slice() {
            return Err(LinkError::ProofMismatch);
        }

        // Build Packet 3: encrypt our address as proof.
        let proof_nonce_3 = counter_nonce(1);
        let encrypted_initiator_addr = cipher
            .encrypt(&proof_nonce_3, self.local_address.as_bytes().as_slice())
            .map_err(|_| LinkError::DecryptionFailed)?;

        let mut packet3 = Vec::with_capacity(PACKET3_LEN);
        packet3.extend_from_slice(&encrypted_initiator_addr);
        // Pad to 48 bytes.
        packet3.resize(PACKET3_LEN, 0);

        let link = MeshLink {
            session_key,
            remote_address: responder_address,
            send_counter: 0,
            recv_counter: 0,
        };

        Ok((link, packet3))
    }
}

impl LinkResponder {
    /// Process Packet 1 from initiator, generate Packet 2.
    ///
    /// Packet 1 layout (80 bytes):
    /// `[initiator_addr: 16][eph_pub: 32][nonce: 24][flags: 8]`
    ///
    /// Packet 2 layout (96 bytes):
    /// `[responder_addr: 16][eph_pub: 32][encrypted_responder_addr: 16+16]`
    pub fn new(
        local_address: MeshAddress,
        packet1: &[u8],
    ) -> Result<(Self, Vec<u8>), LinkError> {
        if packet1.len() != PACKET1_LEN {
            return Err(LinkError::InvalidLength {
                expected: PACKET1_LEN,
                got: packet1.len(),
            });
        }

        // Parse Packet 1.
        let mut initiator_addr_bytes = [0u8; 16];
        initiator_addr_bytes.copy_from_slice(&packet1[..16]);
        let remote_address = MeshAddress::from_bytes(initiator_addr_bytes);

        let mut initiator_eph_pub_bytes = [0u8; 32];
        initiator_eph_pub_bytes.copy_from_slice(&packet1[16..48]);
        let initiator_eph_pub = X25519Public::from(initiator_eph_pub_bytes);

        let mut nonce = [0u8; 24];
        nonce.copy_from_slice(&packet1[48..72]);
        // flags at [72..80] — reserved, ignored.

        // Generate our ephemeral keypair.
        let eph_secret = EphemeralSecret::random_from_rng(OsRng);
        let eph_public = X25519Public::from(&eph_secret);

        // Compute shared secret (consumes eph_secret).
        let shared_secret = eph_secret.diffie_hellman(&initiator_eph_pub);

        // Derive session key.
        let session_key = derive_session_key(shared_secret.as_bytes(), &nonce);

        // Build Packet 2: our address + our eph_pub + encrypted proof of our address.
        let cipher = ChaCha20Poly1305::new((&session_key).into());
        let proof_nonce = counter_nonce(0);
        let encrypted_proof = cipher
            .encrypt(&proof_nonce, local_address.as_bytes().as_slice())
            .map_err(|_| LinkError::DecryptionFailed)?;

        let mut packet2 = Vec::with_capacity(PACKET2_LEN);
        packet2.extend_from_slice(local_address.as_bytes());
        packet2.extend_from_slice(eph_public.as_bytes());
        packet2.extend_from_slice(&encrypted_proof);
        // Pad to PACKET2_LEN for fixed-size framing on constrained transports.
        packet2.resize(PACKET2_LEN, 0);

        let responder = Self {
            remote_address,
            session_key,
        };

        Ok((responder, packet2))
    }

    /// Process Packet 3 from initiator, return completed link.
    ///
    /// Packet 3 layout (48 bytes):
    /// `[encrypted_initiator_addr: 16+16][padding: 16]`
    pub fn complete(self, packet3: &[u8]) -> Result<MeshLink, LinkError> {
        if packet3.len() != PACKET3_LEN {
            return Err(LinkError::InvalidLength {
                expected: PACKET3_LEN,
                got: packet3.len(),
            });
        }

        // The encrypted proof is the first 32 bytes (16 plaintext + 16 tag).
        let encrypted_proof = &packet3[..32];

        let cipher = ChaCha20Poly1305::new((&self.session_key).into());
        let proof_nonce = counter_nonce(1);
        let decrypted_proof = cipher
            .decrypt(&proof_nonce, encrypted_proof)
            .map_err(|_| LinkError::DecryptionFailed)?;

        let mut expected_addr = [0u8; 16];
        expected_addr.copy_from_slice(self.remote_address.as_bytes());

        if decrypted_proof.as_slice() != expected_addr.as_slice() {
            return Err(LinkError::ProofMismatch);
        }

        Ok(MeshLink {
            session_key: self.session_key,
            remote_address: self.remote_address,
            send_counter: 0,
            recv_counter: 0,
        })
    }
}

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

    fn test_address(byte: u8) -> MeshAddress {
        MeshAddress::from_public_key(&[byte; 32])
    }

    #[test]
    fn full_handshake_roundtrip() {
        let addr_a = test_address(1);
        let addr_b = test_address(2);

        // Initiator creates Packet 1.
        let (initiator, packet1) = LinkInitiator::new(addr_a);
        assert_eq!(packet1.len(), PACKET1_LEN);

        // Responder processes Packet 1, creates Packet 2.
        let (responder, packet2) = LinkResponder::new(addr_b, &packet1).expect("responder::new");
        assert_eq!(packet2.len(), PACKET2_LEN);

        // Initiator processes Packet 2, creates Packet 3, gets link.
        let (link_a, packet3) = initiator
            .process_response(&packet2)
            .expect("initiator::process_response");
        assert_eq!(packet3.len(), PACKET3_LEN);

        // Responder processes Packet 3, gets link.
        let link_b = responder.complete(&packet3).expect("responder::complete");

        // Both sides should have the same session key.
        assert_eq!(link_a.session_key(), link_b.session_key());

        // Check remote addresses.
        assert_eq!(link_a.remote_address(), addr_b);
        assert_eq!(link_b.remote_address(), addr_a);
    }

    #[test]
    fn encrypt_decrypt_roundtrip() {
        let addr_a = test_address(10);
        let addr_b = test_address(20);

        let (initiator, packet1) = LinkInitiator::new(addr_a);
        let (responder, packet2) = LinkResponder::new(addr_b, &packet1).expect("responder");
        let (mut link_a, packet3) = initiator.process_response(&packet2).expect("initiator");
        let mut link_b = responder.complete(&packet3).expect("complete");

        let plaintext = b"hello constrained mesh";
        let ciphertext = link_a.encrypt(plaintext).expect("encrypt");
        let decrypted = link_b.decrypt(&ciphertext).expect("decrypt");
        assert_eq!(decrypted, plaintext);

        // Reverse direction.
        let plaintext2 = b"hello back";
        let ciphertext2 = link_b.encrypt(plaintext2).expect("encrypt");
        let decrypted2 = link_a.decrypt(&ciphertext2).expect("decrypt");
        assert_eq!(decrypted2, plaintext2);
    }

    #[test]
    fn wrong_key_fails_decrypt() {
        let addr_a = test_address(30);
        let addr_b = test_address(40);

        let (initiator, packet1) = LinkInitiator::new(addr_a);
        let (responder, packet2) = LinkResponder::new(addr_b, &packet1).expect("responder");
        let (mut link_a, packet3) = initiator.process_response(&packet2).expect("initiator");
        let _link_b = responder.complete(&packet3).expect("complete");

        let ciphertext = link_a.encrypt(b"secret").expect("encrypt");

        // Create a link with a different session key.
        let mut fake_link = MeshLink {
            session_key: [0xFFu8; 32],
            remote_address: addr_a,
            send_counter: 0,
            recv_counter: 0,
        };

        let result = fake_link.decrypt(&ciphertext);
        assert!(result.is_err(), "decryption with wrong key must fail");
    }

    #[test]
    fn counter_increments() {
        let addr_a = test_address(50);
        let addr_b = test_address(60);

        let (initiator, packet1) = LinkInitiator::new(addr_a);
        let (responder, packet2) = LinkResponder::new(addr_b, &packet1).expect("responder");
        let (mut link_a, packet3) = initiator.process_response(&packet2).expect("initiator");
        let mut link_b = responder.complete(&packet3).expect("complete");

        assert_eq!(link_a.messages_sent(), 0);
        assert_eq!(link_b.messages_received(), 0);

        link_a.encrypt(b"msg1").expect("encrypt");
        assert_eq!(link_a.messages_sent(), 1);

        link_a.encrypt(b"msg2").expect("encrypt");
        assert_eq!(link_a.messages_sent(), 2);

        // Decrypt two messages on the other side.
        // We need fresh ciphertexts — re-do with proper counter tracking.
        let addr_c = test_address(70);
        let addr_d = test_address(80);
        let (init2, p1) = LinkInitiator::new(addr_c);
        let (resp2, p2) = LinkResponder::new(addr_d, &p1).expect("responder");
        let (mut la, p3) = init2.process_response(&p2).expect("initiator");
        let mut lb = resp2.complete(&p3).expect("complete");

        let ct1 = la.encrypt(b"msg1").expect("encrypt");
        let ct2 = la.encrypt(b"msg2").expect("encrypt");

        lb.decrypt(&ct1).expect("decrypt");
        assert_eq!(lb.messages_received(), 1);

        lb.decrypt(&ct2).expect("decrypt");
        assert_eq!(lb.messages_received(), 2);
    }

    #[test]
    fn packet_sizes() {
        let addr = test_address(90);

        let (_initiator, packet1) = LinkInitiator::new(addr);
        assert_eq!(packet1.len(), 80, "packet 1 must be 80 bytes");

        // Complete a handshake to check packet 2 and 3 sizes.
        let addr_b = test_address(91);
        let (init, p1) = LinkInitiator::new(addr);
        let (resp, p2) = LinkResponder::new(addr_b, &p1).expect("responder");
        assert_eq!(p2.len(), 96, "packet 2 must be 96 bytes");

        let (_link, p3) = init.process_response(&p2).expect("initiator");
        assert_eq!(p3.len(), 48, "packet 3 must be 48 bytes");

        // Verify responder can complete.
        resp.complete(&p3).expect("complete");
    }
}