path: root/inheritance/src/contract.rs

use serde::{Deserialize, Serialize};
use crate::unlock::UnlockSchedule;
use blake3::Hasher;

/// Time-locked inheritance contract for OnionCoin
/// WORLD'S FIRST native blockchain inheritance system!
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InheritanceContract {
    /// Contract ID (hash of owner + creation time)
    pub contract_id: [u8; 32],

    /// Owner's public key
    pub owner: [u8; 32],

    /// Beneficiaries with their inheritance percentages
    pub beneficiaries: Vec<Beneficiary>,

    /// Configuration
    pub config: InheritanceConfig,

    /// Last heartbeat timestamp
    pub last_heartbeat: i64,

    /// Contract creation time
    pub created_at: i64,

    /// Current status
    pub status: ContractStatus,

    /// Unlock schedule (optional for progressive unlock)
    pub unlock_schedule: Option<UnlockSchedule>,

    /// Encrypted data (for privacy)
    pub encrypted_metadata: Option<Vec<u8>>,

    /// Dispute flag (owner can dispute if alive)
    pub dispute_active: bool,

    /// Total amount locked in contract
    pub locked_amount: u64,
}

/// A beneficiary who will receive inheritance
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Beneficiary {
    /// Beneficiary's public key
    pub pubkey: [u8; 32],

    /// Percentage of inheritance (0-100)
    pub percentage: u8,

    /// Optional encrypted name/info
    pub encrypted_info: Option<Vec<u8>>,

    /// Has this beneficiary been notified?
    pub notified: bool,

    /// Amount unlocked so far
    pub unlocked_amount: u64,
}

impl Beneficiary {
    pub fn new(pubkey: [u8; 32], percentage: u8) -> Result<Self, InheritanceError> {
        if percentage > 100 {
            return Err(InheritanceError::InvalidPercentage(percentage));
        }

        Ok(Self {
            pubkey,
            percentage,
            encrypted_info: None,
            notified: false,
            unlocked_amount: 0,
        })
    }

    /// Calculate absolute amount for this beneficiary
    pub fn calculate_amount(&self, total: u64) -> u64 {
        (total as f64 * (self.percentage as f64 / 100.0)) as u64
    }
}

/// Contract configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InheritanceConfig {
    /// Heartbeat interval (seconds) - how often owner must check in
    /// Default: 90 days
    pub heartbeat_interval: i64,

    /// Grace period after missed heartbeat (seconds)
    /// Default: 30 days
    pub grace_period: i64,

    /// Enable progressive unlock
    pub progressive_unlock: bool,

    /// Require Tor-only access for heartbeat
    pub tor_only: bool,

    /// Enable encrypted beneficiary list
    pub encrypted_beneficiaries: bool,

    /// Require multi-signature for unlock
    pub multi_sig_required: bool,

    /// Number of signatures required (if multi_sig enabled)
    pub required_signatures: u32,

    /// Enable Shamir secret sharing
    pub shamir_enabled: bool,

    /// Shamir threshold (M of N)
    pub shamir_threshold: Option<(u8, u8)>,
}

impl Default for InheritanceConfig {
    fn default() -> Self {
        Self {
            heartbeat_interval: 90 * 24 * 3600,   // 90 days
            grace_period: 30 * 24 * 3600,         // 30 days
            progressive_unlock: true,              // Default ON
            tor_only: true,                        // Privacy first!
            encrypted_beneficiaries: true,
            multi_sig_required: false,
            required_signatures: 0,
            shamir_enabled: false,
            shamir_threshold: None,
        }
    }
}

impl InheritanceConfig {
    /// Total time before unlock starts
    pub fn total_timeout(&self) -> i64 {
        self.heartbeat_interval + self.grace_period
    }

    /// Create a secure config (maximum protections)
    pub fn secure() -> Self {
        Self {
            heartbeat_interval: 180 * 24 * 3600,  // 6 months
            grace_period: 60 * 24 * 3600,         // 2 months
            progressive_unlock: true,
            tor_only: true,
            encrypted_beneficiaries: true,
            multi_sig_required: true,
            required_signatures: 2,
            shamir_enabled: true,
            shamir_threshold: Some((3, 5)),       // 3 of 5
        }
    }

    /// Create a quick config (for testing or urgent cases)
    pub fn quick() -> Self {
        Self {
            heartbeat_interval: 30 * 24 * 3600,   // 30 days
            grace_period: 7 * 24 * 3600,          // 7 days
            progressive_unlock: true,
            tor_only: false,
            encrypted_beneficiaries: false,
            multi_sig_required: false,
            required_signatures: 0,
            shamir_enabled: false,
            shamir_threshold: None,
        }
    }
}

/// Contract status
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ContractStatus {
    /// Contract is active, owner is alive
    Active,

    /// Heartbeat missed, grace period active
    GracePeriod,

    /// Grace period expired, unlocking in progress
    Unlocking,

    /// Fully unlocked, inheritance distributed
    Completed,

    /// Owner disputed (proved they're alive)
    Disputed,

    /// Contract cancelled by owner
    Cancelled,
}

impl InheritanceContract {
    /// Create a new inheritance contract
    pub fn new(
        owner: [u8; 32],
        beneficiaries: Vec<Beneficiary>,
        config: InheritanceConfig,
        locked_amount: u64,
        current_time: i64,
    ) -> Result<Self, InheritanceError> {
        // Validate beneficiaries
        Self::validate_beneficiaries(&beneficiaries)?;

        // Generate contract ID
        let contract_id = Self::generate_contract_id(&owner, current_time);

        Ok(Self {
            contract_id,
            owner,
            beneficiaries,
            config,
            last_heartbeat: current_time,
            created_at: current_time,
            status: ContractStatus::Active,
            unlock_schedule: None,
            encrypted_metadata: None,
            dispute_active: false,
            locked_amount,
        })
    }

    /// Generate unique contract ID
    fn generate_contract_id(owner: &[u8; 32], created_at: i64) -> [u8; 32] {
        let mut hasher = Hasher::new();
        hasher.update(owner);
        hasher.update(&created_at.to_le_bytes());
        *hasher.finalize().as_bytes()
    }

    /// Validate beneficiaries list
    fn validate_beneficiaries(beneficiaries: &[Beneficiary]) -> Result<(), InheritanceError> {
        if beneficiaries.is_empty() {
            return Err(InheritanceError::NoBeneficiaries);
        }

        // Check total percentage = 100%
        let total: u32 = beneficiaries.iter().map(|b| b.percentage as u32).sum();
        if total != 100 {
            return Err(InheritanceError::InvalidTotalPercentage(total));
        }

        // Check for duplicate beneficiaries
        let mut seen = std::collections::HashSet::new();
        for b in beneficiaries {
            if !seen.insert(b.pubkey) {
                return Err(InheritanceError::DuplicateBeneficiary);
            }
        }

        Ok(())
    }

    /// Record a heartbeat (owner is alive)
    pub fn heartbeat(&mut self, current_time: i64) -> Result<(), InheritanceError> {
        // Only owner can heartbeat
        if self.status == ContractStatus::Cancelled {
            return Err(InheritanceError::ContractCancelled);
        }

        if self.status == ContractStatus::Completed {
            return Err(InheritanceError::ContractCompleted);
        }

        self.last_heartbeat = current_time;

        // Reset status if was in grace period
        if self.status == ContractStatus::GracePeriod {
            self.status = ContractStatus::Active;
        }

        Ok(())
    }

    /// Check if heartbeat timeout has been reached
    pub fn is_timeout(&self, current_time: i64) -> bool {
        let elapsed = current_time - self.last_heartbeat;
        elapsed >= self.config.heartbeat_interval
    }

    /// Check if grace period has expired
    pub fn is_grace_expired(&self, current_time: i64) -> bool {
        let elapsed = current_time - self.last_heartbeat;
        elapsed >= self.config.total_timeout()
    }

    /// Update contract status based on time
    pub fn update_status(&mut self, current_time: i64) {
        if self.is_grace_expired(current_time) {
            if self.status != ContractStatus::Unlocking
                && self.status != ContractStatus::Completed
                && self.status != ContractStatus::Cancelled {
                self.status = ContractStatus::Unlocking;
            }
        } else if self.is_timeout(current_time) {
            if self.status == ContractStatus::Active {
                self.status = ContractStatus::GracePeriod;
            }
        }
    }

    /// File a dispute (owner proves they're alive)
    pub fn file_dispute(&mut self, current_time: i64) -> Result<(), InheritanceError> {
        if self.status == ContractStatus::Completed {
            return Err(InheritanceError::ContractCompleted);
        }

        self.dispute_active = true;
        self.status = ContractStatus::Disputed;
        self.last_heartbeat = current_time; // Reset heartbeat

        Ok(())
    }

    /// Cancel contract (owner decides to cancel)
    pub fn cancel(&mut self) -> Result<u64, InheritanceError> {
        if self.status == ContractStatus::Completed {
            return Err(InheritanceError::ContractCompleted);
        }

        self.status = ContractStatus::Cancelled;

        // Return locked amount to owner
        let amount = self.locked_amount;
        self.locked_amount = 0;

        Ok(amount)
    }

    /// Calculate time until unlock
    pub fn time_until_unlock(&self, current_time: i64) -> i64 {
        let total_timeout = self.config.total_timeout();
        let elapsed = current_time - self.last_heartbeat;
        (total_timeout - elapsed).max(0)
    }

    /// Get human-readable status
    pub fn status_description(&self, current_time: i64) -> String {
        match self.status {
            ContractStatus::Active => {
                let days = self.time_until_unlock(current_time) / (24 * 3600);
                format!("Active - {} days until timeout", days)
            }
            ContractStatus::GracePeriod => {
                let days = self.time_until_unlock(current_time) / (24 * 3600);
                format!("Grace Period - {} days remaining", days)
            }
            ContractStatus::Unlocking => "Unlocking in progress".to_string(),
            ContractStatus::Completed => "Inheritance distributed".to_string(),
            ContractStatus::Disputed => "Disputed by owner".to_string(),
            ContractStatus::Cancelled => "Cancelled".to_string(),
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum InheritanceError {
    #[error("No beneficiaries specified")]
    NoBeneficiaries,

    #[error("Invalid beneficiary percentage: {0} (must be 0-100)")]
    InvalidPercentage(u8),

    #[error("Total percentage must be 100%, got {0}%")]
    InvalidTotalPercentage(u32),

    #[error("Duplicate beneficiary detected")]
    DuplicateBeneficiary,

    #[error("Contract already cancelled")]
    ContractCancelled,

    #[error("Contract already completed")]
    ContractCompleted,

    #[error("Heartbeat timeout not reached")]
    TimeoutNotReached,

    #[error("Grace period not expired")]
    GracePeriodNotExpired,

    #[error("Insufficient locked amount")]
    InsufficientAmount,

    #[error("Unauthorized access")]
    Unauthorized,
}

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

    #[test]
    fn test_beneficiary_creation() {
        let beneficiary = Beneficiary::new([1u8; 32], 50).unwrap();
        assert_eq!(beneficiary.percentage, 50);

        let invalid = Beneficiary::new([1u8; 32], 101);
        assert!(invalid.is_err());
    }

    #[test]
    fn test_beneficiary_amount_calculation() {
        let beneficiary = Beneficiary::new([1u8; 32], 25).unwrap();
        assert_eq!(beneficiary.calculate_amount(1000), 250);
    }

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

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

        assert!(contract.is_ok());
    }

    #[test]
    fn test_invalid_total_percentage() {
        let beneficiaries = vec![
            Beneficiary::new([1u8; 32], 60).unwrap(),
            Beneficiary::new([2u8; 32], 30).unwrap(), // Total = 90%, invalid
        ];

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

        assert!(matches!(
            contract.unwrap_err(),
            InheritanceError::InvalidTotalPercentage(90)
        ));
    }

    #[test]
    fn test_heartbeat() {
        let beneficiaries = vec![Beneficiary::new([1u8; 32], 100).unwrap()];

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

        // Record heartbeat
        contract.heartbeat(100).unwrap();
        assert_eq!(contract.last_heartbeat, 100);
    }

    #[test]
    fn test_timeout_detection() {
        let beneficiaries = vec![Beneficiary::new([1u8; 32], 100).unwrap()];

        let config = InheritanceConfig {
            heartbeat_interval: 1000,
            grace_period: 500,
            ..Default::default()
        };

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

        assert!(!contract.is_timeout(500));
        assert!(contract.is_timeout(1001));
    }

    #[test]
    fn test_grace_period() {
        let beneficiaries = vec![Beneficiary::new([1u8; 32], 100).unwrap()];

        let config = InheritanceConfig {
            heartbeat_interval: 1000,
            grace_period: 500,
            ..Default::default()
        };

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

        assert!(!contract.is_grace_expired(1400));
        assert!(contract.is_grace_expired(1501));
    }

    #[test]
    fn test_status_updates() {
        let beneficiaries = vec![Beneficiary::new([1u8; 32], 100).unwrap()];

        let config = InheritanceConfig {
            heartbeat_interval: 100,
            grace_period: 50,
            ..Default::default()
        };

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

        // Initially active
        assert_eq!(contract.status, ContractStatus::Active);

        // After timeout, should enter grace period
        contract.update_status(101);
        assert_eq!(contract.status, ContractStatus::GracePeriod);

        // After grace period, should start unlocking
        contract.update_status(151);
        assert_eq!(contract.status, ContractStatus::Unlocking);
    }

    #[test]
    fn test_dispute() {
        let beneficiaries = vec![Beneficiary::new([1u8; 32], 100).unwrap()];

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

        contract.file_dispute(1000).unwrap();

        assert_eq!(contract.status, ContractStatus::Disputed);
        assert!(contract.dispute_active);
        assert_eq!(contract.last_heartbeat, 1000);
    }

    #[test]
    fn test_cancel_contract() {
        let beneficiaries = vec![Beneficiary::new([1u8; 32], 100).unwrap()];

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

        let returned = contract.cancel().unwrap();

        assert_eq!(returned, 10000);
        assert_eq!(contract.status, ContractStatus::Cancelled);
        assert_eq!(contract.locked_amount, 0);
    }
}