kaccy-core 0.2.0

//! MEV Mitigation Strategies
//!
//! This module implements various strategies to mitigate MEV extraction,
//! protecting users from sandwich attacks, front-running, and other MEV exploits.

use crate::error::CoreError;
use rust_decimal::Decimal;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use std::time::{Duration, SystemTime};

/// Types of MEV mitigation strategies
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MitigationStrategy {
    /// Bundle submission - transactions executed together atomically
    BundleSubmission,
    /// Transaction ordering protection - fair ordering rules
    FairOrdering,
    /// MEV rebate - share MEV profits with users
    MevRebate,
    /// Privacy-preserving execution - hide transaction details
    PrivateExecution,
    /// Time-delayed execution - randomized execution time
    TimeDelayed,
    /// Threshold encryption - decrypt only when threshold reached
    ThresholdEncryption,
}

/// Configuration for MEV rebate mechanism
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MevRebateConfig {
    /// Percentage of detected MEV to rebate to victims (0-100)
    pub rebate_percentage: Decimal,
    /// Minimum MEV amount to trigger rebate
    pub min_rebate_amount: Decimal,
    /// Duration to wait before processing rebate
    pub processing_delay: Duration,
}

impl Default for MevRebateConfig {
    fn default() -> Self {
        Self {
            rebate_percentage: Decimal::new(50, 0),      // 50% rebate
            min_rebate_amount: Decimal::new(1, 3),       // 0.001 BTC
            processing_delay: Duration::from_secs(3600), // 1 hour
        }
    }
}

/// MEV rebate record
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MevRebate {
    /// Unique identifier for this rebate
    pub rebate_id: String,
    /// User who was harmed by MEV extraction
    pub victim_user_id: String,
    /// User who performed MEV extraction
    pub attacker_user_id: String,
    /// Total MEV amount extracted
    pub mev_amount: Decimal,
    /// Amount to be rebated to the victim
    pub rebate_amount: Decimal,
    /// Transaction IDs involved in the MEV event
    pub transaction_ids: Vec<String>,
    /// When this rebate record was created
    pub created_at: SystemTime,
    /// When this rebate was processed
    pub processed_at: Option<SystemTime>,
    /// Current processing status
    pub status: RebateStatus,
}

/// Status of an MEV rebate
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RebateStatus {
    /// Rebate is queued and awaiting processing
    Pending,
    /// Rebate is currently being processed
    Processing,
    /// Rebate has been successfully paid out
    Completed,
    /// Rebate processing failed
    Failed,
}

/// MEV rebate manager
///
/// Manages the distribution of MEV rebates to victims of MEV extraction
pub struct MevRebateManager {
    /// Rebate configuration
    config: MevRebateConfig,
    /// All rebates indexed by rebate ID
    pending_rebates: HashMap<String, MevRebate>,
    /// Monotonic counter used to generate unique rebate IDs
    rebate_counter: u64,
}

impl MevRebateManager {
    /// Create a new MEV rebate manager with the given configuration
    pub fn new(config: MevRebateConfig) -> Self {
        Self {
            config,
            pending_rebates: HashMap::new(),
            rebate_counter: 0,
        }
    }

    /// Create a new MEV rebate manager with default configuration
    pub fn with_defaults() -> Self {
        Self::new(MevRebateConfig::default())
    }

    /// Creates a new rebate for a victim of MEV extraction
    pub fn create_rebate(
        &mut self,
        victim_user_id: String,
        attacker_user_id: String,
        mev_amount: Decimal,
        transaction_ids: Vec<String>,
    ) -> Result<MevRebate, CoreError> {
        if mev_amount < self.config.min_rebate_amount {
            return Err(CoreError::Validation(
                "MEV amount below minimum threshold".to_string(),
            ));
        }

        let rebate_amount = mev_amount * self.config.rebate_percentage / Decimal::new(100, 0);

        self.rebate_counter += 1;
        let rebate_id = format!("rebate_{}", self.rebate_counter);

        let rebate = MevRebate {
            rebate_id: rebate_id.clone(),
            victim_user_id,
            attacker_user_id,
            mev_amount,
            rebate_amount,
            transaction_ids,
            created_at: SystemTime::now(),
            processed_at: None,
            status: RebateStatus::Pending,
        };

        self.pending_rebates.insert(rebate_id, rebate.clone());
        Ok(rebate)
    }

    /// Processes pending rebates that have passed the processing delay
    pub fn process_pending_rebates(&mut self) -> Vec<MevRebate> {
        let now = SystemTime::now();
        let mut processed = Vec::new();

        for rebate in self.pending_rebates.values_mut() {
            if rebate.status != RebateStatus::Pending {
                continue;
            }

            let elapsed = now
                .duration_since(rebate.created_at)
                .unwrap_or(Duration::ZERO);

            if elapsed >= self.config.processing_delay {
                rebate.status = RebateStatus::Processing;
                processed.push(rebate.clone());
            }
        }

        processed
    }

    /// Marks a rebate as completed
    pub fn complete_rebate(&mut self, rebate_id: &str) -> Result<(), CoreError> {
        let rebate = self
            .pending_rebates
            .get_mut(rebate_id)
            .ok_or_else(|| CoreError::NotFound(format!("Rebate {} not found", rebate_id)))?;

        rebate.status = RebateStatus::Completed;
        rebate.processed_at = Some(SystemTime::now());
        Ok(())
    }

    /// Marks a rebate as failed
    pub fn fail_rebate(&mut self, rebate_id: &str) -> Result<(), CoreError> {
        let rebate = self
            .pending_rebates
            .get_mut(rebate_id)
            .ok_or_else(|| CoreError::NotFound(format!("Rebate {} not found", rebate_id)))?;

        rebate.status = RebateStatus::Failed;
        Ok(())
    }

    /// Gets all rebates for a specific victim
    pub fn get_victim_rebates(&self, user_id: &str) -> Vec<&MevRebate> {
        self.pending_rebates
            .values()
            .filter(|r| r.victim_user_id == user_id)
            .collect()
    }

    /// Gets total rebate amount for a victim
    pub fn get_total_rebate_amount(&self, user_id: &str) -> Decimal {
        self.pending_rebates
            .values()
            .filter(|r| r.victim_user_id == user_id && r.status == RebateStatus::Completed)
            .map(|r| r.rebate_amount)
            .sum()
    }
}

/// Transaction bundle for atomic execution
///
/// Bundles multiple transactions together to prevent sandwich attacks
/// by ensuring they execute atomically in a specific order.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransactionBundle {
    /// Unique identifier for this bundle
    pub bundle_id: String,
    /// Ordered list of transactions in this bundle
    pub transactions: Vec<BundledTransaction>,
    /// Earliest acceptable execution time
    pub min_timestamp: Option<SystemTime>,
    /// Latest acceptable execution time
    pub max_timestamp: Option<SystemTime>,
    /// Whether to revert all transactions if any one fails
    pub revert_on_failure: bool,
}

/// A single transaction inside a bundle
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BundledTransaction {
    /// Transaction identifier
    pub tx_id: String,
    /// User who submitted this transaction
    pub user_id: String,
    /// Position of this transaction within the bundle (zero-indexed)
    pub order_index: usize,
    /// Raw encoded transaction data
    pub data: Vec<u8>,
}

impl TransactionBundle {
    /// Create an empty transaction bundle with the given ID
    pub fn new(bundle_id: String) -> Self {
        Self {
            bundle_id,
            transactions: Vec::new(),
            min_timestamp: None,
            max_timestamp: None,
            revert_on_failure: true,
        }
    }

    /// Adds a transaction to the bundle
    pub fn add_transaction(&mut self, tx: BundledTransaction) -> Result<(), CoreError> {
        // Ensure order indices are sequential
        let expected_index = self.transactions.len();
        if tx.order_index != expected_index {
            return Err(CoreError::Validation(format!(
                "Invalid order index. Expected {}, got {}",
                expected_index, tx.order_index
            )));
        }

        self.transactions.push(tx);
        Ok(())
    }

    /// Validates the bundle
    pub fn validate(&self) -> Result<(), CoreError> {
        if self.transactions.is_empty() {
            return Err(CoreError::Validation("Bundle cannot be empty".to_string()));
        }

        // Check timestamp constraints
        if let (Some(min), Some(max)) = (self.min_timestamp, self.max_timestamp) {
            if min > max {
                return Err(CoreError::Validation(
                    "min_timestamp cannot be greater than max_timestamp".to_string(),
                ));
            }
        }

        Ok(())
    }

    /// Sets time constraints for bundle execution
    pub fn with_time_constraints(
        mut self,
        min_timestamp: SystemTime,
        max_timestamp: SystemTime,
    ) -> Self {
        self.min_timestamp = Some(min_timestamp);
        self.max_timestamp = Some(max_timestamp);
        self
    }
}

/// Fair ordering manager
///
/// Implements fair transaction ordering to prevent front-running
#[derive(Debug)]
pub struct FairOrderingManager {
    /// Transactions waiting to be included in the next batch
    pending_queue: VecDeque<PendingTransaction>,
    /// How long to accumulate transactions before ordering and dispatching them
    batch_interval: Duration,
    /// Timestamp of the most recently dispatched batch
    last_batch_time: SystemTime,
}

#[derive(Debug, Clone)]
struct PendingTransaction {
    tx_id: String,
    #[allow(dead_code)]
    user_id: String,
    received_at: SystemTime,
    #[allow(dead_code)]
    priority_score: Decimal,
}

impl FairOrderingManager {
    /// Create a new fair ordering manager with the specified batch interval
    pub fn new(batch_interval: Duration) -> Self {
        Self {
            pending_queue: VecDeque::new(),
            batch_interval,
            last_batch_time: SystemTime::now(),
        }
    }

    /// Create a fair ordering manager with a 10-second batch interval
    pub fn with_defaults() -> Self {
        Self::new(Duration::from_secs(10)) // 10-second batches
    }

    /// Adds a transaction to the pending queue
    pub fn submit_transaction(&mut self, tx_id: String, user_id: String) -> Result<(), CoreError> {
        let tx = PendingTransaction {
            tx_id,
            user_id,
            received_at: SystemTime::now(),
            priority_score: Decimal::ZERO,
        };

        self.pending_queue.push_back(tx);
        Ok(())
    }

    /// Processes a batch of transactions using fair ordering
    pub fn process_batch(&mut self) -> Result<Vec<String>, CoreError> {
        let now = SystemTime::now();
        let elapsed = now
            .duration_since(self.last_batch_time)
            .unwrap_or(Duration::ZERO);

        if elapsed < self.batch_interval {
            return Ok(Vec::new()); // Not time for next batch yet
        }

        // Collect all transactions in the current batch window
        let mut batch = Vec::new();
        while let Some(tx) = self.pending_queue.pop_front() {
            batch.push(tx);
        }

        if batch.is_empty() {
            return Ok(Vec::new());
        }

        // Sort by received time (FIFO) - this is the fair ordering
        batch.sort_by(|a, b| a.received_at.cmp(&b.received_at));

        // Extract transaction IDs in fair order
        let ordered_tx_ids: Vec<String> = batch.iter().map(|tx| tx.tx_id.clone()).collect();

        self.last_batch_time = now;
        Ok(ordered_tx_ids)
    }

    /// Gets the number of pending transactions
    pub fn pending_count(&self) -> usize {
        self.pending_queue.len()
    }

    /// Gets time until next batch
    pub fn time_until_next_batch(&self) -> Duration {
        let elapsed = SystemTime::now()
            .duration_since(self.last_batch_time)
            .unwrap_or(Duration::ZERO);

        self.batch_interval.saturating_sub(elapsed)
    }
}

/// Privacy-preserving execution engine
///
/// Provides mechanisms for executing transactions without revealing details
/// until after execution, preventing MEV extraction.
#[derive(Debug)]
pub struct PrivateExecutionEngine {
    /// Encrypted transactions awaiting decryption and execution
    encrypted_pool: HashMap<String, EncryptedTransaction>,
}

#[derive(Debug, Clone)]
struct EncryptedTransaction {
    #[allow(dead_code)]
    tx_id: String,
    encrypted_data: Vec<u8>,
    #[allow(dead_code)]
    submitted_at: SystemTime,
    #[allow(dead_code)]
    decryption_key_hash: String,
}

impl PrivateExecutionEngine {
    /// Create an empty private execution engine
    pub fn new() -> Self {
        Self {
            encrypted_pool: HashMap::new(),
        }
    }

    /// Submits an encrypted transaction
    pub fn submit_encrypted_transaction(
        &mut self,
        tx_id: String,
        encrypted_data: Vec<u8>,
        decryption_key_hash: String,
    ) -> Result<(), CoreError> {
        let encrypted_tx = EncryptedTransaction {
            tx_id: tx_id.clone(),
            encrypted_data,
            submitted_at: SystemTime::now(),
            decryption_key_hash,
        };

        self.encrypted_pool.insert(tx_id, encrypted_tx);
        Ok(())
    }

    /// Decrypts and executes a transaction
    /// In a real implementation, this would verify the decryption key
    pub fn decrypt_and_execute(
        &mut self,
        tx_id: &str,
        _decryption_key: &[u8],
    ) -> Result<Vec<u8>, CoreError> {
        let encrypted_tx = self
            .encrypted_pool
            .remove(tx_id)
            .ok_or_else(|| CoreError::NotFound(format!("Transaction {} not found", tx_id)))?;

        // In a real implementation, verify the key hash and decrypt
        // For now, we just return the encrypted data as placeholder
        Ok(encrypted_tx.encrypted_data)
    }

    /// Gets the number of encrypted transactions in the pool
    pub fn pool_size(&self) -> usize {
        self.encrypted_pool.len()
    }
}

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

/// MEV protection coordinator
///
/// Coordinates multiple MEV mitigation strategies
pub struct MevProtectionCoordinator {
    /// Manager for MEV rebates to victims
    rebate_manager: MevRebateManager,
    /// Fair ordering component
    fair_ordering: FairOrderingManager,
    /// Private execution component
    private_execution: PrivateExecutionEngine,
    /// Which mitigation strategies are active
    enabled_strategies: Vec<MitigationStrategy>,
}

impl MevProtectionCoordinator {
    /// Create a coordinator with the specified set of active strategies
    pub fn new(enabled_strategies: Vec<MitigationStrategy>) -> Self {
        Self {
            rebate_manager: MevRebateManager::with_defaults(),
            fair_ordering: FairOrderingManager::with_defaults(),
            private_execution: PrivateExecutionEngine::new(),
            enabled_strategies,
        }
    }

    /// Create a coordinator with all available strategies enabled
    pub fn with_all_strategies() -> Self {
        Self::new(vec![
            MitigationStrategy::FairOrdering,
            MitigationStrategy::MevRebate,
            MitigationStrategy::PrivateExecution,
        ])
    }

    /// Checks if a specific strategy is enabled
    pub fn is_strategy_enabled(&self, strategy: MitigationStrategy) -> bool {
        self.enabled_strategies.contains(&strategy)
    }

    /// Gets the rebate manager
    pub fn rebate_manager(&mut self) -> &mut MevRebateManager {
        &mut self.rebate_manager
    }

    /// Gets the fair ordering manager
    pub fn fair_ordering(&mut self) -> &mut FairOrderingManager {
        &mut self.fair_ordering
    }

    /// Gets the private execution engine
    pub fn private_execution(&mut self) -> &mut PrivateExecutionEngine {
        &mut self.private_execution
    }
}

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

    #[test]
    fn test_mev_rebate_creation() {
        let mut manager = MevRebateManager::with_defaults();

        let rebate = manager
            .create_rebate(
                "victim".to_string(),
                "attacker".to_string(),
                Decimal::new(100, 2),
                vec!["tx1".to_string(), "tx2".to_string()],
            )
            .unwrap();

        assert_eq!(rebate.victim_user_id, "victim");
        assert_eq!(rebate.rebate_amount, Decimal::new(50, 2)); // 50% of 1.00
        assert_eq!(rebate.status, RebateStatus::Pending);
    }

    #[test]
    fn test_rebate_below_minimum() {
        let mut manager = MevRebateManager::with_defaults();

        let result = manager.create_rebate(
            "victim".to_string(),
            "attacker".to_string(),
            Decimal::new(1, 4), // 0.0001 BTC, below minimum
            vec!["tx1".to_string()],
        );

        assert!(result.is_err());
    }

    #[test]
    fn test_transaction_bundle() {
        let mut bundle = TransactionBundle::new("bundle1".to_string());

        let tx1 = BundledTransaction {
            tx_id: "tx1".to_string(),
            user_id: "user1".to_string(),
            order_index: 0,
            data: vec![1, 2, 3],
        };

        let tx2 = BundledTransaction {
            tx_id: "tx2".to_string(),
            user_id: "user2".to_string(),
            order_index: 1,
            data: vec![4, 5, 6],
        };

        assert!(bundle.add_transaction(tx1).is_ok());
        assert!(bundle.add_transaction(tx2).is_ok());
        assert!(bundle.validate().is_ok());
        assert_eq!(bundle.transactions.len(), 2);
    }

    #[test]
    fn test_fair_ordering() {
        let mut manager = FairOrderingManager::new(Duration::from_millis(100));

        manager
            .submit_transaction("tx1".to_string(), "user1".to_string())
            .unwrap();
        std::thread::sleep(Duration::from_millis(50));
        manager
            .submit_transaction("tx2".to_string(), "user2".to_string())
            .unwrap();

        std::thread::sleep(Duration::from_millis(100));

        let ordered = manager.process_batch().unwrap();
        assert_eq!(ordered.len(), 2);
        assert_eq!(ordered[0], "tx1"); // First in, first out
        assert_eq!(ordered[1], "tx2");
    }

    #[test]
    fn test_private_execution_engine() {
        let mut engine = PrivateExecutionEngine::new();

        let encrypted_data = vec![1, 2, 3, 4, 5];
        engine
            .submit_encrypted_transaction(
                "tx1".to_string(),
                encrypted_data.clone(),
                "key_hash".to_string(),
            )
            .unwrap();

        assert_eq!(engine.pool_size(), 1);

        let decrypted = engine.decrypt_and_execute("tx1", b"key").unwrap();
        assert_eq!(decrypted, encrypted_data);
        assert_eq!(engine.pool_size(), 0);
    }

    #[test]
    fn test_mev_protection_coordinator() {
        let coordinator = MevProtectionCoordinator::with_all_strategies();

        assert!(coordinator.is_strategy_enabled(MitigationStrategy::FairOrdering));
        assert!(coordinator.is_strategy_enabled(MitigationStrategy::MevRebate));
        assert!(coordinator.is_strategy_enabled(MitigationStrategy::PrivateExecution));
    }

    #[test]
    fn test_rebate_completion() {
        let mut manager = MevRebateManager::with_defaults();

        let rebate = manager
            .create_rebate(
                "victim".to_string(),
                "attacker".to_string(),
                Decimal::new(100, 2),
                vec!["tx1".to_string()],
            )
            .unwrap();

        manager.complete_rebate(&rebate.rebate_id).unwrap();

        let victim_rebates = manager.get_victim_rebates("victim");
        assert_eq!(victim_rebates[0].status, RebateStatus::Completed);
    }
}