blvm-node 0.1.74

//! Congestion Control Manager for Transaction Batching and Fee Optimization
//!
//! Uses CTV covenants to manage transaction throughput and fees:
//! - Transaction batching for reduced fees
//! - Fee rate commitments in covenants
//! - Deferred execution for optimal timing
//! - Dynamic fee adjustment based on network conditions

use crate::node::mempool::MempoolManager;
use crate::payment::covenant::{CovenantEngine, CovenantProof};
use crate::payment::processor::PaymentError;
use crate::rpc::errors::STORAGE_NOT_AVAILABLE_MSG;
use crate::storage::Storage;
use crate::utils::current_timestamp;
use blvm_protocol::payment::PaymentOutput;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tracing::{debug, info};

/// Transaction priority
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub enum TransactionPriority {
    /// Can wait for optimal conditions
    Low = 1,
    /// Standard priority
    Normal = 2,
    /// Should batch soon
    High = 3,
    /// Batch immediately
    Urgent = 4,
}

/// Batch configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchConfig {
    /// Maximum transactions per batch
    pub max_batch_size: usize,
    /// Target fee rate (sat/vbyte)
    pub target_fee_rate: u64,
    /// Maximum wait time before forcing batch (seconds)
    pub max_wait_time: u64,
    /// Minimum fee rate to accept (sat/vbyte)
    pub min_fee_rate: u64,
}

impl Default for BatchConfig {
    fn default() -> Self {
        Self {
            max_batch_size: 100,
            target_fee_rate: 1,  // 1 sat/vbyte target
            max_wait_time: 3600, // 1 hour max wait
            min_fee_rate: 1,     // 1 sat/vbyte minimum
        }
    }
}

/// Pending transaction in batch
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PendingTransaction {
    pub tx_id: String,
    pub outputs: Vec<PaymentOutput>,
    pub priority: TransactionPriority,
    pub created_at: u64,
    pub deadline: Option<u64>, // Optional deadline (Unix timestamp)
}

/// Transaction batch
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransactionBatch {
    pub batch_id: String,
    pub transactions: Vec<PendingTransaction>,
    pub covenant_template: Option<CovenantProof>,
    pub target_fee_rate: u64,
    pub created_at: u64,
    pub ready_to_broadcast: bool,
    pub broadcast_at: Option<u64>,
}

/// Network congestion metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CongestionMetrics {
    /// Current mempool size (transactions)
    pub mempool_size: usize,
    /// Average fee rate in mempool (sat/vbyte)
    pub avg_fee_rate: u64,
    /// Median fee rate in mempool (sat/vbyte)
    pub median_fee_rate: u64,
    /// Estimated blocks until confirmation at current fee rate
    pub estimated_blocks: u32,
    /// Timestamp when metrics were collected
    pub collected_at: u64,
}

/// Congestion Control Manager
pub struct CongestionManager {
    covenant_engine: Arc<CovenantEngine>,
    mempool_manager: Option<Arc<MempoolManager>>,
    storage: Option<Arc<Storage>>,
    batches: HashMap<String, TransactionBatch>,
    config: BatchConfig,
}

impl CongestionManager {
    /// Create a new congestion manager
    pub fn new(
        covenant_engine: Arc<CovenantEngine>,
        mempool_manager: Option<Arc<MempoolManager>>,
        storage: Option<Arc<Storage>>,
        config: BatchConfig,
    ) -> Self {
        let mut manager = Self {
            covenant_engine,
            mempool_manager,
            storage,
            batches: HashMap::new(),
            config,
        };
        // Load existing batches from storage
        if let Err(e) = manager.load_all_batches() {
            tracing::warn!("Failed to load batches from storage: {}", e);
        }
        manager
    }

    /// Load all batches from storage into memory.
    fn load_all_batches(&mut self) -> Result<(), PaymentError> {
        let storage = self
            .storage
            .as_ref()
            .ok_or_else(|| PaymentError::ProcessingError(STORAGE_NOT_AVAILABLE_MSG.to_string()))?;

        let batches_tree = storage.open_tree("batches").map_err(|e| {
            PaymentError::ProcessingError(format!("Failed to open batches tree: {}", e))
        })?;

        let mut loaded = 0usize;
        for entry in batches_tree.iter() {
            let (_key, value) = entry.map_err(|e| {
                PaymentError::ProcessingError(format!("Failed to read batch from storage: {e}"))
            })?;
            let batch: TransactionBatch = bincode::deserialize(&value).map_err(|e| {
                PaymentError::ProcessingError(format!("Failed to deserialize batch: {e}"))
            })?;
            self.batches.insert(batch.batch_id.clone(), batch);
            loaded += 1;
        }

        if loaded > 0 {
            tracing::info!("Loaded {} batches from storage", loaded);
        }
        Ok(())
    }

    /// Save batch state to storage
    fn save_batch(&self, batch: &TransactionBatch) -> Result<(), PaymentError> {
        if let Some(storage) = &self.storage {
            let batches_tree = storage.open_tree("batches").map_err(|e| {
                PaymentError::ProcessingError(format!("Failed to open batches tree: {}", e))
            })?;

            let key = batch.batch_id.as_bytes();
            let value = bincode::serialize(batch).map_err(|e| {
                PaymentError::ProcessingError(format!("Failed to serialize batch: {}", e))
            })?;

            batches_tree.insert(key, &value).map_err(|e| {
                PaymentError::ProcessingError(format!("Failed to save batch: {}", e))
            })?;
        }
        Ok(())
    }

    /// Create a new transaction batch
    ///
    /// # Arguments
    ///
    /// * `batch_id` - Unique identifier for the batch
    /// * `target_fee_rate` - Target fee rate for the batch
    ///
    /// # Returns
    ///
    /// Batch ID
    pub fn create_batch(&mut self, batch_id: &str, target_fee_rate: Option<u64>) -> String {
        let target = target_fee_rate.unwrap_or(self.config.target_fee_rate);
        let created_at = crate::utils::current_timestamp();

        let batch = TransactionBatch {
            batch_id: batch_id.to_string(),
            transactions: Vec::new(),
            covenant_template: None,
            target_fee_rate: target,
            created_at,
            ready_to_broadcast: false,
            broadcast_at: None,
        };

        self.batches.insert(batch_id.to_string(), batch);

        info!(
            "Created transaction batch: {} with target fee rate: {}",
            batch_id, target
        );

        batch_id.to_string()
    }

    /// Add transaction to batch
    ///
    /// # Arguments
    ///
    /// * `batch_id` - Batch ID
    /// * `tx` - Pending transaction to add
    ///
    /// # Returns
    ///
    /// Updated batch or error
    pub fn add_to_batch(
        &mut self,
        batch_id: &str,
        tx: PendingTransaction,
    ) -> Result<(), PaymentError> {
        let (batch_clone, tx_count) = {
            let batch = self.batches.get_mut(batch_id).ok_or_else(|| {
                PaymentError::ProcessingError(format!("Batch {} not found", batch_id))
            })?;

            if batch.transactions.len() >= self.config.max_batch_size {
                return Err(PaymentError::ProcessingError(format!(
                    "Batch {} is full: {} transactions",
                    batch_id, self.config.max_batch_size
                )));
            }

            batch.transactions.push(tx);
            let batch_clone = batch.clone();
            let tx_count = batch.transactions.len();
            (batch_clone, tx_count)
        };

        // Save batch to storage
        if let Err(e) = self.save_batch(&batch_clone) {
            tracing::warn!("Failed to save batch {} to storage: {}", batch_id, e);
        }

        // Update covenant template if batch is ready
        if tx_count >= self.config.max_batch_size {
            self.update_batch_covenant(batch_id)?;
        }

        Ok(())
    }

    /// Drop one transaction from a batch. The remaining transactions stay.
    pub fn remove_from_batch(&mut self, batch_id: &str, tx_id: &str) -> Result<(), PaymentError> {
        let batch_clone = {
            let batch = self.batches.get_mut(batch_id).ok_or_else(|| {
                PaymentError::ProcessingError(format!("Batch {batch_id} not found"))
            })?;
            let before = batch.transactions.len();
            batch.transactions.retain(|tx| tx.tx_id != tx_id);
            if batch.transactions.len() == before {
                return Err(PaymentError::ProcessingError(format!(
                    "transaction {tx_id} is not in batch {batch_id}"
                )));
            }
            batch.clone()
        };
        if let Err(e) = self.save_batch(&batch_clone) {
            tracing::warn!("Failed to save batch {batch_id} after removal: {e}");
        }
        Ok(())
    }

    /// Check current network congestion
    ///
    /// # Returns
    ///
    /// Congestion metrics
    pub fn check_congestion(&self) -> Result<CongestionMetrics, PaymentError> {
        let mempool_manager = self.mempool_manager.as_ref().ok_or_else(|| {
            PaymentError::ProcessingError("Mempool manager not available".to_string())
        })?;

        let mempool_size = mempool_manager.size();

        let mut fee_rates: Vec<u64> = mempool_manager
            .fee_rates_for_storage(self.storage.as_deref())
            .into_iter()
            .filter(|rate| *rate > 0)
            .collect();

        let (avg_fee_rate, median_fee_rate, estimated_blocks) = if fee_rates.is_empty() {
            // No pooled rate was measured. Do not publish the configured minimum as a measurement.
            (0, 0, 0)
        } else {
            fee_rates.sort_unstable();
            let sum: u64 = fee_rates.iter().sum();
            let avg = sum / fee_rates.len() as u64;
            let median = fee_rates[fee_rates.len() / 2];
            let estimated = if avg >= self.config.target_fee_rate {
                1
            } else {
                (self.config.target_fee_rate / avg.max(1)) as u32
            };
            (avg, median, estimated)
        };

        Ok(CongestionMetrics {
            mempool_size,
            avg_fee_rate,
            median_fee_rate,
            estimated_blocks,
            collected_at: current_timestamp(),
        })
    }

    /// Determine if batch should be broadcast
    ///
    /// # Arguments
    ///
    /// * `batch_id` - Batch ID
    ///
    /// # Returns
    ///
    /// `true` if batch should be broadcast, `false` otherwise
    pub fn should_broadcast(&self, batch_id: &str) -> Result<bool, PaymentError> {
        let batch = self.batches.get(batch_id).ok_or_else(|| {
            PaymentError::ProcessingError(format!("Batch {} not found", batch_id))
        })?;

        // Check if batch is empty
        if batch.transactions.is_empty() {
            return Ok(false);
        }

        // Check if batch is full
        if batch.transactions.len() >= self.config.max_batch_size {
            return Ok(true);
        }

        // Check if any transaction has urgent priority
        if batch
            .transactions
            .iter()
            .any(|tx| tx.priority == TransactionPriority::Urgent)
        {
            return Ok(true);
        }

        // Check if max wait time has passed
        let now = current_timestamp();
        if now - batch.created_at >= self.config.max_wait_time {
            return Ok(true);
        }

        // Check if any transaction deadline has passed
        if batch.transactions.iter().any(|tx| {
            if let Some(deadline) = tx.deadline {
                now >= deadline
            } else {
                false
            }
        }) {
            return Ok(true);
        }

        // A measured pool whose average is at or below the target can broadcast.
        // A missing mempool or an empty rate list is not evidence that fees are low.
        if let Ok(metrics) = self.check_congestion() {
            if metrics.avg_fee_rate > 0 && metrics.avg_fee_rate <= batch.target_fee_rate {
                return Ok(true);
            }
        }

        Ok(false)
    }

    /// Broadcast batch when conditions are optimal
    ///
    /// # Arguments
    ///
    /// * `batch_id` - Batch ID
    ///
    /// # Returns
    ///
    /// Covenant proof for the batch
    pub fn broadcast_batch(&mut self, batch_id: &str) -> Result<CovenantProof, PaymentError> {
        #[cfg(not(feature = "ctv"))]
        {
            return Err(PaymentError::FeatureNotEnabled(
                "Congestion control requires CTV feature".to_string(),
            ));
        }

        #[cfg(feature = "ctv")]
        {
            let (needs_covenant, tx_count, target_fee) = {
                let batch = self.batches.get_mut(batch_id).ok_or_else(|| {
                    PaymentError::ProcessingError(format!("Batch {} not found", batch_id))
                })?;

                if batch.transactions.is_empty() {
                    return Err(PaymentError::ProcessingError("Batch is empty".to_string()));
                }

                (
                    batch.covenant_template.is_none(),
                    batch.transactions.len(),
                    batch.target_fee_rate,
                )
            };

            // Update covenant template if not already created
            if needs_covenant {
                self.update_batch_covenant(batch_id)?;
            }

            // Get covenant and update batch state
            let covenant = {
                let batch = self.batches.get_mut(batch_id).ok_or_else(|| {
                    PaymentError::ProcessingError(format!("Batch {} not found", batch_id))
                })?;

                batch.covenant_template.as_ref().ok_or_else(|| {
                    PaymentError::ProcessingError("Batch covenant not created".to_string())
                })?;

                batch.covenant_template.clone().ok_or_else(|| {
                    PaymentError::ProcessingError("Batch covenant not created".to_string())
                })?
            };

            let ready = self.should_broadcast(batch_id)?;
            if !ready {
                return Ok(covenant);
            }

            self.enqueue_batch_transaction(&covenant)?;
            let stored = {
                let batch = self.batches.get_mut(batch_id).ok_or_else(|| {
                    PaymentError::ProcessingError(format!("Batch {} not found", batch_id))
                })?;
                batch.ready_to_broadcast = true;
                batch.broadcast_at = Some(current_timestamp());
                batch.clone()
            };
            if let Err(e) = self.save_batch(&stored) {
                tracing::warn!("Failed to save batch {batch_id} after broadcast: {e}");
            }

            info!(
                "Batch {} ready to broadcast with {} transactions at fee rate {} sat/vbyte",
                batch_id, tx_count, target_fee
            );

            Ok(covenant)
        }
    }

    /// Put the covenant template into the mempool. The template's inputs are the spends.
    #[cfg(feature = "ctv")]
    fn enqueue_batch_transaction(&self, covenant: &CovenantProof) -> Result<(), PaymentError> {
        use blvm_protocol::{OutPoint, Transaction, TransactionInput, TransactionOutput};

        let mempool = self.mempool_manager.as_ref().ok_or_else(|| {
            PaymentError::ProcessingError("Mempool manager not available".to_string())
        })?;
        let template = &covenant.transaction_template;
        let tx = Transaction {
            version: template.version as u64,
            inputs: template
                .inputs
                .iter()
                .map(|input| TransactionInput {
                    prevout: OutPoint {
                        hash: input.prevout_hash,
                        index: input.prevout_index,
                    },
                    script_sig: Vec::new(),
                    sequence: input.sequence as u64,
                })
                .collect::<Vec<_>>()
                .into(),
            outputs: template
                .outputs
                .iter()
                .map(|output| TransactionOutput {
                    value: output.value as i64,
                    script_pubkey: output.script_pubkey.clone(),
                })
                .collect::<Vec<_>>()
                .into(),
            lock_time: template.lock_time as u64,
        };
        let accepted = mempool.add_transaction(tx).map_err(|e| {
            PaymentError::ProcessingError(format!("batch transaction was not accepted: {e}"))
        })?;
        if !accepted {
            return Err(PaymentError::ProcessingError(
                "batch transaction was not accepted".to_string(),
            ));
        }
        Ok(())
    }

    /// Adjust fee rate based on network conditions
    ///
    /// # Arguments
    ///
    /// * `batch_id` - Batch ID
    ///
    /// # Returns
    ///
    /// Updated target fee rate
    pub fn adjust_fee_rate(&mut self, batch_id: &str) -> Result<u64, PaymentError> {
        let metrics = self.check_congestion()?;

        let batch = self.batches.get_mut(batch_id).ok_or_else(|| {
            PaymentError::ProcessingError(format!("Batch {} not found", batch_id))
        })?;

        // Adjust fee rate based on network conditions
        let new_fee_rate = if metrics.avg_fee_rate > batch.target_fee_rate {
            // Network is congested, increase fee rate
            metrics.avg_fee_rate.max(self.config.min_fee_rate)
        } else {
            // Network is not congested, use target
            batch.target_fee_rate
        };

        batch.target_fee_rate = new_fee_rate;

        debug!(
            "Adjusted fee rate for batch {} to {} sat/vbyte",
            batch_id, new_fee_rate
        );

        Ok(new_fee_rate)
    }

    /// Get batch by ID
    pub fn get_batch(&self, batch_id: &str) -> Option<&TransactionBatch> {
        self.batches.get(batch_id)
    }

    /// List all batches
    pub fn list_batches(&self) -> Vec<&TransactionBatch> {
        self.batches.values().collect()
    }

    /// Update batch covenant template
    fn update_batch_covenant(&mut self, batch_id: &str) -> Result<(), PaymentError> {
        let batch = self.batches.get_mut(batch_id).ok_or_else(|| {
            PaymentError::ProcessingError(format!("Batch {} not found", batch_id))
        })?;

        // Combine all transaction outputs
        let mut all_outputs = Vec::new();
        for tx in &batch.transactions {
            all_outputs.extend(tx.outputs.clone());
        }

        // Create covenant template
        let covenant = self.covenant_engine.create_payment_covenant(
            &format!("{}_batch", batch_id),
            &all_outputs,
            None,
        )?;

        batch.covenant_template = Some(covenant);

        Ok(())
    }
}