path: root/network/src/propagation.rs

use crate::dandelion::{DandelionRouter, DandelionPhase, RoutingDecision, PeerId};
use onioncoin_timing::{TimingObfuscator, obfuscation::ObfuscationConfig, DelayStrategy};
use onioncoin_core::Transaction;
use async_trait::async_trait;
use std::sync::Arc;
use tokio::sync::Mutex;

/// Manages transaction propagation with timing obfuscation
pub struct PropagationManager {
    /// Dandelion++ router
    dandelion: Arc<Mutex<DandelionRouter<Vec<u8>>>>,

    /// Timing obfuscator
    obfuscator: Arc<Mutex<TimingObfuscator<Transaction>>>,

    /// Connected peers
    peers: Arc<Mutex<Vec<PeerId>>>,
}

impl PropagationManager {
    pub fn new() -> Self {
        let config = ObfuscationConfig::default();

        Self {
            dandelion: Arc::new(Mutex::new(DandelionRouter::new(Default::default()))),
            obfuscator: Arc::new(Mutex::new(TimingObfuscator::new(config))),
            peers: Arc::new(Mutex::new(Vec::new())),
        }
    }

    /// Receive a new transaction from wallet/user
    pub async fn receive_transaction(&self, tx: Transaction) -> PropagationResult {
        let tx_id = tx.id();

        // Register with Dandelion router
        {
            let mut dandelion = self.dandelion.lock().await;
            dandelion.register_transaction(tx_id.to_vec());
        }

        // Add to mixing pool
        {
            let mut obfuscator = self.obfuscator.lock().await;
            obfuscator.add_to_pool(tx);
        }

        PropagationResult::Queued
    }

    /// Process pending transactions for propagation
    pub async fn process_pending(&self) -> Vec<(Transaction, RoutingDecision)> {
        let mut results = Vec::new();

        // Check if should release batch from mixing pool
        let should_release = {
            let obfuscator = self.obfuscator.lock().await;
            obfuscator.should_release_batch()
        };

        if !should_release {
            return results;
        }

        // Release batch
        let batch = {
            let mut obfuscator = self.obfuscator.lock().await;
            obfuscator.release_batch()
        };

        // Route each transaction
        for tx in batch {
            let tx_id = tx.id();
            let routing = self.route_transaction(&tx_id).await;
            results.push((tx, routing));
        }

        results
    }

    /// Route a single transaction
    async fn route_transaction(&self, tx_id: &[u8]) -> RoutingDecision {
        let mut dandelion = self.dandelion.lock().await;
        let phase = dandelion.process_hop(&tx_id.to_vec());

        match phase {
            DandelionPhase::Stem => {
                // Select random peer for STEM
                let peer = self.select_random_peer().await;
                let delay = DelayStrategy::dandelion_stem();
                RoutingDecision::stem(peer, delay)
            }

            DandelionPhase::Fluff => {
                // Broadcast to random subset of peers
                let peers = self.peers.lock().await;
                let subset_size = (peers.len() as f64).sqrt() as usize;
                RoutingDecision::fluff_random(subset_size.max(1))
            }
        }
    }

    /// Select random peer for STEM routing
    async fn select_random_peer(&self) -> PeerId {
        use rand::seq::SliceRandom;

        let peers = self.peers.lock().await;
        peers
            .choose(&mut rand::thread_rng())
            .cloned()
            .unwrap_or_else(|| "default".to_string())
    }

    /// Add peer to connection pool
    pub async fn add_peer(&self, peer: PeerId) {
        let mut peers = self.peers.lock().await;
        if !peers.contains(&peer) {
            peers.push(peer);
        }
    }

    /// Remove peer from connection pool
    pub async fn remove_peer(&self, peer: &PeerId) {
        let mut peers = self.peers.lock().await;
        peers.retain(|p| p != peer);
    }

    /// Get number of connected peers
    pub async fn peer_count(&self) -> usize {
        self.peers.lock().await.len()
    }
}

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

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PropagationResult {
    /// Transaction queued for propagation
    Queued,

    /// Transaction propagated immediately
    Propagated,

    /// Transaction rejected
    Rejected,
}

/// Strategy for propagating transactions
#[async_trait]
pub trait PropagationStrategy: Send + Sync {
    /// Decide how to propagate a transaction
    async fn route(&self, tx: &Transaction) -> RoutingDecision;

    /// Apply timing delays
    async fn apply_delay(&self, delay: &DelayStrategy);
}

/// Standard propagation strategy (Dandelion++ with timing obfuscation)
pub struct StandardPropagation {
    manager: Arc<PropagationManager>,
}

impl StandardPropagation {
    pub fn new(manager: Arc<PropagationManager>) -> Self {
        Self { manager }
    }
}

#[async_trait]
impl PropagationStrategy for StandardPropagation {
    async fn route(&self, tx: &Transaction) -> RoutingDecision {
        let tx_id = tx.id();
        self.manager.route_transaction(&tx_id).await
    }

    async fn apply_delay(&self, delay: &DelayStrategy) {
        delay.sleep().await;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use onioncoin_timing::TimingMetadata;
    use chrono::Utc;

    fn create_test_tx() -> Transaction {
        use onioncoin_core::transaction::*;

        let seed = [42u8; 32];
        let timing = TimingMetadata::new(Utc::now(), &seed).unwrap();

        Transaction {
            version: 1,
            inputs: vec![TransactionInput {
                previous_output: OutPoint {
                    txid: [1u8; 32],
                    vout: 0,
                },
                key_image: [2u8; 32],
            }],
            outputs: vec![TransactionOutput {
                encrypted_amount: vec![0u8; 32],
                stealth_address: [3u8; 32],
                commitment: [4u8; 32],
            }],
            ring_signatures: vec![RingSignature::new(11)],
            range_proof: vec![0u8; 64],
            encrypted_fee: 1000,
            timing,
        }
    }

    #[tokio::test]
    async fn test_propagation_manager() {
        let manager = PropagationManager::new();
        let tx = create_test_tx();

        let result = manager.receive_transaction(tx).await;
        assert_eq!(result, PropagationResult::Queued);
    }

    #[tokio::test]
    async fn test_peer_management() {
        let manager = PropagationManager::new();

        manager.add_peer("peer1".to_string()).await;
        manager.add_peer("peer2".to_string()).await;

        assert_eq!(manager.peer_count().await, 2);

        manager.remove_peer(&"peer1".to_string()).await;
        assert_eq!(manager.peer_count().await, 1);
    }
}