//! Shamir's (t,n) threshold secret sharing.

use crate::polynomial::{eval_poly, lagrange_interpolate_zero, PRIME};

/// A share in Shamir's secret sharing scheme.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Share {
    /// The x-coordinate (evaluator index) of this share.
    pub x: u64,
    /// The y-coordinate (polynomial evaluation at x) of this share.
    pub y: u64,
}

/// Shamir's Secret Sharing implementation.
pub struct ShamirSecretSharing {
    /// Threshold number of shares needed to reconstruct.
    threshold: usize,
    /// Total number of shares to generate.
    num_shares: usize,
}

impl ShamirSecretSharing {
    /// Create a new (t, n) Shamir secret sharing scheme.
    /// `threshold` shares are needed to reconstruct the secret.
    /// `num_shares` shares will be generated.
    pub fn new(threshold: usize, num_shares: usize) -> Self {
        assert!(threshold > 0, "Threshold must be at least 1");
        assert!(threshold <= num_shares, "Threshold cannot exceed num_shares");
        Self { threshold, num_shares }
    }

    /// Split a secret into shares.
    pub fn split(&self, secret: u64) -> Vec<Share> {
        assert!(secret < PRIME, "Secret must be less than PRIME");
        self.split_with_rng(secret, simple_rng)
    }

    /// Split with a custom RNG function (for testing).
    fn split_with_rng<F>(&self, secret: u64, mut rng: F) -> Vec<Share>
    where
        F: FnMut() -> u64,
    {
        // Generate random polynomial coefficients
        // coeffs[0] = secret, coeffs[1..threshold] are random
        let mut coeffs = vec![secret];
        for _ in 1..self.threshold {
            coeffs.push(rng() % PRIME);
        }

        // Evaluate polynomial at x = 1, 2, ..., n
        (1..=self.num_shares)
            .map(|x| {
                let y = eval_poly(&coeffs, x as u64);
                Share { x: x as u64, y }
            })
            .collect()
    }

    /// Reconstruct the secret from a subset of shares.
    /// Requires at least `threshold` shares.
    pub fn reconstruct(&self, shares: &[Share]) -> u64 {
        assert!(shares.len() >= self.threshold, "Not enough shares");
        let points: Vec<(u64, u64)> = shares[..self.threshold]
            .iter()
            .map(|s| (s.x, s.y))
            .collect();
        lagrange_interpolate_zero(&points)
    }

    /// Reconstruct from exactly the given shares (no threshold check).
    pub fn reconstruct_from(&self, shares: &[Share]) -> u64 {
        assert!(!shares.is_empty(), "Need at least one share");
        let points: Vec<(u64, u64)> = shares.iter().map(|s| (s.x, s.y)).collect();
        lagrange_interpolate_zero(&points)
    }

    /// Get the threshold.
    pub fn threshold(&self) -> usize {
        self.threshold
    }

    /// Get the number of shares.
    pub fn num_shares(&self) -> usize {
        self.num_shares
    }
}

/// Simple deterministic RNG for share generation.
fn simple_rng() -> u64 {
    use core::cell::Cell;
    thread_local! {
        static STATE: Cell<u64> = const { Cell::new(0x853c49e6748fea9b) };
    }
    STATE.with(|s| {
        let mut v = s.get();
        v ^= v << 13;
        v ^= v >> 7;
        v ^= v << 17;
        s.set(v);
        v
    })
}

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

    #[test]
    fn test_shamir_2_of_3() {
        let sss = ShamirSecretSharing::new(2, 3);
        let secret = 12345u64;
        let shares = sss.split(secret);
        assert_eq!(shares.len(), 3);
        // Any 2 shares should work
        assert_eq!(sss.reconstruct(&shares[0..2]), secret);
        assert_eq!(sss.reconstruct(&shares[1..3]), secret);
        assert_eq!(sss.reconstruct(&[shares[0].clone(), shares[2].clone()]), secret);
    }

    #[test]
    fn test_shamir_3_of_5() {
        let sss = ShamirSecretSharing::new(3, 5);
        let secret = 999888777u64;
        let shares = sss.split(secret);
        // Any 3 shares should work
        assert_eq!(sss.reconstruct(&shares[0..3]), secret);
        assert_eq!(sss.reconstruct(&shares[2..5]), secret);
        assert_eq!(sss.reconstruct(&[shares[0].clone(), shares[2].clone(), shares[4].clone()]), secret);
    }

    #[test]
    fn test_shamir_secret_zero() {
        let sss = ShamirSecretSharing::new(2, 4);
        let shares = sss.split(0);
        assert_eq!(sss.reconstruct(&shares[0..2]), 0);
    }

    #[test]
    fn test_shamir_threshold_property() {
        let sss = ShamirSecretSharing::new(3, 5);
        let secret = 42u64;
        let shares = sss.split(secret);
        // Only 2 shares should NOT give the secret (with high probability)
        let reconstructed = sss.reconstruct_from(&shares[0..2]);
        // With 2 shares from a degree-2 polynomial, reconstruction gives a different value
        // (unless very unlucky with the random polynomial)
        assert_ne!(reconstructed, secret);
    }

    #[test]
    fn test_shamir_perfect_secrecy() {
        // A single share from a (2,n) scheme reveals nothing about the secret
        let sss = ShamirSecretSharing::new(2, 5);
        let secret1 = 100u64;
        let secret2 = 200u64;
        // Each single share is a point on a random line.
        // The share value doesn't determine the secret uniquely.
        let shares1 = sss.split(secret1);
        let shares2 = sss.split(secret2);
        // The shares could have any value regardless of secret
        // (this is a probabilistic property, but both sets should be non-trivial)
        assert!(!shares1.iter().all(|s| s.y == 0));
        assert!(!shares2.iter().all(|s| s.y == 0));
    }

    #[test]
    fn test_shamir_1_of_n() {
        let sss = ShamirSecretSharing::new(1, 5);
        let secret = 42u64;
        let shares = sss.split(secret);
        // Any single share gives the secret (degenerate case)
        for share in &shares {
            assert_eq!(sss.reconstruct(std::slice::from_ref(share)), secret);
        }
    }

    #[test]
    #[should_panic(expected = "Not enough shares")]
    fn test_shamir_too_few_shares() {
        let sss = ShamirSecretSharing::new(3, 5);
        let shares = sss.split(42u64);
        let _ = sss.reconstruct(&shares[0..2]);
    }
}
