path: root/inheritance/src/unlock.rs

use serde::{Deserialize, Serialize};
use crate::contract::{InheritanceContract, ContractStatus};

/// Progressive unlock tier
/// This is OnionCoin's INNOVATIVE feature: gradual unlock gives owner time to recover!
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UnlockTier {
    /// Tier number (1, 2, 3...)
    pub tier: u32,

    /// Days after grace period expiry
    pub days_after_expiry: i64,

    /// Percentage of total to unlock at this tier
    pub unlock_percentage: u8,

    /// Has this tier been unlocked?
    pub unlocked: bool,

    /// Timestamp when unlocked
    pub unlock_time: Option<i64>,
}

impl UnlockTier {
    pub fn new(tier: u32, days_after_expiry: i64, unlock_percentage: u8) -> Self {
        Self {
            tier,
            days_after_expiry,
            unlock_percentage,
            unlocked: false,
            unlock_time: None,
        }
    }

    /// Check if this tier should be unlocked
    pub fn should_unlock(&self, time_since_expiry: i64) -> bool {
        !self.unlocked && time_since_expiry >= self.days_after_expiry * 24 * 3600
    }

    /// Mark tier as unlocked
    pub fn mark_unlocked(&mut self, current_time: i64) {
        self.unlocked = true;
        self.unlock_time = Some(current_time);
    }
}

/// Complete unlock schedule with multiple tiers
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UnlockSchedule {
    /// All unlock tiers
    pub tiers: Vec<UnlockTier>,

    /// Grace period expiry time
    pub grace_expiry_time: i64,

    /// Total unlocked so far (percentage)
    pub total_unlocked_percentage: u8,
}

impl UnlockSchedule {
    /// Create standard 4-tier progressive unlock
    /// Tier 1 (30 days): 10%
    /// Tier 2 (60 days): 25% (35% total)
    /// Tier 3 (90 days): 35% (70% total)
    /// Tier 4 (120 days): 30% (100% total)
    pub fn standard(grace_expiry_time: i64) -> Self {
        let tiers = vec![
            UnlockTier::new(1, 30, 10),  // 10% after 30 days
            UnlockTier::new(2, 60, 25),  // 25% after 60 days (35% total)
            UnlockTier::new(3, 90, 35),  // 35% after 90 days (70% total)
            UnlockTier::new(4, 120, 30), // 30% after 120 days (100% total)
        ];

        Self {
            tiers,
            grace_expiry_time,
            total_unlocked_percentage: 0,
        }
    }

    /// Create aggressive 3-tier unlock (faster)
    /// Tier 1 (14 days): 25%
    /// Tier 2 (30 days): 35% (60% total)
    /// Tier 3 (60 days): 40% (100% total)
    pub fn aggressive(grace_expiry_time: i64) -> Self {
        let tiers = vec![
            UnlockTier::new(1, 14, 25),  // 25% after 14 days
            UnlockTier::new(2, 30, 35),  // 35% after 30 days
            UnlockTier::new(3, 60, 40),  // 40% after 60 days
        ];

        Self {
            tiers,
            grace_expiry_time,
            total_unlocked_percentage: 0,
        }
    }

    /// Create conservative 5-tier unlock (slower, more chances to recover)
    /// Tier 1 (60 days): 5%
    /// Tier 2 (90 days): 10% (15% total)
    /// Tier 3 (120 days): 20% (35% total)
    /// Tier 4 (180 days): 30% (65% total)
    /// Tier 5 (365 days): 35% (100% total)
    pub fn conservative(grace_expiry_time: i64) -> Self {
        let tiers = vec![
            UnlockTier::new(1, 60, 5),    // 5% after 60 days
            UnlockTier::new(2, 90, 10),   // 10% after 90 days
            UnlockTier::new(3, 120, 20),  // 20% after 120 days
            UnlockTier::new(4, 180, 30),  // 30% after 180 days
            UnlockTier::new(5, 365, 35),  // 35% after 365 days
        ];

        Self {
            tiers,
            grace_expiry_time,
            total_unlocked_percentage: 0,
        }
    }

    /// Custom unlock schedule
    pub fn custom(grace_expiry_time: i64, tiers: Vec<UnlockTier>) -> Result<Self, UnlockError> {
        // Validate total percentage = 100%
        let total: u32 = tiers.iter().map(|t| t.unlock_percentage as u32).sum();
        if total != 100 {
            return Err(UnlockError::InvalidTotalPercentage(total));
        }

        // Validate tiers are in order
        for i in 1..tiers.len() {
            if tiers[i].days_after_expiry <= tiers[i - 1].days_after_expiry {
                return Err(UnlockError::InvalidTierOrder);
            }
        }

        Ok(Self {
            tiers,
            grace_expiry_time,
            total_unlocked_percentage: 0,
        })
    }

    /// Process unlocks for current time
    pub fn process_unlocks(&mut self, current_time: i64) -> Vec<UnlockEvent> {
        let time_since_expiry = current_time - self.grace_expiry_time;
        let mut events = Vec::new();

        for tier in &mut self.tiers {
            if tier.should_unlock(time_since_expiry) {
                tier.mark_unlocked(current_time);
                self.total_unlocked_percentage += tier.unlock_percentage;

                events.push(UnlockEvent {
                    tier: tier.tier,
                    percentage: tier.unlock_percentage,
                    timestamp: current_time,
                });
            }
        }

        events
    }

    /// Get next unlock tier
    pub fn next_unlock(&self) -> Option<&UnlockTier> {
        self.tiers.iter().find(|t| !t.unlocked)
    }

    /// Get time until next unlock
    pub fn time_until_next_unlock(&self, current_time: i64) -> Option<i64> {
        self.next_unlock().map(|tier| {
            let target_time = self.grace_expiry_time + tier.days_after_expiry * 24 * 3600;
            (target_time - current_time).max(0)
        })
    }

    /// Check if fully unlocked
    pub fn is_fully_unlocked(&self) -> bool {
        self.total_unlocked_percentage >= 100
    }

    /// Get summary of unlock progress
    pub fn progress_summary(&self, current_time: i64) -> String {
        let mut summary = format!("Unlocked: {}%\n", self.total_unlocked_percentage);

        for tier in &self.tiers {
            let status = if tier.unlocked {
                format!("✓ Unlocked at {}",
                    chrono::DateTime::from_timestamp(tier.unlock_time.unwrap(), 0)
                        .map(|dt| dt.format("%Y-%m-%d").to_string())
                        .unwrap_or_else(|| "unknown".to_string()))
            } else {
                let time_until = self.grace_expiry_time + tier.days_after_expiry * 24 * 3600 - current_time;
                if time_until > 0 {
                    format!("⏳ In {} days", time_until / (24 * 3600))
                } else {
                    "⏳ Ready to unlock".to_string()
                }
            };

            summary.push_str(&format!(
                "Tier {}: {}% - {}\n",
                tier.tier, tier.unlock_percentage, status
            ));
        }

        summary
    }
}

/// Event when a tier is unlocked
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UnlockEvent {
    pub tier: u32,
    pub percentage: u8,
    pub timestamp: i64,
}

/// Manager for progressive unlock
pub struct ProgressiveUnlock;

impl ProgressiveUnlock {
    /// Initialize progressive unlock for a contract
    pub fn initialize(
        contract: &mut InheritanceContract,
        schedule_type: UnlockScheduleType,
    ) -> Result<(), UnlockError> {
        if !contract.config.progressive_unlock {
            return Err(UnlockError::ProgressiveUnlockDisabled);
        }

        if contract.status != ContractStatus::Unlocking {
            return Err(UnlockError::ContractNotUnlocking);
        }

        let grace_expiry = contract.last_heartbeat + contract.config.total_timeout();

        let schedule = match schedule_type {
            UnlockScheduleType::Standard => UnlockSchedule::standard(grace_expiry),
            UnlockScheduleType::Aggressive => UnlockSchedule::aggressive(grace_expiry),
            UnlockScheduleType::Conservative => UnlockSchedule::conservative(grace_expiry),
            UnlockScheduleType::Custom(tiers) => {
                UnlockSchedule::custom(grace_expiry, tiers)?
            }
        };

        contract.unlock_schedule = Some(schedule);
        Ok(())
    }

    /// Process unlocks and distribute to beneficiaries
    pub fn process(
        contract: &mut InheritanceContract,
        current_time: i64,
    ) -> Result<Vec<Distribution>, UnlockError> {
        let schedule = contract
            .unlock_schedule
            .as_mut()
            .ok_or(UnlockError::NoSchedule)?;

        let events = schedule.process_unlocks(current_time);

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

        let mut distributions = Vec::new();

        for event in events {
            // Calculate amount to unlock for this tier
            let tier_amount = (contract.locked_amount as f64 * (event.percentage as f64 / 100.0)) as u64;

            // Distribute to beneficiaries
            for beneficiary in &mut contract.beneficiaries {
                let beneficiary_amount = beneficiary.calculate_amount(tier_amount);
                beneficiary.unlocked_amount += beneficiary_amount;

                distributions.push(Distribution {
                    beneficiary: beneficiary.pubkey,
                    amount: beneficiary_amount,
                    tier: event.tier,
                    timestamp: event.timestamp,
                });
            }
        }

        // Mark as completed if fully unlocked
        if schedule.is_fully_unlocked() {
            contract.status = ContractStatus::Completed;
        }

        Ok(distributions)
    }

    /// Get unlock status
    pub fn status(contract: &InheritanceContract, current_time: i64) -> String {
        if let Some(schedule) = &contract.unlock_schedule {
            schedule.progress_summary(current_time)
        } else {
            "No unlock schedule".to_string()
        }
    }
}

/// Type of unlock schedule
#[derive(Debug, Clone)]
pub enum UnlockScheduleType {
    Standard,
    Aggressive,
    Conservative,
    Custom(Vec<UnlockTier>),
}

/// Distribution to a beneficiary
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Distribution {
    pub beneficiary: [u8; 32],
    pub amount: u64,
    pub tier: u32,
    pub timestamp: i64,
}

#[derive(Debug, thiserror::Error)]
pub enum UnlockError {
    #[error("Invalid total percentage: {0}% (must be 100%)")]
    InvalidTotalPercentage(u32),

    #[error("Unlock tiers must be in chronological order")]
    InvalidTierOrder,

    #[error("Progressive unlock is disabled for this contract")]
    ProgressiveUnlockDisabled,

    #[error("Contract is not in unlocking state")]
    ContractNotUnlocking,

    #[error("No unlock schedule configured")]
    NoSchedule,
}

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

    #[test]
    fn test_unlock_tier() {
        let mut tier = UnlockTier::new(1, 30, 10);

        assert!(!tier.unlocked);
        assert!(tier.should_unlock(31 * 24 * 3600));
        assert!(!tier.should_unlock(29 * 24 * 3600));

        tier.mark_unlocked(1000);
        assert!(tier.unlocked);
        assert_eq!(tier.unlock_time, Some(1000));
    }

    #[test]
    fn test_standard_schedule() {
        let schedule = UnlockSchedule::standard(0);

        assert_eq!(schedule.tiers.len(), 4);
        assert_eq!(schedule.total_unlocked_percentage, 0);

        let total: u32 = schedule.tiers.iter().map(|t| t.unlock_percentage as u32).sum();
        assert_eq!(total, 100);
    }

    #[test]
    fn test_schedule_processing() {
        let mut schedule = UnlockSchedule::standard(0);

        // 31 days after expiry - should unlock tier 1 (10%)
        let events = schedule.process_unlocks(31 * 24 * 3600);

        assert_eq!(events.len(), 1);
        assert_eq!(events[0].tier, 1);
        assert_eq!(events[0].percentage, 10);
        assert_eq!(schedule.total_unlocked_percentage, 10);
    }

    #[test]
    fn test_multiple_tier_unlock() {
        let mut schedule = UnlockSchedule::standard(0);

        // 91 days after expiry - should unlock tiers 1, 2, and 3
        let events = schedule.process_unlocks(91 * 24 * 3600);

        assert_eq!(events.len(), 3);
        assert_eq!(schedule.total_unlocked_percentage, 70); // 10 + 25 + 35
    }

    #[test]
    fn test_full_unlock() {
        let mut schedule = UnlockSchedule::standard(0);

        schedule.process_unlocks(121 * 24 * 3600);

        assert!(schedule.is_fully_unlocked());
        assert_eq!(schedule.total_unlocked_percentage, 100);
    }

    #[test]
    fn test_custom_schedule() {
        let tiers = vec![
            UnlockTier::new(1, 10, 50),
            UnlockTier::new(2, 20, 50),
        ];

        let schedule = UnlockSchedule::custom(0, tiers);
        assert!(schedule.is_ok());
    }

    #[test]
    fn test_invalid_custom_schedule() {
        // Invalid total percentage
        let tiers = vec![
            UnlockTier::new(1, 10, 50),
            UnlockTier::new(2, 20, 40), // Total = 90%
        ];

        let schedule = UnlockSchedule::custom(0, tiers);
        assert!(matches!(
            schedule.unwrap_err(),
            UnlockError::InvalidTotalPercentage(90)
        ));
    }

    #[test]
    fn test_progressive_unlock_with_contract() {
        use crate::contract::Beneficiary;

        let beneficiaries = vec![
            Beneficiary::new([1u8; 32], 60).unwrap(),
            Beneficiary::new([2u8; 32], 40).unwrap(),
        ];

        let mut contract = InheritanceContract::new(
            [0u8; 32],
            beneficiaries,
            InheritanceConfig::default(),
            10000,
            0,
        )
        .unwrap();

        // Simulate unlocking state
        contract.status = ContractStatus::Unlocking;

        // Initialize progressive unlock
        ProgressiveUnlock::initialize(&mut contract, UnlockScheduleType::Standard).unwrap();

        assert!(contract.unlock_schedule.is_some());

        // Process unlocks after 31 days
        let grace_expiry = contract.config.total_timeout();
        let current_time = grace_expiry + 31 * 24 * 3600;

        let distributions = ProgressiveUnlock::process(&mut contract, current_time).unwrap();

        // Tier 1 unlocks 10% = 1000 ONC
        // Beneficiary 1: 60% of 1000 = 600
        // Beneficiary 2: 40% of 1000 = 400
        assert_eq!(distributions.len(), 2);
        assert_eq!(distributions[0].amount, 600);
        assert_eq!(distributions[1].amount, 400);
    }
}