path: root/inheritance/src/heartbeat.rs

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use crate::contract::InheritanceContract;

/// Manages heartbeat system for inheritance contracts
/// Owner must periodically "check in" to prove they're alive
pub struct HeartbeatManager {
    /// Contract heartbeat tracking
    contracts: HashMap<[u8; 32], HeartbeatRecord>,

    /// Notification settings
    notification_settings: NotificationSettings,
}

/// Record of heartbeat activity for a contract
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HeartbeatRecord {
    /// Contract ID
    pub contract_id: [u8; 32],

    /// Last heartbeat timestamp
    pub last_heartbeat: i64,

    /// Total heartbeats sent
    pub heartbeat_count: u64,

    /// Average interval between heartbeats
    pub avg_interval_days: f64,

    /// Warnings sent to owner
    pub warnings_sent: u32,

    /// Notifications sent to beneficiaries
    pub beneficiary_notifications_sent: u32,
}

/// Notification settings for heartbeat warnings
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NotificationSettings {
    /// Send warning when X% of interval remaining
    pub warning_threshold_percentage: u8,

    /// Send critical warning when in grace period
    pub grace_period_warning: bool,

    /// Notify beneficiaries when grace period starts
    pub notify_beneficiaries_grace: bool,

    /// Notify beneficiaries on each unlock tier
    pub notify_beneficiaries_unlock: bool,
}

impl Default for NotificationSettings {
    fn default() -> Self {
        Self {
            warning_threshold_percentage: 80, // Warn at 80% of interval
            grace_period_warning: true,
            notify_beneficiaries_grace: true,
            notify_beneficiaries_unlock: true,
        }
    }
}

impl HeartbeatManager {
    pub fn new() -> Self {
        Self {
            contracts: HashMap::new(),
            notification_settings: NotificationSettings::default(),
        }
    }

    /// Register a contract for heartbeat monitoring
    pub fn register_contract(&mut self, contract_id: [u8; 32], current_time: i64) {
        let record = HeartbeatRecord {
            contract_id,
            last_heartbeat: current_time,
            heartbeat_count: 0,
            avg_interval_days: 0.0,
            warnings_sent: 0,
            beneficiary_notifications_sent: 0,
        };

        self.contracts.insert(contract_id, record);
    }

    /// Process a heartbeat
    pub fn process_heartbeat(
        &mut self,
        contract: &mut InheritanceContract,
        current_time: i64,
    ) -> Result<HeartbeatStatus, HeartbeatError> {
        // Record heartbeat in contract
        contract.heartbeat(current_time)
            .map_err(|_| HeartbeatError::ContractError)?;

        // Update heartbeat record
        if let Some(record) = self.contracts.get_mut(&contract.contract_id) {
            let interval = current_time - record.last_heartbeat;
            let interval_days = interval as f64 / (24.0 * 3600.0);

            // Update average interval
            let total = record.heartbeat_count as f64;
            record.avg_interval_days =
                (record.avg_interval_days * total + interval_days) / (total + 1.0);

            record.last_heartbeat = current_time;
            record.heartbeat_count += 1;
            record.warnings_sent = 0; // Reset warnings

            Ok(HeartbeatStatus::Success {
                next_required: current_time + contract.config.heartbeat_interval,
                avg_interval_days: record.avg_interval_days,
            })
        } else {
            Err(HeartbeatError::ContractNotRegistered)
        }
    }

    /// Check all contracts and generate notifications
    pub fn check_contracts(
        &mut self,
        contracts: &[&InheritanceContract],
        current_time: i64,
    ) -> Vec<Notification> {
        let mut notifications = Vec::new();

        for contract in contracts {
            if let Some(record) = self.contracts.get_mut(&contract.contract_id) {
                let time_since_heartbeat = current_time - contract.last_heartbeat;
                let interval = contract.config.heartbeat_interval;

                // Calculate percentage of interval elapsed
                let percentage_elapsed = (time_since_heartbeat as f64 / interval as f64 * 100.0) as u8;

                // Check for warnings
                if percentage_elapsed >= self.notification_settings.warning_threshold_percentage
                    && record.warnings_sent == 0
                {
                    notifications.push(Notification {
                        contract_id: contract.contract_id,
                        notification_type: NotificationType::WarningOwner,
                        message: format!(
                            "{}% of heartbeat interval elapsed. Please send heartbeat soon!",
                            percentage_elapsed
                        ),
                        urgency: Urgency::Medium,
                        timestamp: current_time,
                    });
                    record.warnings_sent += 1;
                }

                // Check if grace period started
                if contract.is_timeout(current_time)
                    && !contract.is_grace_expired(current_time)
                    && self.notification_settings.grace_period_warning
                    && record.warnings_sent < 2
                {
                    let grace_remaining = contract.config.grace_period
                        - (time_since_heartbeat - interval);

                    notifications.push(Notification {
                        contract_id: contract.contract_id,
                        notification_type: NotificationType::CriticalOwner,
                        message: format!(
                            "CRITICAL: Grace period active! {} days remaining before unlock!",
                            grace_remaining / (24 * 3600)
                        ),
                        urgency: Urgency::Critical,
                        timestamp: current_time,
                    });

                    // Notify beneficiaries
                    if self.notification_settings.notify_beneficiaries_grace
                        && record.beneficiary_notifications_sent == 0
                    {
                        notifications.push(Notification {
                            contract_id: contract.contract_id,
                            notification_type: NotificationType::InfoBeneficiaries,
                            message: "Inheritance contract has entered grace period. Unlock will begin if owner does not respond.".to_string(),
                            urgency: Urgency::Low,
                            timestamp: current_time,
                        });
                        record.beneficiary_notifications_sent += 1;
                    }

                    record.warnings_sent += 1;
                }

                // Check if unlocking started
                if contract.is_grace_expired(current_time)
                    && record.beneficiary_notifications_sent < 2
                    && self.notification_settings.notify_beneficiaries_unlock
                {
                    notifications.push(Notification {
                        contract_id: contract.contract_id,
                        notification_type: NotificationType::UnlockStarted,
                        message: "Inheritance unlock has begun. Funds will be distributed according to the schedule.".to_string(),
                        urgency: Urgency::High,
                        timestamp: current_time,
                    });
                    record.beneficiary_notifications_sent += 1;
                }
            }
        }

        notifications
    }

    /// Get heartbeat statistics for a contract
    pub fn get_stats(&self, contract_id: &[u8; 32]) -> Option<&HeartbeatRecord> {
        self.contracts.get(contract_id)
    }

    /// Suggest optimal heartbeat frequency for user
    pub fn suggest_frequency(&self, contract_id: &[u8; 32]) -> Option<String> {
        self.contracts.get(contract_id).map(|record| {
            if record.heartbeat_count < 3 {
                "Not enough data yet. Continue sending regular heartbeats.".to_string()
            } else if record.avg_interval_days < 30.0 {
                format!(
                    "You're checking in every {:.0} days. Consider spacing out to 30-60 days.",
                    record.avg_interval_days
                )
            } else if record.avg_interval_days > 60.0 {
                format!(
                    "You're checking in every {:.0} days. Consider more frequent heartbeats (30-45 days) for safety.",
                    record.avg_interval_days
                )
            } else {
                format!(
                    "Good frequency! Checking in every {:.0} days.",
                    record.avg_interval_days
                )
            }
        })
    }

    /// Create simple heartbeat transaction
    /// Just send tiny amount to yourself to prove you're alive
    pub fn create_heartbeat_tx(owner: [u8; 32]) -> HeartbeatTransaction {
        HeartbeatTransaction {
            from: owner,
            to: owner, // Send to yourself
            amount: 1, // Minimum amount (0.00000001 ONC)
            timestamp: chrono::Utc::now().timestamp(),
            heartbeat_marker: true,
        }
    }
}

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

/// Result of heartbeat processing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum HeartbeatStatus {
    Success {
        next_required: i64,
        avg_interval_days: f64,
    },
    Warning {
        message: String,
    },
    Error {
        message: String,
    },
}

/// Notification to send to owner or beneficiaries
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Notification {
    pub contract_id: [u8; 32],
    pub notification_type: NotificationType,
    pub message: String,
    pub urgency: Urgency,
    pub timestamp: i64,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum NotificationType {
    WarningOwner,
    CriticalOwner,
    InfoBeneficiaries,
    UnlockStarted,
    TierUnlocked,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Urgency {
    Low,
    Medium,
    High,
    Critical,
}

/// Simple heartbeat transaction
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HeartbeatTransaction {
    pub from: [u8; 32],
    pub to: [u8; 32],
    pub amount: u64,
    pub timestamp: i64,
    pub heartbeat_marker: bool,
}

#[derive(Debug, thiserror::Error)]
pub enum HeartbeatError {
    #[error("Contract not registered with heartbeat manager")]
    ContractNotRegistered,

    #[error("Contract error occurred")]
    ContractError,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::contract::{InheritanceContract, Beneficiary, InheritanceConfig};

    fn create_test_contract() -> InheritanceContract {
        let beneficiaries = vec![Beneficiary::new([1u8; 32], 100).unwrap()];

        InheritanceContract::new(
            [0u8; 32],
            beneficiaries,
            InheritanceConfig::default(),
            10000,
            0,
        )
        .unwrap()
    }

    #[test]
    fn test_heartbeat_manager_registration() {
        let mut manager = HeartbeatManager::new();
        let contract_id = [1u8; 32];

        manager.register_contract(contract_id, 0);

        assert!(manager.contracts.contains_key(&contract_id));
    }

    #[test]
    fn test_heartbeat_processing() {
        let mut manager = HeartbeatManager::new();
        let mut contract = create_test_contract();

        manager.register_contract(contract.contract_id, 0);

        let result = manager.process_heartbeat(&mut contract, 1000);

        assert!(result.is_ok());
        assert_eq!(contract.last_heartbeat, 1000);

        let record = manager.get_stats(&contract.contract_id).unwrap();
        assert_eq!(record.heartbeat_count, 1);
    }

    #[test]
    fn test_average_interval_calculation() {
        let mut manager = HeartbeatManager::new();
        let mut contract = create_test_contract();

        manager.register_contract(contract.contract_id, 0);

        // First heartbeat at day 30
        manager.process_heartbeat(&mut contract, 30 * 24 * 3600).unwrap();

        // Second heartbeat at day 60
        manager.process_heartbeat(&mut contract, 60 * 24 * 3600).unwrap();

        let record = manager.get_stats(&contract.contract_id).unwrap();
        assert!((record.avg_interval_days - 30.0).abs() < 1.0);
    }

    #[test]
    fn test_notification_generation() {
        let mut manager = HeartbeatManager::new();
        let mut contract = create_test_contract();

        // Set short interval for testing
        contract.config.heartbeat_interval = 100;
        contract.config.grace_period = 50;

        manager.register_contract(contract.contract_id, 0);

        // Check at 85% of interval (should generate warning)
        let notifications = manager.check_contracts(&[&contract], 85);

        assert!(!notifications.is_empty());
        assert_eq!(notifications[0].notification_type, NotificationType::WarningOwner);
        assert_eq!(notifications[0].urgency, Urgency::Medium);
    }

    #[test]
    fn test_grace_period_notifications() {
        let mut manager = HeartbeatManager::new();
        let mut contract = create_test_contract();

        contract.config.heartbeat_interval = 100;
        contract.config.grace_period = 50;

        manager.register_contract(contract.contract_id, 0);

        // Enter grace period
        let notifications = manager.check_contracts(&[&contract], 101);

        // Should have critical warning for owner and info for beneficiaries
        assert!(notifications.len() >= 2);
        assert!(notifications.iter().any(|n| n.notification_type == NotificationType::CriticalOwner));
        assert!(notifications.iter().any(|n| n.notification_type == NotificationType::InfoBeneficiaries));
    }

    #[test]
    fn test_heartbeat_transaction() {
        let owner = [42u8; 32];
        let tx = HeartbeatManager::create_heartbeat_tx(owner);

        assert_eq!(tx.from, owner);
        assert_eq!(tx.to, owner); // Sent to self
        assert_eq!(tx.amount, 1); // Minimal amount
        assert!(tx.heartbeat_marker);
    }

    #[test]
    fn test_frequency_suggestion() {
        let mut manager = HeartbeatManager::new();
        let contract_id = [1u8; 32];

        manager.register_contract(contract_id, 0);

        // Not enough data initially
        let suggestion = manager.suggest_frequency(&contract_id).unwrap();
        assert!(suggestion.contains("Not enough data"));

        // Simulate multiple heartbeats
        let mut record = manager.contracts.get_mut(&contract_id).unwrap();
        record.heartbeat_count = 5;
        record.avg_interval_days = 45.0;

        let suggestion = manager.suggest_frequency(&contract_id).unwrap();
        assert!(suggestion.contains("Good frequency"));
    }
}