kaccy-core 0.2.0

//! Protocol protection and exploit detection system
//!
//! This module provides automatic pause triggers, exploit detection,
//! recovery mechanisms, and compensation distribution for the protocol.

use crate::error::{CoreError, Result};
use chrono::{DateTime, Utc};
use rust_decimal::prelude::*;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use uuid::Uuid;

/// Type of security event
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
pub enum SecurityEventType {
    /// Suspicious price movement detected
    SuspiciousPriceMovement,
    /// Abnormal trading volume
    AbnormalVolume,
    /// Potential flash loan attack
    FlashLoanAttack,
    /// Oracle manipulation detected
    OracleManipulation,
    /// Reentrancy attack attempt
    ReentrancyAttempt,
    /// Unauthorized access attempt
    UnauthorizedAccess,
    /// Excessive slippage
    ExcessiveSlippage,
    /// Rapid successive transactions
    RapidTransactions,
}

/// Severity level of a security event
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub enum SecuritySeverity {
    /// Informational; monitor but no action required
    Low,
    /// Elevated concern; investigate when convenient
    Medium,
    /// Likely attack; prompt response needed
    High,
    /// Active exploit; immediate response required
    Critical,
}

/// Security event record
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SecurityEvent {
    /// Unique identifier for this security event
    pub event_id: Uuid,
    /// Category of the security event
    pub event_type: SecurityEventType,
    /// How severe this event is considered
    pub severity: SecuritySeverity,
    /// Human-readable description of the event
    pub description: String,
    /// Users affected by this event
    pub affected_users: Vec<Uuid>,
    /// Tokens affected by this event
    pub affected_tokens: Vec<Uuid>,
    /// When this event was detected
    pub detected_at: DateTime<Utc>,
    /// When this event was resolved, if applicable
    pub resolved_at: Option<DateTime<Utc>>,
    /// Whether this event has been resolved
    pub is_resolved: bool,
}

impl SecurityEvent {
    /// Create a new unresolved security event
    pub fn new(
        event_type: SecurityEventType,
        severity: SecuritySeverity,
        description: String,
        affected_users: Vec<Uuid>,
        affected_tokens: Vec<Uuid>,
    ) -> Self {
        Self {
            event_id: Uuid::new_v4(),
            event_type,
            severity,
            description,
            affected_users,
            affected_tokens,
            detected_at: Utc::now(),
            resolved_at: None,
            is_resolved: false,
        }
    }

    /// Mark event as resolved
    pub fn resolve(&mut self) {
        self.is_resolved = true;
        self.resolved_at = Some(Utc::now());
    }
}

/// Protocol pause state
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum PauseState {
    /// Protocol is fully operational
    Active,
    /// Specific features paused
    PartialPause,
    /// Protocol fully paused
    FullPause,
}

/// Recovery action to be taken after a security event
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RecoveryAction {
    /// Pause all protocol operations
    PauseProtocol,
    /// Pause trading for a specific token
    PauseToken(Uuid),
    /// Revert a set of recent transactions
    RevertTransactions {
        /// IDs of the transactions to revert
        transaction_ids: Vec<Uuid>,
    },
    /// Send compensation payments to affected users
    CompensateUsers {
        /// Map from user ID to compensation amount in BTC
        user_amounts: HashMap<Uuid, Decimal>,
    },
    /// Withdraw funds to safety in an emergency
    EmergencyWithdrawal {
        /// Token to withdraw
        token_id: Uuid,
        /// Amount to withdraw
        amount: Decimal,
    },
}

/// Compensation claim for affected users
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompensationClaim {
    /// Unique identifier for this compensation claim
    pub claim_id: Uuid,
    /// User who filed this claim
    pub user_id: Uuid,
    /// Security event that caused the loss
    pub security_event_id: Uuid,
    /// Amount of compensation claimed
    pub amount: Decimal,
    /// Description of the loss and basis for the claim
    pub description: String,
    /// Whether the compensation has been paid out
    pub is_paid: bool,
    /// When this claim was created
    pub created_at: DateTime<Utc>,
    /// When the compensation was paid out
    pub paid_at: Option<DateTime<Utc>>,
}

impl CompensationClaim {
    /// Create a new unpaid compensation claim
    pub fn new(
        user_id: Uuid,
        security_event_id: Uuid,
        amount: Decimal,
        description: String,
    ) -> Result<Self> {
        if amount <= Decimal::ZERO {
            return Err(CoreError::Validation(
                "Compensation amount must be positive".to_string(),
            ));
        }

        Ok(Self {
            claim_id: Uuid::new_v4(),
            user_id,
            security_event_id,
            amount,
            description,
            is_paid: false,
            created_at: Utc::now(),
            paid_at: None,
        })
    }

    /// Mark compensation as paid
    pub fn mark_paid(&mut self) {
        self.is_paid = true;
        self.paid_at = Some(Utc::now());
    }
}

/// Protocol protection configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProtectionConfig {
    /// Maximum price change per hour before triggering alert (e.g., 0.2 for 20%)
    pub max_price_change_per_hour: Decimal,
    /// Maximum volume change before triggering alert (multiplier, e.g., 5.0 for 5x)
    pub max_volume_multiplier: Decimal,
    /// Minimum time between transactions for rapid trading detection (seconds)
    pub min_transaction_interval: u64,
    /// Number of rapid transactions to trigger alert
    pub rapid_transaction_threshold: usize,
    /// Enable automatic pause on critical events
    pub auto_pause_enabled: bool,
    /// Compensation pool reserve percentage
    pub compensation_pool_reserve: Decimal,
}

impl Default for ProtectionConfig {
    fn default() -> Self {
        Self {
            max_price_change_per_hour: Decimal::from_str("0.20").unwrap(), // 20%
            max_volume_multiplier: Decimal::from_str("5.0").unwrap(),      // 5x normal
            min_transaction_interval: 5,                                   // 5 seconds
            rapid_transaction_threshold: 10,
            auto_pause_enabled: true,
            compensation_pool_reserve: Decimal::from_str("0.05").unwrap(), // 5%
        }
    }
}

/// Protocol protection manager
pub struct ProtocolProtection {
    /// Protection configuration (thresholds, auto-pause settings)
    config: ProtectionConfig,
    /// Current operational state of the protocol
    pause_state: PauseState,
    /// Active and historical security events indexed by event ID
    security_events: HashMap<Uuid, SecurityEvent>,
    /// Compensation claims indexed by claim ID
    compensation_claims: HashMap<Uuid, CompensationClaim>,
    /// BTC balance reserved for paying compensation claims
    compensation_pool_balance: Decimal,
    /// Token ID -> Price history (timestamp, price)
    price_history: HashMap<Uuid, Vec<(DateTime<Utc>, Decimal)>>,
    /// User ID -> Transaction timestamps
    user_transaction_history: HashMap<Uuid, Vec<DateTime<Utc>>>,
}

impl ProtocolProtection {
    /// Create a new protocol protection manager with the given configuration
    pub fn new(config: ProtectionConfig) -> Self {
        Self {
            config,
            pause_state: PauseState::Active,
            security_events: HashMap::new(),
            compensation_claims: HashMap::new(),
            compensation_pool_balance: Decimal::ZERO,
            price_history: HashMap::new(),
            user_transaction_history: HashMap::new(),
        }
    }

    /// Check for suspicious price movement
    pub fn check_price_movement(
        &mut self,
        token_id: Uuid,
        new_price: Decimal,
    ) -> Result<Option<SecurityEvent>> {
        let now = Utc::now();
        let one_hour_ago = now - chrono::Duration::hours(1);

        // Add new price to history
        self.price_history
            .entry(token_id)
            .or_default()
            .push((now, new_price));

        // Clean up old history (keep last 24 hours)
        let twenty_four_hours_ago = now - chrono::Duration::hours(24);
        if let Some(history) = self.price_history.get_mut(&token_id) {
            history.retain(|(timestamp, _)| *timestamp > twenty_four_hours_ago);

            // Check for rapid price changes in last hour
            let recent_prices: Vec<Decimal> = history
                .iter()
                .filter(|(timestamp, _)| *timestamp > one_hour_ago)
                .map(|(_, price)| *price)
                .collect();

            if recent_prices.len() >= 2 {
                let min_price = recent_prices.iter().min().copied().unwrap();
                let max_price = recent_prices.iter().max().copied().unwrap();

                if min_price > Decimal::ZERO {
                    let price_change = (max_price - min_price) / min_price;

                    if price_change > self.config.max_price_change_per_hour {
                        let event = self.create_security_event(
                            SecurityEventType::SuspiciousPriceMovement,
                            SecuritySeverity::High,
                            format!(
                                "Price changed by {:.2}% in one hour",
                                price_change * Decimal::from(100)
                            ),
                            vec![],
                            vec![token_id],
                        );

                        return Ok(Some(event));
                    }
                }
            }
        }

        Ok(None)
    }

    /// Check for rapid trading (potential bot attack)
    pub fn check_rapid_trading(&mut self, user_id: Uuid) -> Result<Option<SecurityEvent>> {
        let now = Utc::now();

        // Add transaction timestamp
        self.user_transaction_history
            .entry(user_id)
            .or_default()
            .push(now);

        // Clean up old history (keep last hour)
        let one_hour_ago = now - chrono::Duration::hours(1);
        if let Some(history) = self.user_transaction_history.get_mut(&user_id) {
            history.retain(|timestamp| *timestamp > one_hour_ago);

            // Check for rapid transactions
            if history.len() >= self.config.rapid_transaction_threshold {
                // Check if transactions are suspiciously close together
                let mut rapid_count = 0;
                for window in history.windows(2) {
                    if let [t1, t2] = window {
                        let diff = (*t2 - *t1).num_seconds();
                        if diff < self.config.min_transaction_interval as i64 {
                            rapid_count += 1;
                        }
                    }
                }

                if rapid_count >= self.config.rapid_transaction_threshold / 2 {
                    let event = self.create_security_event(
                        SecurityEventType::RapidTransactions,
                        SecuritySeverity::Medium,
                        format!("User performed {} rapid transactions", rapid_count),
                        vec![user_id],
                        vec![],
                    );

                    return Ok(Some(event));
                }
            }
        }

        Ok(None)
    }

    /// Create and record a security event
    fn create_security_event(
        &mut self,
        event_type: SecurityEventType,
        severity: SecuritySeverity,
        description: String,
        affected_users: Vec<Uuid>,
        affected_tokens: Vec<Uuid>,
    ) -> SecurityEvent {
        let event = SecurityEvent::new(
            event_type,
            severity,
            description,
            affected_users,
            affected_tokens,
        );

        // Auto-pause if configured and event is critical
        if self.config.auto_pause_enabled && severity == SecuritySeverity::Critical {
            self.pause_protocol();
        }

        let event_id = event.event_id;
        self.security_events.insert(event_id, event.clone());
        event
    }

    /// Report a security event
    pub fn report_security_event(
        &mut self,
        event_type: SecurityEventType,
        severity: SecuritySeverity,
        description: String,
        affected_users: Vec<Uuid>,
        affected_tokens: Vec<Uuid>,
    ) -> Uuid {
        let event = self.create_security_event(
            event_type,
            severity,
            description,
            affected_users,
            affected_tokens,
        );
        event.event_id
    }

    /// Pause protocol operations
    pub fn pause_protocol(&mut self) {
        self.pause_state = PauseState::FullPause;
    }

    /// Resume protocol operations
    pub fn resume_protocol(&mut self) {
        self.pause_state = PauseState::Active;
    }

    /// Check if protocol is paused
    pub fn is_paused(&self) -> bool {
        matches!(
            self.pause_state,
            PauseState::FullPause | PauseState::PartialPause
        )
    }

    /// Get current pause state
    pub fn get_pause_state(&self) -> PauseState {
        self.pause_state
    }

    /// Create compensation claim for affected user
    pub fn create_compensation_claim(
        &mut self,
        user_id: Uuid,
        security_event_id: Uuid,
        amount: Decimal,
        description: String,
    ) -> Result<Uuid> {
        // Verify security event exists
        if !self.security_events.contains_key(&security_event_id) {
            return Err(CoreError::NotFound("Security event not found".to_string()));
        }

        let claim = CompensationClaim::new(user_id, security_event_id, amount, description)?;
        let claim_id = claim.claim_id;

        self.compensation_claims.insert(claim_id, claim);
        Ok(claim_id)
    }

    /// Process compensation payout
    pub fn payout_compensation(&mut self, claim_id: Uuid) -> Result<Decimal> {
        let claim = self
            .compensation_claims
            .get_mut(&claim_id)
            .ok_or_else(|| CoreError::NotFound("Compensation claim not found".to_string()))?;

        if claim.is_paid {
            return Err(CoreError::InvalidState(
                "Compensation already paid".to_string(),
            ));
        }

        if claim.amount > self.compensation_pool_balance {
            return Err(CoreError::InsufficientBalance {
                required: claim.amount,
                available: self.compensation_pool_balance,
            });
        }

        claim.mark_paid();
        self.compensation_pool_balance -= claim.amount;

        Ok(claim.amount)
    }

    /// Add funds to compensation pool
    pub fn fund_compensation_pool(&mut self, amount: Decimal) -> Result<()> {
        if amount <= Decimal::ZERO {
            return Err(CoreError::Validation("Amount must be positive".to_string()));
        }

        self.compensation_pool_balance += amount;
        Ok(())
    }

    /// Get compensation pool balance
    pub fn get_compensation_pool_balance(&self) -> Decimal {
        self.compensation_pool_balance
    }

    /// Get security event by ID
    pub fn get_security_event(&self, event_id: Uuid) -> Option<&SecurityEvent> {
        self.security_events.get(&event_id)
    }

    /// Get all unresolved security events
    pub fn get_unresolved_events(&self) -> Vec<&SecurityEvent> {
        self.security_events
            .values()
            .filter(|e| !e.is_resolved)
            .collect()
    }

    /// Get security events by severity
    pub fn get_events_by_severity(&self, severity: SecuritySeverity) -> Vec<&SecurityEvent> {
        self.security_events
            .values()
            .filter(|e| e.severity == severity)
            .collect()
    }

    /// Resolve a security event
    pub fn resolve_security_event(&mut self, event_id: Uuid) -> Result<()> {
        let event = self
            .security_events
            .get_mut(&event_id)
            .ok_or_else(|| CoreError::NotFound("Security event not found".to_string()))?;

        event.resolve();
        Ok(())
    }

    /// Get compensation claim
    pub fn get_compensation_claim(&self, claim_id: Uuid) -> Option<&CompensationClaim> {
        self.compensation_claims.get(&claim_id)
    }

    /// Get all compensation claims for a user
    pub fn get_user_compensation_claims(&self, user_id: Uuid) -> Vec<&CompensationClaim> {
        self.compensation_claims
            .values()
            .filter(|c| c.user_id == user_id)
            .collect()
    }

    /// Get protocol statistics
    pub fn get_statistics(&self) -> ProtectionStatistics {
        let total_events = self.security_events.len();
        let resolved_events = self
            .security_events
            .values()
            .filter(|e| e.is_resolved)
            .count();
        let critical_events = self
            .security_events
            .values()
            .filter(|e| e.severity == SecuritySeverity::Critical)
            .count();

        let total_compensation_paid = self
            .compensation_claims
            .values()
            .filter(|c| c.is_paid)
            .map(|c| c.amount)
            .sum();

        ProtectionStatistics {
            total_security_events: total_events,
            resolved_events,
            critical_events,
            total_compensation_claims: self.compensation_claims.len(),
            total_compensation_paid,
            compensation_pool_balance: self.compensation_pool_balance,
            pause_state: self.pause_state,
        }
    }
}

/// Protocol protection statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProtectionStatistics {
    /// Total number of security events recorded
    pub total_security_events: usize,
    /// Number of security events that have been resolved
    pub resolved_events: usize,
    /// Number of security events with Critical severity
    pub critical_events: usize,
    /// Total number of compensation claims filed
    pub total_compensation_claims: usize,
    /// Total BTC paid out in compensation
    pub total_compensation_paid: Decimal,
    /// Current BTC balance in the compensation pool
    pub compensation_pool_balance: Decimal,
    /// Current operational state of the protocol
    pub pause_state: PauseState,
}

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

    #[test]
    fn test_security_event_creation() {
        let event = SecurityEvent::new(
            SecurityEventType::SuspiciousPriceMovement,
            SecuritySeverity::High,
            "Price spike detected".to_string(),
            vec![],
            vec![Uuid::new_v4()],
        );

        assert_eq!(event.severity, SecuritySeverity::High);
        assert!(!event.is_resolved);
    }

    #[test]
    fn test_security_event_resolution() {
        let mut event = SecurityEvent::new(
            SecurityEventType::SuspiciousPriceMovement,
            SecuritySeverity::High,
            "Test".to_string(),
            vec![],
            vec![],
        );

        event.resolve();
        assert!(event.is_resolved);
        assert!(event.resolved_at.is_some());
    }

    #[test]
    fn test_compensation_claim_creation() {
        let user_id = Uuid::new_v4();
        let event_id = Uuid::new_v4();

        let claim = CompensationClaim::new(
            user_id,
            event_id,
            dec!(1000),
            "Compensation for exploit".to_string(),
        );

        assert!(claim.is_ok());
        let claim = claim.unwrap();
        assert!(!claim.is_paid);
        assert_eq!(claim.amount, dec!(1000));
    }

    #[test]
    fn test_protocol_pause() {
        let config = ProtectionConfig::default();
        let mut protection = ProtocolProtection::new(config);

        assert!(!protection.is_paused());

        protection.pause_protocol();
        assert!(protection.is_paused());

        protection.resume_protocol();
        assert!(!protection.is_paused());
    }

    #[test]
    fn test_price_movement_detection() {
        let config = ProtectionConfig::default();
        let mut protection = ProtocolProtection::new(config);
        let token_id = Uuid::new_v4();

        // Add initial price
        let result = protection.check_price_movement(token_id, dec!(100));
        assert!(result.is_ok());
        assert!(result.unwrap().is_none());

        // Add price with large change (should trigger alert)
        let result = protection.check_price_movement(token_id, dec!(150));
        assert!(result.is_ok());
        // First detection might not trigger due to time window
    }

    #[test]
    fn test_rapid_trading_detection() {
        let config = ProtectionConfig::default();
        let threshold = config.rapid_transaction_threshold;
        let mut protection = ProtocolProtection::new(config);
        let user_id = Uuid::new_v4();

        // Simulate rapid transactions
        for _ in 0..threshold {
            let _ = protection.check_rapid_trading(user_id);
        }

        let stats = protection.get_statistics();
        // Should have detected rapid trading (stats should be valid)
        assert!(stats.total_security_events < 100); // Reasonable upper bound for test
    }

    #[test]
    fn test_compensation_workflow() {
        let config = ProtectionConfig::default();
        let mut protection = ProtocolProtection::new(config);

        // Fund compensation pool
        protection.fund_compensation_pool(dec!(10000)).unwrap();
        assert_eq!(protection.get_compensation_pool_balance(), dec!(10000));

        // Create security event
        let event_id = protection.report_security_event(
            SecurityEventType::FlashLoanAttack,
            SecuritySeverity::Critical,
            "Exploit detected".to_string(),
            vec![Uuid::new_v4()],
            vec![],
        );

        // Create compensation claim
        let user_id = Uuid::new_v4();
        let claim_id = protection
            .create_compensation_claim(user_id, event_id, dec!(1000), "Compensation".to_string())
            .unwrap();

        // Payout compensation
        let payout = protection.payout_compensation(claim_id).unwrap();
        assert_eq!(payout, dec!(1000));
        assert_eq!(protection.get_compensation_pool_balance(), dec!(9000));
    }

    #[test]
    fn test_auto_pause_on_critical_event() {
        let config = ProtectionConfig {
            auto_pause_enabled: true,
            ..Default::default()
        };
        let mut protection = ProtocolProtection::new(config);

        // Report critical event
        protection.report_security_event(
            SecurityEventType::FlashLoanAttack,
            SecuritySeverity::Critical,
            "Critical attack detected".to_string(),
            vec![],
            vec![],
        );

        // Should be paused
        assert!(protection.is_paused());
    }
}