//! Lightweight pub/sub broadcast channels for mesh announcements.
//!
//! Each [`BroadcastChannel`] holds a ChaCha20-Poly1305 symmetric key used to
//! encrypt and decrypt messages on that topic. Peers that know the key can
//! subscribe; the key itself is exchanged out-of-band.
//!
//! [`BroadcastManager`] collects channels by topic and provides convenience
//! methods for encrypt/decrypt without exposing raw keys.

use std::collections::HashMap;

use chacha20poly1305::aead::{Aead, AeadCore, KeyInit};
use chacha20poly1305::ChaCha20Poly1305;
use rand::rngs::OsRng;
use zeroize::{Zeroize, ZeroizeOnDrop};

/// A single broadcast channel identified by topic, secured with a symmetric key.
pub struct BroadcastChannel {
    topic: String,
    key: [u8; 32],
}

impl Drop for BroadcastChannel {
    fn drop(&mut self) {
        self.key.zeroize();
    }
}

impl ZeroizeOnDrop for BroadcastChannel {}

impl BroadcastChannel {
    /// Create a new channel with a random ChaCha20-Poly1305 key.
    pub fn new(topic: &str) -> Self {
        let mut key = [0u8; 32];
        rand::RngCore::fill_bytes(&mut OsRng, &mut key);
        Self {
            topic: topic.to_string(),
            key,
        }
    }

    /// Create a channel with a pre-shared key (e.g. received from another peer).
    pub fn with_key(topic: &str, key: [u8; 32]) -> Self {
        Self {
            topic: topic.to_string(),
            key,
        }
    }

    /// Encrypt `plaintext`, returning `nonce || ciphertext`.
    pub fn encrypt(&self, plaintext: &[u8]) -> anyhow::Result<Vec<u8>> {
        let cipher = ChaCha20Poly1305::new((&self.key).into());
        let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
        let ciphertext = cipher
            .encrypt(&nonce, plaintext)
            .map_err(|_| anyhow::anyhow!("ChaCha20Poly1305 encryption failed"))?;
        let mut out = Vec::with_capacity(nonce.len() + ciphertext.len());
        out.extend_from_slice(&nonce);
        out.extend_from_slice(&ciphertext);
        Ok(out)
    }

    /// Decrypt data produced by [`encrypt`](Self::encrypt).
    ///
    /// Expects `nonce (12 bytes) || ciphertext`.
    pub fn decrypt(&self, data: &[u8]) -> anyhow::Result<Vec<u8>> {
        if data.len() < 12 {
            anyhow::bail!("broadcast ciphertext too short (need at least 12-byte nonce)");
        }
        let (nonce_bytes, ciphertext) = data.split_at(12);
        let nonce = chacha20poly1305::Nonce::from_slice(nonce_bytes);
        let cipher = ChaCha20Poly1305::new((&self.key).into());
        cipher
            .decrypt(nonce, ciphertext)
            .map_err(|_| anyhow::anyhow!("broadcast decryption failed (wrong key or corrupted)"))
    }

    /// The topic name for this channel.
    pub fn topic(&self) -> &str {
        &self.topic
    }

    /// The raw 32-byte symmetric key (for sharing with peers out-of-band).
    pub fn key(&self) -> &[u8; 32] {
        &self.key
    }
}

/// Manages a set of broadcast channels keyed by topic.
pub struct BroadcastManager {
    channels: HashMap<String, BroadcastChannel>,
}

impl BroadcastManager {
    /// Create an empty manager.
    pub fn new() -> Self {
        Self {
            channels: HashMap::new(),
        }
    }

    /// Subscribe to a topic with a pre-shared key.
    pub fn subscribe(&mut self, topic: &str, key: [u8; 32]) {
        self.channels
            .insert(topic.to_string(), BroadcastChannel::with_key(topic, key));
    }

    /// Unsubscribe from a topic.
    pub fn unsubscribe(&mut self, topic: &str) {
        self.channels.remove(topic);
    }

    /// Create a new broadcast channel with a random key and return a reference.
    pub fn create_channel(&mut self, topic: &str) -> &BroadcastChannel {
        self.channels
            .insert(topic.to_string(), BroadcastChannel::new(topic));
        self.channels
            .get(topic)
            .expect("just inserted")
    }

    /// Look up a channel by topic.
    pub fn get(&self, topic: &str) -> Option<&BroadcastChannel> {
        self.channels.get(topic)
    }

    /// List all subscribed topics.
    pub fn topics(&self) -> Vec<String> {
        self.channels.keys().cloned().collect()
    }

    /// Encrypt a message on the given topic. Returns `None` if not subscribed.
    pub fn encrypt(&self, topic: &str, plaintext: &[u8]) -> Option<anyhow::Result<Vec<u8>>> {
        self.channels.get(topic).map(|ch| ch.encrypt(plaintext))
    }

    /// Decrypt a message on the given topic. Returns `None` if not subscribed.
    pub fn decrypt(&self, topic: &str, data: &[u8]) -> Option<Vec<u8>> {
        self.channels
            .get(topic)
            .and_then(|ch| ch.decrypt(data).ok())
    }
}

impl Default for BroadcastManager {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn encrypt_decrypt_roundtrip() {
        let ch = BroadcastChannel::new("test-topic");
        let plaintext = b"hello broadcast";
        let encrypted = ch.encrypt(plaintext).expect("encrypt");
        let decrypted = ch.decrypt(&encrypted).expect("decrypt");
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn wrong_key_fails_decrypt() {
        let ch1 = BroadcastChannel::new("topic");
        let ch2 = BroadcastChannel::new("topic"); // different random key
        let encrypted = ch1.encrypt(b"secret").expect("encrypt");
        let result = ch2.decrypt(&encrypted);
        assert!(result.is_err(), "wrong key should fail decryption");
    }

    #[test]
    fn with_key_roundtrip() {
        let key = [42u8; 32];
        let ch = BroadcastChannel::with_key("shared", key);
        let ct = ch.encrypt(b"data").expect("encrypt");
        let ch2 = BroadcastChannel::with_key("shared", key);
        let pt = ch2.decrypt(&ct).expect("same key should decrypt");
        assert_eq!(pt, b"data");
    }

    #[test]
    fn manager_subscribe_unsubscribe() {
        let mut mgr = BroadcastManager::new();
        assert!(mgr.topics().is_empty());

        let key = [1u8; 32];
        mgr.subscribe("alerts", key);
        assert_eq!(mgr.topics().len(), 1);
        assert!(mgr.get("alerts").is_some());

        mgr.unsubscribe("alerts");
        assert!(mgr.topics().is_empty());
        assert!(mgr.get("alerts").is_none());
    }

    #[test]
    fn manager_create_channel() {
        let mut mgr = BroadcastManager::new();
        let ch = mgr.create_channel("news");
        let key = *ch.key();
        assert_eq!(ch.topic(), "news");

        // Encrypt via manager, decrypt manually with the same key.
        let ct = mgr.encrypt("news", b"headline").expect("subscribed").expect("encrypt");
        let ch2 = BroadcastChannel::with_key("news", key);
        let pt = ch2.decrypt(&ct).expect("decrypt");
        assert_eq!(pt, b"headline");
    }

    #[test]
    fn manager_encrypt_decrypt() {
        let mut mgr = BroadcastManager::new();
        mgr.subscribe("ch1", [7u8; 32]);

        let ct = mgr.encrypt("ch1", b"round-trip").expect("subscribed").expect("encrypt");
        let pt = mgr.decrypt("ch1", &ct).expect("decrypt");
        assert_eq!(pt, b"round-trip");

        // Unknown topic returns None.
        assert!(mgr.encrypt("unknown", b"x").is_none());
        assert!(mgr.decrypt("unknown", b"x").is_none());
    }

    #[test]
    fn short_ciphertext_rejected() {
        let ch = BroadcastChannel::new("t");
        let result = ch.decrypt(&[0u8; 5]); // less than 12-byte nonce
        assert!(result.is_err());
    }
}
