path: root/inheritance/src/secret_sharing.rs

use serde::{Deserialize, Serialize};
use sharks::{Sharks, Share};

/// Shamir Secret Sharing for OnionCoin inheritance
/// Allows splitting seed phrase into N shares, requiring M shares to recover
pub struct ShamirShares {
    /// Threshold (M of N)
    threshold: u8,

    /// Total shares (N)
    total_shares: u8,

    /// Share distribution info
    distribution: Vec<ShareInfo>,
}

/// Information about a share
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShareInfo {
    /// Share ID (1-indexed)
    pub share_id: u8,

    /// Recipient of this share (encrypted)
    pub recipient_pubkey: [u8; 32],

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

    /// Has share been claimed?
    pub claimed: bool,

    /// Claim timestamp
    pub claim_time: Option<i64>,
}

impl ShamirShares {
    /// Create new Shamir secret sharing scheme
    pub fn new(threshold: u8, total_shares: u8) -> Result<Self, ShamirError> {
        if threshold > total_shares {
            return Err(ShamirError::InvalidThreshold {
                threshold,
                total_shares,
            });
        }

        if threshold == 0 || total_shares == 0 {
            return Err(ShamirError::ZeroShares);
        }

        // Note: u8 can't be > 255, but keeping check for clarity
        #[allow(unused_comparisons)]
        if total_shares > 255 {
            return Err(ShamirError::TooManyShares(total_shares));
        }

        Ok(Self {
            threshold,
            total_shares,
            distribution: Vec::new(),
        })
    }

    /// Split a secret into shares
    pub fn split_secret(&mut self, secret: &[u8], recipients: &[[u8; 32]]) -> Result<Vec<Vec<u8>>, ShamirError> {
        if recipients.len() != self.total_shares as usize {
            return Err(ShamirError::RecipientCountMismatch {
                expected: self.total_shares,
                got: recipients.len() as u8,
            });
        }

        // Create Sharks instance
        let sharks = Sharks(self.threshold);

        // Generate shares from secret
        let dealer = sharks.dealer(secret);
        let shares: Vec<Share> = dealer.take(self.total_shares as usize).collect();

        // Convert shares to bytes
        let share_bytes: Vec<Vec<u8>> = shares.iter()
            .map(|s| s.y.iter().map(|gf| gf.0).collect())
            .collect();

        // Record distribution
        for (i, recipient) in recipients.iter().enumerate() {
            self.distribution.push(ShareInfo {
                share_id: (i + 1) as u8,
                recipient_pubkey: *recipient,
                encrypted_info: None,
                claimed: false,
                claim_time: None,
            });
        }

        Ok(share_bytes)
    }

    /// Recover secret from shares
    pub fn recover_secret(
        &self,
        shares: &[Vec<u8>],
    ) -> Result<Vec<u8>, ShamirError> {
        if shares.len() < self.threshold as usize {
            return Err(ShamirError::InsufficientShares {
                required: self.threshold,
                provided: shares.len() as u8,
            });
        }

        // Convert bytes back to Share objects
        let share_objects: Result<Vec<Share>, _> = shares
            .iter()
            .map(|s| Share::try_from(s.as_slice()))
            .collect();

        let share_objects = share_objects.map_err(|_| ShamirError::InvalidShare)?;

        // Create Sharks instance and recover
        let sharks = Sharks(self.threshold);

        sharks
            .recover(&share_objects)
            .map_err(|_| ShamirError::RecoveryFailed)
    }

    /// Mark a share as claimed
    pub fn claim_share(&mut self, share_id: u8, current_time: i64) -> Result<(), ShamirError> {
        let share_info = self
            .distribution
            .iter_mut()
            .find(|s| s.share_id == share_id)
            .ok_or(ShamirError::ShareNotFound(share_id))?;

        if share_info.claimed {
            return Err(ShamirError::ShareAlreadyClaimed(share_id));
        }

        share_info.claimed = true;
        share_info.claim_time = Some(current_time);

        Ok(())
    }

    /// Get number of claimed shares
    pub fn claimed_count(&self) -> usize {
        self.distribution.iter().filter(|s| s.claimed).count()
    }

    /// Check if recovery is possible (enough shares claimed)
    pub fn can_recover(&self) -> bool {
        self.claimed_count() >= self.threshold as usize
    }

    /// Get recovery status
    pub fn recovery_status(&self) -> String {
        let claimed = self.claimed_count();
        let required = self.threshold as usize;

        if claimed >= required {
            format!("✓ Recovery possible: {}/{} shares claimed", claimed, required)
        } else {
            format!(
                "⏳ Waiting for shares: {}/{} claimed ({} more needed)",
                claimed,
                required,
                required - claimed
            )
        }
    }

    /// Get configuration
    pub fn config(&self) -> (u8, u8) {
        (self.threshold, self.total_shares)
    }
}

/// Distribution plan for shares
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShareDistribution {
    /// Shamir configuration
    pub threshold: u8,
    pub total_shares: u8,

    /// Share holders (trusted guardians)
    pub guardians: Vec<Guardian>,

    /// Instructions for recovery
    pub recovery_instructions: String,
}

/// A guardian who holds a share
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Guardian {
    pub pubkey: [u8; 32],
    pub name: Option<String>, // Optional encrypted name
    pub contact: Option<String>, // Optional encrypted contact info
    pub share_id: u8,
}

impl ShareDistribution {
    /// Create a standard 3-of-5 distribution
    pub fn standard_3_of_5(guardians: Vec<[u8; 32]>) -> Result<Self, ShamirError> {
        if guardians.len() != 5 {
            return Err(ShamirError::RecipientCountMismatch {
                expected: 5,
                got: guardians.len() as u8,
            });
        }

        let guardians = guardians
            .into_iter()
            .enumerate()
            .map(|(i, pubkey)| Guardian {
                pubkey,
                name: None,
                contact: None,
                share_id: (i + 1) as u8,
            })
            .collect();

        Ok(Self {
            threshold: 3,
            total_shares: 5,
            guardians,
            recovery_instructions: "Contact any 3 of the 5 guardians to recover the secret.".to_string(),
        })
    }

    /// Create a custom distribution
    pub fn custom(
        threshold: u8,
        guardians: Vec<Guardian>,
        instructions: String,
    ) -> Result<Self, ShamirError> {
        if threshold > guardians.len() as u8 {
            return Err(ShamirError::InvalidThreshold {
                threshold,
                total_shares: guardians.len() as u8,
            });
        }

        Ok(Self {
            threshold,
            total_shares: guardians.len() as u8,
            guardians,
            recovery_instructions: instructions,
        })
    }

    /// Generate and distribute shares
    pub fn generate_shares(&self, secret: &[u8]) -> Result<Vec<(Guardian, Vec<u8>)>, ShamirError> {
        let mut shamir = ShamirShares::new(self.threshold, self.total_shares)?;

        let recipients: Vec<[u8; 32]> = self.guardians.iter().map(|g| g.pubkey).collect();

        let shares = shamir.split_secret(secret, &recipients)?;

        Ok(self
            .guardians
            .iter()
            .zip(shares)
            .map(|(g, s)| (g.clone(), s))
            .collect())
    }
}

#[derive(Debug, thiserror::Error)]
pub enum ShamirError {
    #[error("Invalid threshold: {threshold} of {total_shares} (threshold must be ≤ total)")]
    InvalidThreshold { threshold: u8, total_shares: u8 },

    #[error("Threshold and total shares cannot be zero")]
    ZeroShares,

    #[error("Too many shares: {0} (maximum 255)")]
    TooManyShares(u8),

    #[error("Recipient count mismatch: expected {expected}, got {got}")]
    RecipientCountMismatch { expected: u8, got: u8 },

    #[error("Insufficient shares for recovery: need {required}, got {provided}")]
    InsufficientShares { required: u8, provided: u8 },

    #[error("Invalid share data")]
    InvalidShare,

    #[error("Secret recovery failed")]
    RecoveryFailed,

    #[error("Share {0} not found")]
    ShareNotFound(u8),

    #[error("Share {0} already claimed")]
    ShareAlreadyClaimed(u8),
}

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

    #[test]
    fn test_shamir_creation() {
        let shamir = ShamirShares::new(3, 5);
        assert!(shamir.is_ok());

        let invalid = ShamirShares::new(5, 3); // Threshold > total
        assert!(invalid.is_err());
    }

    #[test]
    fn test_secret_splitting_and_recovery() {
        let secret = b"my super secret seed phrase here";
        let recipients = vec![[1u8; 32], [2u8; 32], [3u8; 32], [4u8; 32], [5u8; 32]];

        let mut shamir = ShamirShares::new(3, 5).unwrap();

        // Split secret
        let shares = shamir.split_secret(secret, &recipients).unwrap();
        assert_eq!(shares.len(), 5);

        // Recover with 3 shares
        let recovered = shamir.recover_secret(&shares[0..3]).unwrap();
        assert_eq!(recovered, secret);

        // Recover with different 3 shares
        let recovered2 = shamir.recover_secret(&[shares[1].clone(), shares[3].clone(), shares[4].clone()]).unwrap();
        assert_eq!(recovered2, secret);

        // Recover with all 5 shares
        let recovered3 = shamir.recover_secret(&shares).unwrap();
        assert_eq!(recovered3, secret);
    }

    #[test]
    fn test_insufficient_shares() {
        let secret = b"my super secret seed phrase here";
        let recipients = vec![[1u8; 32], [2u8; 32], [3u8; 32], [4u8; 32], [5u8; 32]];

        let mut shamir = ShamirShares::new(3, 5).unwrap();
        let shares = shamir.split_secret(secret, &recipients).unwrap();

        // Try to recover with only 2 shares
        let result = shamir.recover_secret(&shares[0..2]);
        assert!(result.is_err());
    }

    #[test]
    fn test_share_claiming() {
        let secret = b"test secret";
        let recipients = vec![[1u8; 32], [2u8; 32], [3u8; 32]];

        let mut shamir = ShamirShares::new(2, 3).unwrap();
        shamir.split_secret(secret, &recipients).unwrap();

        assert_eq!(shamir.claimed_count(), 0);
        assert!(!shamir.can_recover());

        // Claim first share
        shamir.claim_share(1, 1000).unwrap();
        assert_eq!(shamir.claimed_count(), 1);
        assert!(!shamir.can_recover());

        // Claim second share
        shamir.claim_share(2, 2000).unwrap();
        assert_eq!(shamir.claimed_count(), 2);
        assert!(shamir.can_recover());

        // Try to claim already claimed share
        let result = shamir.claim_share(1, 3000);
        assert!(result.is_err());
    }

    #[test]
    fn test_share_distribution() {
        let guardians = vec![[1u8; 32], [2u8; 32], [3u8; 32], [4u8; 32], [5u8; 32]];

        let distribution = ShareDistribution::standard_3_of_5(guardians).unwrap();

        assert_eq!(distribution.threshold, 3);
        assert_eq!(distribution.total_shares, 5);
        assert_eq!(distribution.guardians.len(), 5);
    }

    #[test]
    fn test_distribution_share_generation() {
        let guardians = vec![[1u8; 32], [2u8; 32], [3u8; 32], [4u8; 32], [5u8; 32]];

        let distribution = ShareDistribution::standard_3_of_5(guardians).unwrap();
        let secret = b"my wallet seed phrase";

        let shares = distribution.generate_shares(secret).unwrap();

        assert_eq!(shares.len(), 5);

        // Extract just the share bytes
        let share_bytes: Vec<Vec<u8>> = shares.iter().map(|(_, s)| s.clone()).collect();

        // Test recovery with 3 shares
        let shamir = ShamirShares::new(3, 5).unwrap();
        let recovered = shamir.recover_secret(&share_bytes[0..3]).unwrap();

        assert_eq!(recovered, secret);
    }

    #[test]
    fn test_recovery_status() {
        let secret = b"test";
        let recipients = vec![[1u8; 32], [2u8; 32], [3u8; 32]];

        let mut shamir = ShamirShares::new(2, 3).unwrap();
        shamir.split_secret(secret, &recipients).unwrap();

        let status = shamir.recovery_status();
        assert!(status.contains("0/2"));

        shamir.claim_share(1, 0).unwrap();
        let status = shamir.recovery_status();
        assert!(status.contains("1/2"));

        shamir.claim_share(2, 0).unwrap();
        let status = shamir.recovery_status();
        assert!(status.contains("Recovery possible"));
    }
}