kaccy-core 0.2.0

//! Threshold Encryption for Orders
//!
//! Provides comprehensive threshold encryption support including:
//! - Time-based decryption keys
//! - Distributed key generation (DKG)
//! - Encrypted mempool for private transactions
//! - Fair transaction ordering via time-locked encryption

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;

use crate::error::{CoreError, Result};

/// Threshold encryption scheme type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum EncryptionScheme {
    /// Simple threshold scheme (simplified for demo)
    SimpleThreshold,
    /// Shamir's Secret Sharing based
    Shamir,
    /// Verifiable Secret Sharing
    VSS,
}

/// Key share for threshold encryption
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyShare {
    /// Share ID
    pub id: String,
    /// Party ID (who holds this share)
    pub party_id: String,
    /// Share index
    pub index: u32,
    /// Share data (encrypted/encoded)
    pub share_data: Vec<u8>,
    /// Threshold (minimum shares needed)
    pub threshold: u32,
    /// Total number of shares
    pub total_shares: u32,
    /// Creation timestamp
    pub created_at: DateTime<Utc>,
}

impl KeyShare {
    /// Create a new key share
    pub fn new(
        id: String,
        party_id: String,
        index: u32,
        share_data: Vec<u8>,
        threshold: u32,
        total_shares: u32,
    ) -> Result<Self> {
        if threshold > total_shares {
            return Err(CoreError::Validation(
                "Threshold cannot exceed total shares".to_string(),
            ));
        }
        if index >= total_shares {
            return Err(CoreError::Validation(
                "Share index must be less than total shares".to_string(),
            ));
        }

        Ok(Self {
            id,
            party_id,
            index,
            share_data,
            threshold,
            total_shares,
            created_at: Utc::now(),
        })
    }
}

/// Time-locked encryption key
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TimeLockedKey {
    /// Key ID
    pub id: String,
    /// Encrypted key material
    pub encrypted_key: Vec<u8>,
    /// Unlock time
    pub unlock_time: DateTime<Utc>,
    /// Block height for unlocking (alternative to time)
    pub unlock_block: Option<u64>,
    /// Puzzle for time-lock (e.g., hash chain)
    pub time_lock_puzzle: Vec<u8>,
    /// Creation timestamp
    pub created_at: DateTime<Utc>,
}

impl TimeLockedKey {
    /// Create a new time-locked key
    pub fn new(
        id: String,
        encrypted_key: Vec<u8>,
        unlock_time: DateTime<Utc>,
        unlock_block: Option<u64>,
    ) -> Self {
        // Generate a simple time-lock puzzle (hash chain)
        let puzzle = Self::generate_puzzle(&encrypted_key, unlock_time);

        Self {
            id,
            encrypted_key,
            unlock_time,
            unlock_block,
            time_lock_puzzle: puzzle,
            created_at: Utc::now(),
        }
    }

    /// Generate a time-lock puzzle
    fn generate_puzzle(key: &[u8], unlock_time: DateTime<Utc>) -> Vec<u8> {
        let mut hasher = Sha256::new();
        hasher.update(key);
        hasher.update(unlock_time.timestamp().to_le_bytes());
        hasher.finalize().to_vec()
    }

    /// Check if the key can be unlocked
    pub fn can_unlock(&self, current_time: DateTime<Utc>, current_block: Option<u64>) -> bool {
        // Check time-based unlock
        if current_time >= self.unlock_time {
            return true;
        }

        // Check block-based unlock
        if let (Some(unlock_block), Some(current)) = (self.unlock_block, current_block) {
            if current >= unlock_block {
                return true;
            }
        }

        false
    }

    /// Attempt to unlock the key
    pub fn unlock(
        &self,
        current_time: DateTime<Utc>,
        current_block: Option<u64>,
    ) -> Result<Vec<u8>> {
        if !self.can_unlock(current_time, current_block) {
            return Err(CoreError::Validation(
                "Key cannot be unlocked yet".to_string(),
            ));
        }

        // In a real implementation, this would perform cryptographic unlocking
        // For now, we just return the encrypted key (which would be decrypted)
        Ok(self.encrypted_key.clone())
    }
}

/// Encrypted order in the mempool
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptedOrder {
    /// Order ID
    pub id: String,
    /// Encrypted order data
    pub ciphertext: Vec<u8>,
    /// Encryption metadata
    pub encryption_metadata: EncryptionMetadata,
    /// Time-locked key reference
    pub timelock_key_id: String,
    /// Submission timestamp
    pub submitted_at: DateTime<Utc>,
    /// Reveal timestamp (when decrypted)
    pub revealed_at: Option<DateTime<Utc>>,
}

/// Encryption metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptionMetadata {
    /// Encryption scheme used
    pub scheme: EncryptionScheme,
    /// Threshold required for decryption
    pub threshold: u32,
    /// Total key shares
    pub total_shares: u32,
    /// Nonce/IV for encryption
    pub nonce: Vec<u8>,
}

impl EncryptedOrder {
    /// Create a new encrypted order
    pub fn new(
        id: String,
        ciphertext: Vec<u8>,
        scheme: EncryptionScheme,
        threshold: u32,
        total_shares: u32,
        timelock_key_id: String,
    ) -> Self {
        // Generate a random nonce (in real implementation, use proper RNG)
        let nonce = Self::generate_nonce();

        Self {
            id,
            ciphertext,
            encryption_metadata: EncryptionMetadata {
                scheme,
                threshold,
                total_shares,
                nonce,
            },
            timelock_key_id,
            submitted_at: Utc::now(),
            revealed_at: None,
        }
    }

    /// Generate a random nonce
    fn generate_nonce() -> Vec<u8> {
        // In real implementation, use a CSPRNG
        let mut hasher = Sha256::new();
        hasher.update(Utc::now().timestamp_nanos_opt().unwrap_or(0).to_le_bytes());
        hasher.finalize()[0..16].to_vec()
    }

    /// Mark order as revealed
    pub fn mark_revealed(&mut self) {
        self.revealed_at = Some(Utc::now());
    }

    /// Check if order is revealed
    pub fn is_revealed(&self) -> bool {
        self.revealed_at.is_some()
    }
}

/// Distributed Key Generation (DKG) participant
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DKGParticipant {
    /// Participant ID
    pub id: String,
    /// Participant's address
    pub address: String,
    /// Public key (for verification)
    pub public_key: Vec<u8>,
    /// Commitment to secret share
    pub commitment: Vec<u8>,
    /// Status
    pub status: ParticipantStatus,
    /// Joined timestamp
    pub joined_at: DateTime<Utc>,
}

/// Participant status in DKG
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ParticipantStatus {
    /// Registered and waiting
    Registered,
    /// Actively participating
    Active,
    /// Completed their part
    Completed,
    /// Failed or dropped out
    Failed,
}

impl DKGParticipant {
    /// Create a new DKG participant
    pub fn new(id: String, address: String, public_key: Vec<u8>) -> Self {
        // Generate commitment (hash of public key in this simplified version)
        let mut hasher = Sha256::new();
        hasher.update(&public_key);
        let commitment = hasher.finalize().to_vec();

        Self {
            id,
            address,
            public_key,
            commitment,
            status: ParticipantStatus::Registered,
            joined_at: Utc::now(),
        }
    }

    /// Activate the participant
    pub fn activate(&mut self) {
        self.status = ParticipantStatus::Active;
    }

    /// Mark as completed
    pub fn complete(&mut self) {
        self.status = ParticipantStatus::Completed;
    }

    /// Mark as failed
    pub fn fail(&mut self) {
        self.status = ParticipantStatus::Failed;
    }
}

/// DKG session for generating distributed keys
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DKGSession {
    /// Session ID
    pub id: String,
    /// Threshold (minimum participants needed)
    pub threshold: u32,
    /// Participants
    pub participants: HashMap<String, DKGParticipant>,
    /// Generated key shares
    pub key_shares: Vec<KeyShare>,
    /// Public key (aggregated)
    pub public_key: Option<Vec<u8>>,
    /// Session status
    pub status: DKGSessionStatus,
    /// Started timestamp
    pub started_at: DateTime<Utc>,
    /// Completed timestamp
    pub completed_at: Option<DateTime<Utc>>,
}

/// DKG session status
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum DKGSessionStatus {
    /// Collecting participants
    Initializing,
    /// Running DKG protocol
    Running,
    /// Successfully completed
    Completed,
    /// Failed
    Failed,
}

impl DKGSession {
    /// Create a new DKG session
    pub fn new(id: String, threshold: u32) -> Self {
        Self {
            id,
            threshold,
            participants: HashMap::new(),
            key_shares: Vec::new(),
            public_key: None,
            status: DKGSessionStatus::Initializing,
            started_at: Utc::now(),
            completed_at: None,
        }
    }

    /// Add a participant to the session
    pub fn add_participant(&mut self, participant: DKGParticipant) -> Result<()> {
        if self.status != DKGSessionStatus::Initializing {
            return Err(CoreError::Validation(
                "Cannot add participants after session has started".to_string(),
            ));
        }

        if self.participants.contains_key(&participant.id) {
            return Err(CoreError::Validation(format!(
                "Participant {} already in session",
                participant.id
            )));
        }

        self.participants
            .insert(participant.id.clone(), participant);
        Ok(())
    }

    /// Start the DKG protocol
    pub fn start(&mut self) -> Result<()> {
        let participant_count = self.participants.len() as u32;
        if participant_count < self.threshold {
            return Err(CoreError::Validation(format!(
                "Not enough participants: {} < {}",
                participant_count, self.threshold
            )));
        }

        self.status = DKGSessionStatus::Running;

        // Activate all participants
        for participant in self.participants.values_mut() {
            participant.activate();
        }

        Ok(())
    }

    /// Generate key shares (simplified implementation)
    pub fn generate_shares(&mut self) -> Result<()> {
        if self.status != DKGSessionStatus::Running {
            return Err(CoreError::Validation(
                "Session must be running to generate shares".to_string(),
            ));
        }

        let participant_count = self.participants.len() as u32;

        // Generate a "master secret" (simplified)
        let mut hasher = Sha256::new();
        hasher.update(self.id.as_bytes());
        hasher.update(self.started_at.timestamp().to_le_bytes());
        let master_secret = hasher.finalize().to_vec();

        // Generate shares for each participant
        for (idx, (participant_id, participant)) in self.participants.iter_mut().enumerate() {
            // In real implementation, use proper secret sharing (Shamir, etc.)
            let mut share_hasher = Sha256::new();
            share_hasher.update(&master_secret);
            share_hasher.update((idx as u32).to_le_bytes());
            share_hasher.update(participant_id.as_bytes());
            let share_data = share_hasher.finalize().to_vec();

            let key_share = KeyShare::new(
                format!("share-{}-{}", self.id, idx),
                participant_id.clone(),
                idx as u32,
                share_data,
                self.threshold,
                participant_count,
            )?;

            self.key_shares.push(key_share);
            participant.complete();
        }

        // Generate aggregated public key (simplified)
        let mut pk_hasher = Sha256::new();
        pk_hasher.update(&master_secret);
        pk_hasher.update(b"public_key");
        self.public_key = Some(pk_hasher.finalize().to_vec());

        self.status = DKGSessionStatus::Completed;
        self.completed_at = Some(Utc::now());

        Ok(())
    }

    /// Get key share for a specific participant
    pub fn get_share(&self, participant_id: &str) -> Option<&KeyShare> {
        self.key_shares
            .iter()
            .find(|s| s.party_id == participant_id)
    }

    /// Check if session is complete
    pub fn is_complete(&self) -> bool {
        self.status == DKGSessionStatus::Completed
    }
}

/// Encrypted mempool manager
#[derive(Debug)]
pub struct EncryptedMempoolManager {
    /// Encrypted orders
    pub encrypted_orders: HashMap<String, EncryptedOrder>,
    /// Time-locked keys
    pub timelock_keys: HashMap<String, TimeLockedKey>,
    /// DKG sessions
    pub dkg_sessions: HashMap<String, DKGSession>,
    /// Revealed orders (order_id -> plaintext)
    pub revealed_orders: HashMap<String, Vec<u8>>,
}

impl EncryptedMempoolManager {
    /// Create a new encrypted mempool manager
    pub fn new() -> Self {
        Self {
            encrypted_orders: HashMap::new(),
            timelock_keys: HashMap::new(),
            dkg_sessions: HashMap::new(),
            revealed_orders: HashMap::new(),
        }
    }

    /// Add a time-locked key
    pub fn add_timelock_key(&mut self, key: TimeLockedKey) -> Result<()> {
        if self.timelock_keys.contains_key(&key.id) {
            return Err(CoreError::Validation(format!(
                "Time-locked key {} already exists",
                key.id
            )));
        }
        self.timelock_keys.insert(key.id.clone(), key);
        Ok(())
    }

    /// Submit an encrypted order
    pub fn submit_encrypted_order(&mut self, order: EncryptedOrder) -> Result<()> {
        // Verify the time-locked key exists
        if !self.timelock_keys.contains_key(&order.timelock_key_id) {
            return Err(CoreError::Validation(format!(
                "Time-locked key {} not found",
                order.timelock_key_id
            )));
        }

        if self.encrypted_orders.contains_key(&order.id) {
            return Err(CoreError::Validation(format!(
                "Order {} already exists",
                order.id
            )));
        }

        self.encrypted_orders.insert(order.id.clone(), order);
        Ok(())
    }

    /// Try to reveal orders whose time-locks have expired
    pub fn reveal_expired_orders(
        &mut self,
        current_time: DateTime<Utc>,
        current_block: Option<u64>,
    ) -> Result<Vec<String>> {
        let mut revealed_order_ids = Vec::new();

        // Find orders whose time-locks have expired
        for order in self.encrypted_orders.values_mut() {
            if order.is_revealed() {
                continue;
            }

            if let Some(timelock_key) = self.timelock_keys.get(&order.timelock_key_id) {
                if timelock_key.can_unlock(current_time, current_block) {
                    // Unlock and "decrypt" the order
                    if let Ok(_unlocked_key) = timelock_key.unlock(current_time, current_block) {
                        // In real implementation, use unlocked_key to decrypt order.ciphertext
                        // For now, we just mark it as revealed
                        order.mark_revealed();
                        self.revealed_orders
                            .insert(order.id.clone(), order.ciphertext.clone());
                        revealed_order_ids.push(order.id.clone());
                    }
                }
            }
        }

        Ok(revealed_order_ids)
    }

    /// Create a new DKG session
    pub fn create_dkg_session(&mut self, session_id: String, threshold: u32) -> Result<()> {
        if self.dkg_sessions.contains_key(&session_id) {
            return Err(CoreError::Validation(format!(
                "DKG session {} already exists",
                session_id
            )));
        }

        let session = DKGSession::new(session_id.clone(), threshold);
        self.dkg_sessions.insert(session_id, session);
        Ok(())
    }

    /// Get a DKG session
    pub fn get_dkg_session(&self, session_id: &str) -> Option<&DKGSession> {
        self.dkg_sessions.get(session_id)
    }

    /// Get a mutable DKG session
    pub fn get_dkg_session_mut(&mut self, session_id: &str) -> Option<&mut DKGSession> {
        self.dkg_sessions.get_mut(session_id)
    }

    /// Get all encrypted orders
    pub fn get_encrypted_orders(&self) -> Vec<&EncryptedOrder> {
        self.encrypted_orders.values().collect()
    }

    /// Get all revealed orders
    pub fn get_revealed_orders(&self) -> Vec<(&String, &Vec<u8>)> {
        self.revealed_orders.iter().collect()
    }

    /// Get statistics
    pub fn stats(&self) -> MempoolStats {
        let total_encrypted = self.encrypted_orders.len();
        let total_revealed = self
            .encrypted_orders
            .values()
            .filter(|o| o.is_revealed())
            .count();

        MempoolStats {
            total_encrypted_orders: total_encrypted,
            total_revealed_orders: total_revealed,
            pending_orders: total_encrypted - total_revealed,
            active_dkg_sessions: self
                .dkg_sessions
                .values()
                .filter(|s| s.status == DKGSessionStatus::Running)
                .count(),
            total_timelock_keys: self.timelock_keys.len(),
        }
    }
}

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

/// Mempool statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MempoolStats {
    /// Total encrypted orders
    pub total_encrypted_orders: usize,
    /// Total revealed orders
    pub total_revealed_orders: usize,
    /// Pending (not yet revealed) orders
    pub pending_orders: usize,
    /// Active DKG sessions
    pub active_dkg_sessions: usize,
    /// Total time-lock keys
    pub total_timelock_keys: usize,
}

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

    #[test]
    fn test_key_share_creation() {
        let share = KeyShare::new(
            "share-1".to_string(),
            "party-1".to_string(),
            0,
            vec![1, 2, 3],
            2,
            3,
        )
        .unwrap();

        assert_eq!(share.index, 0);
        assert_eq!(share.threshold, 2);
        assert_eq!(share.total_shares, 3);

        // Invalid threshold
        let invalid = KeyShare::new(
            "share-2".to_string(),
            "party-2".to_string(),
            0,
            vec![1, 2, 3],
            4, // threshold > total_shares
            3,
        );
        assert!(invalid.is_err());
    }

    #[test]
    fn test_timelock_key() {
        let unlock_time = Utc::now() + chrono::Duration::seconds(3600);
        let key = TimeLockedKey::new(
            "key-1".to_string(),
            vec![1, 2, 3, 4],
            unlock_time,
            Some(1000),
        );

        // Cannot unlock yet
        assert!(!key.can_unlock(Utc::now(), Some(500)));

        // Can unlock after time
        let future_time = Utc::now() + chrono::Duration::seconds(7200);
        assert!(key.can_unlock(future_time, Some(500)));

        // Can unlock after block height
        assert!(key.can_unlock(Utc::now(), Some(1000)));
    }

    #[test]
    fn test_dkg_session() {
        let mut session = DKGSession::new("session-1".to_string(), 2);

        // Add participants
        let p1 = DKGParticipant::new("p1".to_string(), "addr1".to_string(), vec![1, 2, 3]);
        let p2 = DKGParticipant::new("p2".to_string(), "addr2".to_string(), vec![4, 5, 6]);
        let p3 = DKGParticipant::new("p3".to_string(), "addr3".to_string(), vec![7, 8, 9]);

        session.add_participant(p1).unwrap();
        session.add_participant(p2).unwrap();
        session.add_participant(p3).unwrap();

        // Start and generate shares
        session.start().unwrap();
        assert_eq!(session.status, DKGSessionStatus::Running);

        session.generate_shares().unwrap();
        assert!(session.is_complete());
        assert_eq!(session.key_shares.len(), 3);
        assert!(session.public_key.is_some());
    }

    #[test]
    fn test_encrypted_mempool() {
        let mut manager = EncryptedMempoolManager::new();

        // Add time-locked key
        let unlock_time = Utc::now() + chrono::Duration::seconds(100);
        let timelock = TimeLockedKey::new("tl-1".to_string(), vec![1, 2, 3], unlock_time, None);
        manager.add_timelock_key(timelock).unwrap();

        // Submit encrypted order
        let order = EncryptedOrder::new(
            "order-1".to_string(),
            vec![4, 5, 6],
            EncryptionScheme::SimpleThreshold,
            2,
            3,
            "tl-1".to_string(),
        );
        manager.submit_encrypted_order(order).unwrap();

        let stats = manager.stats();
        assert_eq!(stats.total_encrypted_orders, 1);
        assert_eq!(stats.pending_orders, 1);

        // Cannot reveal yet
        let revealed = manager.reveal_expired_orders(Utc::now(), None).unwrap();
        assert_eq!(revealed.len(), 0);

        // Can reveal after time
        let future = Utc::now() + chrono::Duration::seconds(200);
        let revealed = manager.reveal_expired_orders(future, None).unwrap();
        assert_eq!(revealed.len(), 1);
    }
}