path: root/network/src/dandelion.rs

use onioncoin_timing::{DelayStrategy, delay::DelaySequence};
use rand::Rng;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Dandelion++ phases for transaction propagation
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum DandelionPhase {
    /// STEM phase: forward to single peer with delays
    Stem,
    /// FLUFF phase: broadcast to all peers
    Fluff,
}

/// Router for Dandelion++ protocol
#[derive(Debug)]
pub struct DandelionRouter<TxId> {
    /// Current phase for each transaction
    phases: HashMap<TxId, DandelionPhase>,

    /// Hop count for STEM phase
    stem_hops: HashMap<TxId, u32>,

    /// Configuration
    config: DandelionConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DandelionConfig {
    /// Minimum STEM hops before transitioning to FLUFF
    pub min_stem_hops: u32,

    /// Maximum STEM hops before forcing FLUFF
    pub max_stem_hops: u32,

    /// Probability of transitioning to FLUFF at each hop (0.0 - 1.0)
    pub fluff_probability: f64,

    /// Enable delays between hops
    pub use_delays: bool,
}

impl Default for DandelionConfig {
    fn default() -> Self {
        Self {
            min_stem_hops: 1,
            max_stem_hops: 4,
            fluff_probability: 0.25, // 25% chance per hop after min
            use_delays: true,
        }
    }
}

impl<TxId: std::hash::Hash + Eq + Clone> DandelionRouter<TxId> {
    pub fn new(config: DandelionConfig) -> Self {
        Self {
            phases: HashMap::new(),
            stem_hops: HashMap::new(),
            config,
        }
    }

    /// Register a new transaction (starts in STEM phase)
    pub fn register_transaction(&mut self, tx_id: TxId) {
        self.phases.insert(tx_id.clone(), DandelionPhase::Stem);
        self.stem_hops.insert(tx_id, 0);
    }

    /// Process a hop and determine if should transition to FLUFF
    pub fn process_hop(&mut self, tx_id: &TxId) -> DandelionPhase {
        let current_phase = self.phases.get(tx_id).copied().unwrap_or(DandelionPhase::Stem);

        if current_phase == DandelionPhase::Fluff {
            return DandelionPhase::Fluff;
        }

        // Increment hop count
        let hops = self.stem_hops.entry(tx_id.clone()).or_insert(0);
        *hops += 1;

        // Force FLUFF if max hops reached
        if *hops >= self.config.max_stem_hops {
            self.phases.insert(tx_id.clone(), DandelionPhase::Fluff);
            return DandelionPhase::Fluff;
        }

        // If past min hops, randomly transition to FLUFF
        if *hops >= self.config.min_stem_hops {
            let mut rng = rand::thread_rng();
            if rng.gen_bool(self.config.fluff_probability) {
                self.phases.insert(tx_id.clone(), DandelionPhase::Fluff);
                return DandelionPhase::Fluff;
            }
        }

        DandelionPhase::Stem
    }

    /// Get current phase for a transaction
    pub fn get_phase(&self, tx_id: &TxId) -> Option<DandelionPhase> {
        self.phases.get(tx_id).copied()
    }

    /// Force transition to FLUFF phase
    pub fn force_fluff(&mut self, tx_id: &TxId) {
        self.phases.insert(tx_id.clone(), DandelionPhase::Fluff);
    }

    /// Get number of STEM hops for a transaction
    pub fn get_stem_hops(&self, tx_id: &TxId) -> u32 {
        self.stem_hops.get(tx_id).copied().unwrap_or(0)
    }

    /// Clean up old transactions
    pub fn cleanup(&mut self, tx_ids: &[TxId]) {
        for tx_id in tx_ids {
            self.phases.remove(tx_id);
            self.stem_hops.remove(tx_id);
        }
    }

    /// Generate delay sequence for STEM phase
    pub fn create_stem_delays(&self) -> DelaySequence {
        DelaySequence::dandelion_stem()
    }
}

/// Represents a routing decision for Dandelion++
#[derive(Debug, Clone)]
pub struct RoutingDecision {
    /// Phase to use
    pub phase: DandelionPhase,

    /// Target peers (1 for STEM, multiple for FLUFF)
    pub target_peers: TargetPeers,

    /// Delay before forwarding
    pub delay: Option<DelayStrategy>,
}

#[derive(Debug, Clone)]
pub enum TargetPeers {
    /// Single peer (for STEM)
    Single(PeerId),

    /// Multiple peers (for FLUFF broadcast)
    Multiple(Vec<PeerId>),

    /// Random subset of peers
    RandomSubset { count: usize },
}

/// Placeholder for peer identifier
pub type PeerId = String;

impl RoutingDecision {
    /// Create STEM routing decision
    pub fn stem(peer: PeerId, delay: DelayStrategy) -> Self {
        Self {
            phase: DandelionPhase::Stem,
            target_peers: TargetPeers::Single(peer),
            delay: Some(delay),
        }
    }

    /// Create FLUFF routing decision
    pub fn fluff(peers: Vec<PeerId>, delay: Option<DelayStrategy>) -> Self {
        Self {
            phase: DandelionPhase::Fluff,
            target_peers: TargetPeers::Multiple(peers),
            delay,
        }
    }

    /// Create random subset FLUFF
    pub fn fluff_random(count: usize) -> Self {
        Self {
            phase: DandelionPhase::Fluff,
            target_peers: TargetPeers::RandomSubset { count },
            delay: None,
        }
    }
}

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

    #[test]
    fn test_dandelion_registration() {
        let mut router = DandelionRouter::<u64>::new(DandelionConfig::default());
        router.register_transaction(1);

        assert_eq!(router.get_phase(&1), Some(DandelionPhase::Stem));
        assert_eq!(router.get_stem_hops(&1), 0);
    }

    #[test]
    fn test_dandelion_hop_progression() {
        let mut router = DandelionRouter::<u64>::new(DandelionConfig::default());
        router.register_transaction(1);

        // Process hops
        for _ in 0..3 {
            router.process_hop(&1);
        }

        assert!(router.get_stem_hops(&1) >= 3);
    }

    #[test]
    fn test_dandelion_force_fluff() {
        let config = DandelionConfig {
            min_stem_hops: 1,
            max_stem_hops: 10,
            fluff_probability: 0.0, // Never random transition
            use_delays: true,
        };

        let mut router = DandelionRouter::<u64>::new(config);
        router.register_transaction(1);

        // Should stay in STEM
        router.process_hop(&1);
        assert_eq!(router.get_phase(&1), Some(DandelionPhase::Stem));

        // Force FLUFF
        router.force_fluff(&1);
        assert_eq!(router.get_phase(&1), Some(DandelionPhase::Fluff));
    }

    #[test]
    fn test_dandelion_max_hops() {
        let config = DandelionConfig {
            min_stem_hops: 1,
            max_stem_hops: 3,
            fluff_probability: 0.0,
            use_delays: true,
        };

        let mut router = DandelionRouter::<u64>::new(config);
        router.register_transaction(1);

        // Process max hops
        for _ in 0..3 {
            router.process_hop(&1);
        }

        // Should transition to FLUFF at max hops
        assert_eq!(router.get_phase(&1), Some(DandelionPhase::Fluff));
    }

    #[test]
    fn test_routing_decision() {
        let stem_decision = RoutingDecision::stem(
            "peer1".to_string(),
            DelayStrategy::dandelion_stem(),
        );

        assert_eq!(stem_decision.phase, DandelionPhase::Stem);
        assert!(stem_decision.delay.is_some());

        let fluff_decision = RoutingDecision::fluff(
            vec!["peer1".to_string(), "peer2".to_string()],
            None,
        );

        assert_eq!(fluff_decision.phase, DandelionPhase::Fluff);
    }
}