path: root/inheritance/src/recovery.rs

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

/// Manages dispute resolution and anti-scam protections
pub struct RecoveryManager {
    /// Active disputes
    disputes: HashMap<[u8; 32], Dispute>,

    /// Recovery attempts
    recovery_attempts: HashMap<[u8; 32], Vec<RecoveryAttempt>>,

    /// Configuration
    config: RecoveryConfig,
}

/// Configuration for recovery and dispute system
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecoveryConfig {
    /// Allow owner to dispute unlock
    pub allow_disputes: bool,

    /// Maximum disputes allowed
    pub max_disputes: u32,

    /// Cooldown period between disputes (seconds)
    pub dispute_cooldown: i64,

    /// Enable multi-signature for disputes
    pub multi_sig_disputes: bool,

    /// Required confirmations for recovery
    pub required_confirmations: u32,

    /// Enable anti-scam checks
    pub anti_scam_enabled: bool,
}

impl Default for RecoveryConfig {
    fn default() -> Self {
        Self {
            allow_disputes: true,
            max_disputes: 3,                    // Max 3 disputes
            dispute_cooldown: 7 * 24 * 3600,    // 7 days between disputes
            multi_sig_disputes: false,
            required_confirmations: 1,
            anti_scam_enabled: true,
        }
    }
}

/// A dispute filed by the owner
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Dispute {
    /// Contract ID
    pub contract_id: [u8; 32],

    /// Who filed the dispute
    pub filed_by: [u8; 32],

    /// When the dispute was filed
    pub filed_at: i64,

    /// Dispute reason
    pub reason: DisputeReason,

    /// Dispute status
    pub status: DisputeStatus,

    /// Evidence provided
    pub evidence: Vec<Evidence>,

    /// Resolution
    pub resolution: Option<DisputeResolution>,

    /// Number of previous disputes
    pub dispute_count: u32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum DisputeReason {
    /// Owner is alive and well
    OwnerAlive,

    /// Temporary incapacitation (hospital, etc.)
    TemporaryIncapacitation,

    /// Technical issue preventing heartbeat
    TechnicalIssue,

    /// Suspected fraudulent beneficiary
    FraudSuspected,

    /// Other reason
    Other(String),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum DisputeStatus {
    /// Pending review
    Pending,

    /// Under investigation
    UnderReview,

    /// Resolved in favor of owner
    ResolvedOwner,

    /// Resolved in favor of beneficiaries
    ResolvedBeneficiaries,

    /// Rejected (too many disputes)
    Rejected,
}

/// Evidence for a dispute
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Evidence {
    /// Evidence type
    pub evidence_type: EvidenceType,

    /// Evidence data (encrypted)
    pub data: Vec<u8>,

    /// Timestamp
    pub submitted_at: i64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum EvidenceType {
    /// Cryptographic proof of life (signed message)
    ProofOfLife,

    /// Medical certificate (encrypted)
    MedicalCertificate,

    /// Technical logs
    TechnicalLogs,

    /// Witness statement
    WitnessStatement,

    /// Other
    Other,
}

/// Resolution of a dispute
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DisputeResolution {
    /// Resolution type
    pub resolution_type: ResolutionType,

    /// Actions taken
    pub actions: Vec<ResolutionAction>,

    /// Resolved at
    pub resolved_at: i64,

    /// Notes
    pub notes: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ResolutionType {
    /// Owner proven alive, reset heartbeat
    ResetHeartbeat,

    /// Extend grace period
    ExtendGracePeriod { additional_days: i64 },

    /// Cancel unlock, resume normal operation
    CancelUnlock,

    /// Proceed with unlock (owner dispute invalid)
    ProceedWithUnlock,

    /// Freeze contract pending investigation
    FreezeContract,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ResolutionAction {
    HeartbeatReset,
    GracePeriodExtended { days: i64 },
    ContractCancelled,
    UnlockProceeded,
    ContractFrozen,
}

/// Attempt to recover/access contract
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecoveryAttempt {
    /// Who attempted recovery
    pub attempted_by: [u8; 32],

    /// When
    pub attempted_at: i64,

    /// Method used
    pub method: RecoveryMethod,

    /// Success or failure
    pub success: bool,

    /// Suspicious indicators
    pub suspicious: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RecoveryMethod {
    /// Normal progressive unlock
    ProgressiveUnlock,

    /// Shamir secret sharing
    ShamirRecovery { shares_used: u8 },

    /// Emergency recovery key
    EmergencyKey,

    /// Dispute resolution
    DisputeResolution,
}

impl RecoveryManager {
    pub fn new(config: RecoveryConfig) -> Self {
        Self {
            disputes: HashMap::new(),
            recovery_attempts: HashMap::new(),
            config,
        }
    }

    /// File a dispute
    pub fn file_dispute(
        &mut self,
        contract: &mut InheritanceContract,
        filed_by: [u8; 32],
        reason: DisputeReason,
        current_time: i64,
    ) -> Result<(), RecoveryError> {
        if !self.config.allow_disputes {
            return Err(RecoveryError::DisputesNotAllowed);
        }

        // Check if owner
        if filed_by != contract.owner {
            return Err(RecoveryError::Unauthorized);
        }

        // Check max disputes
        let previous_disputes = self.disputes.values()
            .filter(|d| d.contract_id == contract.contract_id)
            .count() as u32;

        if previous_disputes >= self.config.max_disputes {
            return Err(RecoveryError::TooManyDisputes(previous_disputes));
        }

        // Check cooldown
        if let Some(last_dispute) = self.get_last_dispute(&contract.contract_id) {
            let time_since_last = current_time - last_dispute.filed_at;
            if time_since_last < self.config.dispute_cooldown {
                return Err(RecoveryError::DisputeCooldownActive {
                    remaining: self.config.dispute_cooldown - time_since_last,
                });
            }
        }

        // Create dispute
        let dispute = Dispute {
            contract_id: contract.contract_id,
            filed_by,
            filed_at: current_time,
            reason,
            status: DisputeStatus::Pending,
            evidence: Vec::new(),
            resolution: None,
            dispute_count: previous_disputes,
        };

        // File dispute in contract
        contract.file_dispute(current_time)?;

        self.disputes.insert(contract.contract_id, dispute);

        Ok(())
    }

    /// Add evidence to dispute
    pub fn add_evidence(
        &mut self,
        contract_id: &[u8; 32],
        evidence: Evidence,
    ) -> Result<(), RecoveryError> {
        let dispute = self.disputes.get_mut(contract_id)
            .ok_or(RecoveryError::DisputeNotFound)?;

        if dispute.status != DisputeStatus::Pending
            && dispute.status != DisputeStatus::UnderReview
        {
            return Err(RecoveryError::DisputeAlreadyResolved);
        }

        dispute.evidence.push(evidence);
        dispute.status = DisputeStatus::UnderReview;

        Ok(())
    }

    /// Resolve a dispute
    pub fn resolve_dispute(
        &mut self,
        contract: &mut InheritanceContract,
        resolution: DisputeResolution,
        current_time: i64,
    ) -> Result<(), RecoveryError> {
        let dispute = self.disputes.get_mut(&contract.contract_id)
            .ok_or(RecoveryError::DisputeNotFound)?;

        // Apply resolution actions
        for action in &resolution.actions {
            match action {
                ResolutionAction::HeartbeatReset => {
                    contract.heartbeat(current_time)?;
                    contract.status = ContractStatus::Active;
                }
                ResolutionAction::GracePeriodExtended { days } => {
                    contract.config.grace_period += days * 24 * 3600;
                }
                ResolutionAction::ContractCancelled => {
                    contract.cancel()?;
                }
                ResolutionAction::UnlockProceeded => {
                    contract.status = ContractStatus::Unlocking;
                    contract.dispute_active = false;
                }
                ResolutionAction::ContractFrozen => {
                    // Implementation would freeze contract
                }
            }
        }

        dispute.status = match resolution.resolution_type {
            ResolutionType::ResetHeartbeat | ResolutionType::ExtendGracePeriod { .. } | ResolutionType::CancelUnlock => {
                DisputeStatus::ResolvedOwner
            }
            ResolutionType::ProceedWithUnlock => {
                DisputeStatus::ResolvedBeneficiaries
            }
            ResolutionType::FreezeContract => {
                DisputeStatus::UnderReview
            }
        };

        dispute.resolution = Some(resolution);

        Ok(())
    }

    /// Record a recovery attempt
    pub fn record_recovery_attempt(
        &mut self,
        contract_id: [u8; 32],
        attempt: RecoveryAttempt,
    ) {
        self.recovery_attempts
            .entry(contract_id)
            .or_insert_with(Vec::new)
            .push(attempt);
    }

    /// Check for suspicious activity
    pub fn check_suspicious_activity(&self, contract_id: &[u8; 32]) -> SuspiciousActivityReport {
        let attempts = self.recovery_attempts.get(contract_id);

        let mut report = SuspiciousActivityReport {
            contract_id: *contract_id,
            suspicious: false,
            indicators: Vec::new(),
        };

        if let Some(attempts) = attempts {
            // Check for multiple failed attempts
            let failed_attempts = attempts.iter().filter(|a| !a.success).count();
            if failed_attempts > 5 {
                report.suspicious = true;
                report.indicators.push(format!(
                    "{} failed recovery attempts",
                    failed_attempts
                ));
            }

            // Check for attempts from unknown parties
            let unique_attemptors: std::collections::HashSet<_> =
                attempts.iter().map(|a| a.attempted_by).collect();

            if unique_attemptors.len() > 3 {
                report.suspicious = true;
                report.indicators.push(format!(
                    "{} different parties attempted recovery",
                    unique_attemptors.len()
                ));
            }

            // Check for suspicious timing (many attempts in short time)
            if attempts.len() > 10 {
                let first = attempts.first().unwrap().attempted_at;
                let last = attempts.last().unwrap().attempted_at;
                let timespan_hours = (last - first) / 3600;

                if timespan_hours < 24 {
                    report.suspicious = true;
                    report.indicators.push(format!(
                        "{} attempts in {} hours",
                        attempts.len(),
                        timespan_hours
                    ));
                }
            }
        }

        report
    }

    /// Get last dispute for a contract
    fn get_last_dispute(&self, contract_id: &[u8; 32]) -> Option<&Dispute> {
        self.disputes.get(contract_id)
    }

    /// Get dispute statistics
    pub fn get_dispute_stats(&self, contract_id: &[u8; 32]) -> DisputeStats {
        let total_disputes = self.disputes.values()
            .filter(|d| d.contract_id == *contract_id)
            .count();

        let resolved_owner = self.disputes.values()
            .filter(|d| d.contract_id == *contract_id && d.status == DisputeStatus::ResolvedOwner)
            .count();

        let resolved_beneficiaries = self.disputes.values()
            .filter(|d| d.contract_id == *contract_id && d.status == DisputeStatus::ResolvedBeneficiaries)
            .count();

        DisputeStats {
            total_disputes,
            resolved_owner,
            resolved_beneficiaries,
            pending: total_disputes - resolved_owner - resolved_beneficiaries,
        }
    }
}

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

/// Report of suspicious activity
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SuspiciousActivityReport {
    pub contract_id: [u8; 32],
    pub suspicious: bool,
    pub indicators: Vec<String>,
}

/// Statistics about disputes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DisputeStats {
    pub total_disputes: usize,
    pub resolved_owner: usize,
    pub resolved_beneficiaries: usize,
    pub pending: usize,
}

#[derive(Debug, thiserror::Error)]
pub enum RecoveryError {
    #[error("Disputes are not allowed for this contract")]
    DisputesNotAllowed,

    #[error("Unauthorized: only owner can file disputes")]
    Unauthorized,

    #[error("Too many disputes filed: {0}")]
    TooManyDisputes(u32),

    #[error("Dispute cooldown active: {remaining} seconds remaining")]
    DisputeCooldownActive { remaining: i64 },

    #[error("Dispute not found")]
    DisputeNotFound,

    #[error("Dispute already resolved")]
    DisputeAlreadyResolved,

    #[error("Contract error: {0}")]
    ContractError(String),
}

impl From<crate::contract::InheritanceError> for RecoveryError {
    fn from(e: crate::contract::InheritanceError) -> Self {
        RecoveryError::ContractError(e.to_string())
    }
}

#[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_file_dispute() {
        let mut manager = RecoveryManager::default();
        let mut contract = create_test_contract();

        let result = manager.file_dispute(
            &mut contract,
            [0u8; 32], // Owner
            DisputeReason::OwnerAlive,
            1000,
        );

        assert!(result.is_ok());
        assert!(contract.dispute_active);
        assert_eq!(contract.status, ContractStatus::Disputed);
    }

    #[test]
    fn test_unauthorized_dispute() {
        let mut manager = RecoveryManager::default();
        let mut contract = create_test_contract();

        let result = manager.file_dispute(
            &mut contract,
            [99u8; 32], // Not owner
            DisputeReason::OwnerAlive,
            1000,
        );

        assert!(result.is_err());
    }

    #[test]
    fn test_max_disputes() {
        let mut config = RecoveryConfig::default();
        config.max_disputes = 2;
        config.dispute_cooldown = 0; // No cooldown for testing

        let mut manager = RecoveryManager::new(config);
        let mut contract = create_test_contract();

        // First dispute - OK
        manager.file_dispute(&mut contract, [0u8; 32], DisputeReason::OwnerAlive, 1000).unwrap();

        // Reset contract for second dispute
        contract.dispute_active = false;
        contract.status = ContractStatus::Active;

        // Second dispute - OK
        manager.file_dispute(&mut contract, [0u8; 32], DisputeReason::OwnerAlive, 2000).unwrap();

        // Reset for third dispute
        contract.dispute_active = false;
        contract.status = ContractStatus::Active;

        // Third dispute - Should fail (max 2)
        let result = manager.file_dispute(&mut contract, [0u8; 32], DisputeReason::OwnerAlive, 3000);
        assert!(result.is_err());
    }

    #[test]
    fn test_add_evidence() {
        let mut manager = RecoveryManager::default();
        let mut contract = create_test_contract();

        manager.file_dispute(&mut contract, [0u8; 32], DisputeReason::OwnerAlive, 1000).unwrap();

        let evidence = Evidence {
            evidence_type: EvidenceType::ProofOfLife,
            data: vec![1, 2, 3],
            submitted_at: 1100,
        };

        let result = manager.add_evidence(&contract.contract_id, evidence);
        assert!(result.is_ok());

        let dispute = manager.disputes.get(&contract.contract_id).unwrap();
        assert_eq!(dispute.evidence.len(), 1);
        assert_eq!(dispute.status, DisputeStatus::UnderReview);
    }

    #[test]
    fn test_resolve_dispute_reset_heartbeat() {
        let mut manager = RecoveryManager::default();
        let mut contract = create_test_contract();

        manager.file_dispute(&mut contract, [0u8; 32], DisputeReason::OwnerAlive, 1000).unwrap();

        let resolution = DisputeResolution {
            resolution_type: ResolutionType::ResetHeartbeat,
            actions: vec![ResolutionAction::HeartbeatReset],
            resolved_at: 2000,
            notes: None,
        };

        manager.resolve_dispute(&mut contract, resolution, 2000).unwrap();

        assert_eq!(contract.last_heartbeat, 2000);
        assert_eq!(contract.status, ContractStatus::Active);

        let dispute = manager.disputes.get(&contract.contract_id).unwrap();
        assert_eq!(dispute.status, DisputeStatus::ResolvedOwner);
    }

    #[test]
    fn test_record_recovery_attempt() {
        let mut manager = RecoveryManager::default();
        let contract_id = [1u8; 32];

        let attempt = RecoveryAttempt {
            attempted_by: [2u8; 32],
            attempted_at: 1000,
            method: RecoveryMethod::ProgressiveUnlock,
            success: false,
            suspicious: false,
        };

        manager.record_recovery_attempt(contract_id, attempt);

        let attempts = manager.recovery_attempts.get(&contract_id).unwrap();
        assert_eq!(attempts.len(), 1);
    }

    #[test]
    fn test_suspicious_activity_detection() {
        let mut manager = RecoveryManager::default();
        let contract_id = [1u8; 32];

        // Record many failed attempts
        for i in 0..10 {
            let attempt = RecoveryAttempt {
                attempted_by: [(i % 5) as u8; 32], // Different attemptors
                attempted_at: 1000 + i * 100,
                method: RecoveryMethod::EmergencyKey,
                success: false,
                suspicious: false,
            };
            manager.record_recovery_attempt(contract_id, attempt);
        }

        let report = manager.check_suspicious_activity(&contract_id);

        assert!(report.suspicious);
        assert!(!report.indicators.is_empty());
    }

    #[test]
    fn test_dispute_stats() {
        let mut manager = RecoveryManager::default();
        let mut contract = create_test_contract();

        manager.file_dispute(&mut contract, [0u8; 32], DisputeReason::OwnerAlive, 1000).unwrap();

        let stats = manager.get_dispute_stats(&contract.contract_id);

        assert_eq!(stats.total_disputes, 1);
        assert_eq!(stats.pending, 1);
        assert_eq!(stats.resolved_owner, 0);
    }
}